BUG: Add support to replace partitions in date-partitioned tables (googleapis#47)

Author: mremes

docs/source/changelog.rst

@@ -14,6 +14,8 @@ Changelog
 - Use the `google-auth <https://google-auth.readthedocs.io/en/latest/>`__ library for authentication because ``oauth2client`` is deprecated. (:issue:`39`)
 - :func:`read_gbq` now has a ``auth_local_webserver`` boolean argument for controlling whether to use web server or console flow when getting user credentials. Replaces `--noauth_local_webserver` command line argument. (:issue:`35`)
 - :func:`read_gbq` now displays the BigQuery Job ID and standard price in verbose output. (:issue:`70` and :issue:`71`)
+- ``read_gbq`` now displays the BigQuery Job ID and standard price in verbose output. (:issue:`70` and :issue:`71`)
+- Add support to replace partitions in `date-partitioned tables <https://cloud.google.com/bigquery/docs/partitioned-tables>`__. Partition must be specified with a partition decorator separator (``$``). (:issue:`47`)
 
 0.1.6 / 2017-05-03
 ------------------

pandas_gbq/gbq.py

@@ -769,26 +769,6 @@ def schema_is_subset(self, dataset_id, table_id, schema):
 
         return all(field in fields_remote for field in fields_local)
 
-    def delete_and_recreate_table(self, dataset_id, table_id, table_schema):
-        delay = 0
-
-        # Changes to table schema may take up to 2 minutes as of May 2015 See
-        # `Issue 191
-        # <https://code.google.com/p/google-bigquery/issues/detail?id=191>`__
-        # Compare previous schema with new schema to determine if there should
-        # be a 120 second delay
-
-        if not self.verify_schema(dataset_id, table_id, table_schema):
-            self._print('The existing table has a different schema. Please '
-                        'wait 2 minutes. See Google BigQuery issue #191')
-            delay = 120
-
-        table = _Table(self.project_id, dataset_id,
-                       private_key=self.private_key)
-        table.delete(table_id)
-        table.create(table_id, table_schema)
-        sleep(delay)
-
 
 def _parse_data(schema, rows):
     # see:
@@ -1053,17 +1033,32 @@ def to_gbq(dataframe, destination_table, project_id, chunksize=10000,
                                      "already exists. "
                                      "Change the if_exists parameter to "
                                      "append or replace data.")
-        elif if_exists == 'replace':
-            connector.delete_and_recreate_table(
-                dataset_id, table_id, table_schema)
-        elif if_exists == 'append':
+        else:
+            delay = 0
             if not connector.schema_is_subset(dataset_id,
                                               table_id,
                                               table_schema):
-                raise InvalidSchema("Please verify that the structure and "
-                                    "data types in the DataFrame match the "
-                                    "schema of the destination table.")
+                if if_exists == 'append' \
+                        or table.partition_decorator in table_id:
+                    raise InvalidSchema("Please verify that the structure "
+                                        "and data types in the DataFrame "
+                                        "match the schema of the destination "
+                                        "table.")
+                elif if_exists == 'replace':
+                    table._print('The existing table has a different schema. '
+                                 'Please wait 2 minutes. See Google BigQuery '
+                                 'issue #191')
+                    delay = 120
+            if if_exists == 'replace':
+                table.delete(table_id)
+                if table.partition_decorator not in table_id:
+                    table.create(table_id, table_schema)
+                    sleep(delay)
+
     else:
+        if table.partition_decorator in table_id:
+            raise TableCreationError("Cannot create a partition without the "
+                                     "main table.")
         table.create(table_id, table_schema)
 
     connector.load_data(dataframe, dataset_id, table_id, chunksize)
@@ -1108,6 +1103,8 @@ def _generate_bq_schema(df, default_type='STRING'):
 
 class _Table(GbqConnector):
 
+    partition_decorator = '$'
+
     def __init__(self, project_id, dataset_id, reauth=False, verbose=False,
                  private_key=None):
         try:
@@ -1144,7 +1141,7 @@ def exists(self, table_id):
             else:
                 self.process_http_error(ex)
 
-    def create(self, table_id, schema):
+    def create(self, table_id, schema, date_partitioned=False):
         """ Create a table in Google BigQuery given a table and schema
 
         Parameters
@@ -1154,6 +1151,8 @@ def create(self, table_id, schema):
         schema : str
             Use the generate_bq_schema to generate your table schema from a
             dataframe.
+        date_partitioned: boolean
+            Whether table is to be created as a date partitioned table.
         """
 
         if self.exists(table_id):
@@ -1173,6 +1172,9 @@ def create(self, table_id, schema):
                 'datasetId': self.dataset_id
             }
         }
+        if date_partitioned:
+            # The only type supported is DAY
+            body.update({'timePartitioning': {'type': 'DAY'}})
 
         try:
             self.service.tables().insert(

pandas_gbq/tests/test_gbq.py

@@ -20,6 +20,7 @@
 
 
 TABLE_ID = 'new_test'
+DPT_TABLE_ID = 'dpt_test'
 
 
 def _skip_if_no_project_id():
@@ -946,6 +947,8 @@ def setup_method(self, method):
                                     private_key=_get_private_key_path())
         self.destination_table = "{0}{1}.{2}".format(self.dataset_prefix, "1",
                                                      TABLE_ID)
+        self.destination_dpt = "{0}{1}.{2}".format(self.dataset_prefix, "1",
+                                                   DPT_TABLE_ID)
         self.dataset.create(self.dataset_prefix + "1")
 
     @classmethod
@@ -1080,6 +1083,94 @@ def test_upload_data_if_table_exists_replace(self):
                               private_key=_get_private_key_path())
         assert result['num_rows'][0] == 5
 
+    def test_upload_data_if_table_exists_replace_dpt_partition(self):
+        # Issue #47; tests that 'replace' is done by the subsequent call
+        test_dpt_suffix = "20170101"
+        test_size = 10
+        df = make_mixed_dataframe_v2(test_size)
+        df_different_schema = tm.makeMixedDataFrame()
+
+        dpt_partition = self.destination_dpt + '$' + test_dpt_suffix
+        self.table.create(DPT_TABLE_ID, gbq._generate_bq_schema(df))
+        gbq.to_gbq(df, dpt_partition, _get_project_id(),
+                   private_key=_get_private_key_path())
+
+        gbq.to_gbq(df_different_schema, dpt_partition,
+                   _get_project_id(), if_exists='replace',
+                   private_key=_get_private_key_path())
+
+        sleep(30)
+
+        # Test whole table
+        result0 = gbq.read_gbq("SELECT COUNT(*) AS num_rows FROM {0}"
+                               .format(self.destination_dpt),
+                               project_id=_get_project_id(),
+                               private_key=_get_private_key_path())
+        assert result0['num_rows'][0] == 5
+
+        # Test destination partition
+        result1 = gbq.read_gbq("SELECT COUNT(*) AS num_rows FROM {0}"
+                               .format(dpt_partition),
+                               project_id=_get_project_id(),
+                               private_key=_get_private_key_path())
+        assert result1['num_rows'][0] == 5
+        self.table.delete(DPT_TABLE_ID)
+
+    def test_upload_data_if_table_exists_append_dpt_partition(self):
+        # Issue #47; tests that 'append' appends to an existing partition
+        test_dpt_suffix = "20170101"
+        test_size = 10
+        df = make_mixed_dataframe_v2(test_size)
+
+        dpt_partition = self.destination_dpt + '$' + test_dpt_suffix
+        self.table.create(DPT_TABLE_ID, gbq._generate_bq_schema(df))
+
+        result0 = gbq.read_gbq("SELECT COUNT(*) AS num_rows FROM {0}"
+                               .format(dpt_partition),
+                               project_id=_get_project_id(),
+                               private_key=_get_private_key_path())
+        assert result0['num_rows'][0] == 5
+
+        gbq.to_gbq(df, dpt_partition,
+                   _get_project_id(), if_exists='append',
+                   private_key=_get_private_key_path())
+
+        result1 = gbq.read_gbq("SELECT COUNT(*) AS num_rows FROM {0}"
+                               .format(dpt_partition),
+                               project_id=_get_project_id(),
+                               private_key=_get_private_key_path())
+
+        assert result1['num_rows'][0] == 15
+
+        sleep(30)
+
+        # Test whole table
+        result0 = gbq.read_gbq("SELECT COUNT(*) AS num_rows FROM {0}"
+                               .format(self.destination_dpt),
+                               project_id=_get_project_id(),
+                               private_key=_get_private_key_path())
+        assert result0['num_rows'][0] == 5
+
+        # Test destination partition
+        result1 = gbq.read_gbq("SELECT COUNT(*) AS num_rows FROM {0}"
+                               .format(dpt_partition),
+                               project_id=_get_project_id(),
+                               private_key=_get_private_key_path())
+        assert result1['num_rows'][0] == 10
+        self.table.delete(DPT_TABLE_ID)
+
+    def test_table_creation_error_raised_when_dpt_does_not_exist(self):
+        test_dpt_suffix = "20170101"
+        test_size = 10
+        df = make_mixed_dataframe_v2(test_size)
+        dpt_partition = 'dataset.foobar' + '$' + test_dpt_suffix
+
+        with pytest.raises(gbq.TableCreationError):
+            gbq.to_gbq(df,
+                       dpt_partition,
+                       project_id=_get_project_id(),
+                       private_key=_get_private_key_path())
+
     def test_upload_data_if_table_exists_raises_value_error(self):
         test_id = "4"
         test_size = 10