Add system test for table snapshots

Author: plamut

tests/system/test_client.py

@@ -2316,6 +2316,75 @@ def test_parameterized_types_round_trip(self):
 
         self.assertEqual(tuple(s._key()[:2] for s in table2.schema), fields)
 
+    def test_table_snapshots(self):
+        from google.cloud.bigquery import CopyJobConfig
+        from google.cloud.bigquery import OperationType
+
+        client = Config.CLIENT
+
+        source_table_path = f"{client.project}.{Config.DATASET}.test_table"
+        snapshot_table_path = f"{source_table_path}_snapshot"
+
+        # Create the table before loading so that the column order is predictable.
+        schema = [
+            bigquery.SchemaField("foo", "INTEGER"),
+            bigquery.SchemaField("bar", "STRING"),
+        ]
+        source_table = helpers.retry_403(Config.CLIENT.create_table)(
+            Table(source_table_path, schema=schema)
+        )
+        self.to_delete.insert(0, source_table)
+
+        # Populate the table with initial data.
+        rows = [{"foo": 1, "bar": "one"}, {"foo": 2, "bar": "two"}]
+        load_job = Config.CLIENT.load_table_from_json(rows, source_table)
+        load_job.result()
+
+        # Now create a snapshot before modifying the original table data.
+        copy_config = CopyJobConfig()
+        copy_config.operation_type = OperationType.SNAPSHOT
+
+        copy_job = client.copy_table(
+            sources=source_table_path,
+            destination=snapshot_table_path,
+            job_config=copy_config,
+        )
+        copy_job.result()
+
+        snapshot_table = client.get_table(snapshot_table_path)
+        self.to_delete.insert(0, snapshot_table)
+
+        # Modify data in original table.
+        sql = f'INSERT INTO `{source_table_path}`(foo, bar) VALUES (3, "three")'
+        query_job = client.query(sql)
+        query_job.result()
+
+        # List rows from the source table and compare them to rows from the snapshot.
+        rows_iter = client.list_rows(source_table_path)
+        rows = sorted(row.values() for row in rows_iter)
+        assert rows == [(1, "one"), (2, "two"), (3, "three")]
+
+        rows_iter = client.list_rows(snapshot_table_path)
+        rows = sorted(row.values() for row in rows_iter)
+        assert rows == [(1, "one"), (2, "two")]
+
+        # Now restore the table from the snapshot and it should again contain the old
+        # set of rows.
+        copy_config = CopyJobConfig()
+        copy_config.operation_type = OperationType.RESTORE
+        copy_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
+
+        copy_job = client.copy_table(
+            sources=snapshot_table_path,
+            destination=source_table_path,
+            job_config=copy_config,
+        )
+        copy_job.result()
+
+        rows_iter = client.list_rows(source_table_path)
+        rows = sorted(row.values() for row in rows_iter)
+        assert rows == [(1, "one"), (2, "two")]
+
     def temp_dataset(self, dataset_id, location=None):
         project = Config.CLIENT.project
         dataset_ref = bigquery.DatasetReference(project, dataset_id)