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

Author: mremes

docs/source/changelog.rst

@@ -6,6 +6,7 @@ Changelog
 
 - Resolve issue where the optional ``--noauth_local_webserver`` command line argument would not be propagated during the authentication process. (:issue:`35`)
 - Drop support for Python 3.4 (:issue:`40`)
+- Add support to replace partitions in date-partitioned tables (:issue:`47`)
 
 0.1.6 / 2017-05-03
 ------------------

pandas_gbq/gbq.py

@@ -840,17 +840,27 @@ 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.verify_schema(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)
 
 
@@ -893,6 +903,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:

pandas_gbq/tests/test_gbq.py

@@ -21,6 +21,7 @@
 
 
 TABLE_ID = 'new_test'
+DPT_TABLE_ID = 'dpt_test'
 
 
 _IMPORTS = False
@@ -407,7 +408,7 @@ def test_should_return_bigquery_strings_as_python_strings(self):
     def test_to_gbq_should_fail_if_invalid_table_name_passed(self):
         with tm.assertRaises(gbq.NotFoundException):
             gbq.to_gbq(DataFrame(), 'invalid_table_name', project_id="1234")
-
+    
     def test_to_gbq_with_no_project_id_given_should_fail(self):
         with tm.assertRaises(TypeError):
             gbq.to_gbq(DataFrame(), 'dataset.tablename')
@@ -996,6 +997,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_date_partitioned_table = "{0}{1}.{2}".format(self.dataset_prefix, "1",
+                                                                      DPT_TABLE_ID)
         self.dataset.create(self.dataset_prefix + "1")
 
     @classmethod
@@ -1094,6 +1097,79 @@ 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_date_partitioned_table + '$' + test_dpt_suffix
+
+        gbq.to_gbq(df, dpt_partition, _get_project_id(),
+                   chunksize=10000, 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_date_partitioned_table),
+                               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
+
+    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_date_partitioned_table + '$' + test_dpt_suffix
+
+        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_date_partitioned_table),
+                               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
+
     def test_upload_data_if_table_exists_raises_value_error(self):
         test_id = "4"
         test_size = 10
@@ -1117,7 +1193,7 @@ def test_google_upload_errors_should_raise_exception(self):
         with tm.assertRaises(gbq.StreamingInsertError):
             gbq.to_gbq(bad_df, self.destination_table + test_id,
                        _get_project_id(), private_key=_get_private_key_path())
-
+    
     def test_generate_schema(self):
         df = tm.makeMixedDataFrame()
         schema = gbq._generate_bq_schema(df)