Add partition support to BigQuery

gcloud/bigquery/table.py

@@ -230,6 +230,56 @@ def table_type(self):
         """
         return self._properties.get('type')
 
+    @property
+    def partitioning_type(self):
+        """Time partitioning of the table.
+        :rtype: str, or ``NoneType``
+        :returns: Returns type if the table is partitioned, None otherwise.
+        """
+        return self._properties.get('timePartitioning', {}).get('type')
+
+    @partitioning_type.setter
+    def partitioning_type(self, value):
+        """Update the partitioning type of the table
+
+        :type value: str
+        :param value: partitioning type only "DAY" is currently supported
+        """
+        if not (isinstance(value, six.string_types)
+                and value.upper() == "DAY") and value is not None:
+            raise ValueError("value must be one of ['DAY', None]")
+
+        self._properties.setdefault('timePartitioning', {})
+        if value is not None:
+            self._properties['timePartitioning']['type'] = value.upper()
+
+    @property
+    def partition_expiration(self):
+        """Expiration time in ms for a partition
+        :rtype: int, or ``NoneType``
+        :returns: Returns the time in ms for partition expiration
+        """
+        expiry = None
+        if "timePartitioning" in self._properties:
+            time_part = self._properties.get("timePartitioning")
+            expiry = time_part.get("expirationMs")
+        return expiry
+
+    @partition_expiration.setter
+    def partition_expiration(self, value):
+        """Update the experation time in ms for a partition
+
+        :type value: int
+        :param value: partition experiation time in ms
+        """
+        if not isinstance(value, int):
+            raise ValueError("must be an integer representing millisseconds")
+        try:
+            self._properties["timePartitioning"]["expirationMs"] = value
+        except KeyError:
+            self._properties['timePartitioning'] = {'type': "DAY"}
+            self._properties["timePartitioning"]["expirationMs"] = value
+
     @property
     def description(self):
         """Description of the table.
@@ -348,6 +398,22 @@ def view_query(self):
         """Delete SQL query defining the table as a view."""
         self._properties.pop('view', None)
 
+    def list_partitions(self, client=None):
+        """List the partitions in a table.
+
+        :type client: :class:`gcloud.bigquery.client.Client` or ``NoneType``
+        :param client: the client to use.  If not passed, falls back to the
+                       ``client`` stored on the current dataset.
+
+        :rtype: list
+        :returns: a list of time partitions
+        """
+        query = self._require_client(client).run_sync_query(
+            'SELECT partition_id from [%s.%s$__PARTITIONS_SUMMARY__]' %
+            (self.dataset_name, self.name))
+        query.run()
+        return [row[0] for row in query.rows]
+
     @classmethod
     def from_api_repr(cls, resource, dataset):
         """Factory:  construct a table given its API representation
@@ -423,6 +489,9 @@ def _build_resource(self):
         if self.location is not None:
             resource['location'] = self.location
 
+        if self.partitioning_type is not None:
+            resource['timePartitioning'] = self._properties['timePartitioning']
+
         if self.view_query is not None:
             view = resource['view'] = {}
             view['query'] = self.view_query

gcloud/bigquery/test_table.py

@@ -524,6 +524,154 @@ def test_create_w_bound_client(self):
         self.assertEqual(req['data'], SENT)
         self._verifyResourceProperties(table, RESOURCE)
 
+    def test_create_w_partition_no_expire(self):
+        from gcloud.bigquery.table import SchemaField
+        PATH = 'projects/%s/datasets/%s/tables' % (self.PROJECT, self.DS_NAME)
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+
+        self.assertEqual(table.partitioning_type, None)
+        table.partitioning_type = "DAY"
+        self.assertEqual(table.partitioning_type, "DAY")
+        table.create()
+
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'POST')
+        self.assertEqual(req['path'], '/%s' % PATH)
+        SENT = {
+            'tableReference': {
+                'projectId': self.PROJECT,
+                'datasetId': self.DS_NAME,
+                'tableId': self.TABLE_NAME},
+            'timePartitioning': {'type': 'DAY'},
+            'schema': {'fields': [
+                {'name': 'full_name', 'type': 'STRING', 'mode': 'REQUIRED'},
+                {'name': 'age', 'type': 'INTEGER', 'mode': 'REQUIRED'}]},
+        }
+        self.assertEqual(req['data'], SENT)
+        self._verifyResourceProperties(table, RESOURCE)
+
+    def test_create_w_partition_and_expire(self):
+        from gcloud.bigquery.table import SchemaField
+        PATH = 'projects/%s/datasets/%s/tables' % (self.PROJECT, self.DS_NAME)
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+        self.assertEqual(table.partition_expiration, None)
+        table.partition_expiration = 100
+        self.assertEqual(table.partitioning_type, "DAY")
+        self.assertEqual(table.partition_expiration, 100)
+        table.create()
+
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'POST')
+        self.assertEqual(req['path'], '/%s' % PATH)
+        SENT = {
+            'tableReference': {
+                'projectId': self.PROJECT,
+                'datasetId': self.DS_NAME,
+                'tableId': self.TABLE_NAME},
+            'timePartitioning': {'type': 'DAY', 'expirationMs': 100},
+            'schema': {'fields': [
+                {'name': 'full_name', 'type': 'STRING', 'mode': 'REQUIRED'},
+                {'name': 'age', 'type': 'INTEGER', 'mode': 'REQUIRED'}]},
+        }
+        self.assertEqual(req['data'], SENT)
+        self._verifyResourceProperties(table, RESOURCE)
+
+    def test_partition_type_setter_none_type(self):
+        from gcloud.bigquery.table import SchemaField
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+        self.assertEqual(table.partitioning_type, None)
+        table.partitioning_type = None
+        self.assertEqual(table.partitioning_type, None)
+
+    def test_partition_type_setter_bad_type(self):
+        from gcloud.bigquery.table import SchemaField
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+        with self.assertRaises(ValueError):
+            table.partitioning_type = 123
+
+    def test_partition_type_setter_unknown_value(self):
+        from gcloud.bigquery.table import SchemaField
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+        with self.assertRaises(ValueError):
+            table.partitioning_type = "HASH"
+
+    def test_partition_experiation_bad_type(self):
+        from gcloud.bigquery.table import SchemaField
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+        with self.assertRaises(ValueError):
+            table.partition_expiration = "NEVER"
+
+    def test_partition_expiration(self):
+        from gcloud.bigquery.table import SchemaField
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+        self.assertEqual(table.partition_expiration, None)
+        table.partition_expiration = 100
+        self.assertEqual(table.partitioning_type, "DAY")
+        self.assertEqual(table.partition_expiration, 100)
+
+    def test_list_partitions(self):
+        from gcloud.bigquery.table import SchemaField
+        RESOURCE = self._makeResource()
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        table = self._makeOne(self.TABLE_NAME, dataset,
+                              schema=[full_name, age])
+        self.assertEqual(table.list_partitions(), [20160804, 20160805])
+
     def test_create_w_alternate_client(self):
         import datetime
         from gcloud._helpers import UTC
@@ -1669,6 +1817,18 @@ def __init__(self, project='project', connection=None):
     def job_from_resource(self, resource):  # pylint: disable=unused-argument
         return self._job
 
+    def run_sync_query(self, q=None):  # pylint: disable=unused-argument
+        return _Query(self)
+
+
+class _Query(object):
+
+    def __init__(self, client=None):  # pylint: disable=unused-argument
+        self.rows = []
+
+    def run(self):
+        self.rows = [(20160804, None), (20160805, None)]
+
 
 class _Dataset(object):