feat: add support for table snapshots

Author: plamut

google/cloud/bigquery/__init__.py

@@ -60,6 +60,7 @@
 from google.cloud.bigquery.job import ExtractJobConfig
 from google.cloud.bigquery.job import LoadJob
 from google.cloud.bigquery.job import LoadJobConfig
+from google.cloud.bigquery.job import OperationType
 from google.cloud.bigquery.job import QueryJob
 from google.cloud.bigquery.job import QueryJobConfig
 from google.cloud.bigquery.job import QueryPriority
@@ -86,6 +87,7 @@
 from google.cloud.bigquery.table import PartitionRange
 from google.cloud.bigquery.table import RangePartitioning
 from google.cloud.bigquery.table import Row
+from google.cloud.bigquery.table import SnapshotDefinition
 from google.cloud.bigquery.table import Table
 from google.cloud.bigquery.table import TableReference
 from google.cloud.bigquery.table import TimePartitioningType
@@ -114,6 +116,7 @@
     "PartitionRange",
     "RangePartitioning",
     "Row",
+    "SnapshotDefinition",
     "TimePartitioning",
     "TimePartitioningType",
     # Jobs
@@ -153,6 +156,7 @@
     "ExternalSourceFormat",
     "Encoding",
     "KeyResultStatementKind",
+    "OperationType",
     "QueryPriority",
     "SchemaUpdateOption",
     "SourceFormat",

google/cloud/bigquery/job/__init__.py

@@ -25,6 +25,7 @@
 from google.cloud.bigquery.job.base import UnknownJob
 from google.cloud.bigquery.job.copy_ import CopyJob
 from google.cloud.bigquery.job.copy_ import CopyJobConfig
+from google.cloud.bigquery.job.copy_ import OperationType
 from google.cloud.bigquery.job.extract import ExtractJob
 from google.cloud.bigquery.job.extract import ExtractJobConfig
 from google.cloud.bigquery.job.load import LoadJob
@@ -59,6 +60,7 @@
     "UnknownJob",
     "CopyJob",
     "CopyJobConfig",
+    "OperationType",
     "ExtractJob",
     "ExtractJobConfig",
     "LoadJob",

google/cloud/bigquery/job/copy_.py

@@ -14,6 +14,8 @@
 
 """Classes for copy jobs."""
 
+from typing import Optional
+
 from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 from google.cloud.bigquery import _helpers
 from google.cloud.bigquery.table import TableReference
@@ -23,6 +25,25 @@
 from google.cloud.bigquery.job.base import _JobReference
 
 
+class OperationType:
+    """Different operation types supported in table copy job.
+
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#operationtype
+    """
+
+    OPERATION_TYPE_UNSPECIFIED = "OPERATION_TYPE_UNSPECIFIED"
+    """Unspecified operation type."""
+
+    COPY = "COPY"
+    """The source and destination table have the same table type."""
+
+    SNAPSHOT = "SNAPSHOT"
+    """The source table type is TABLE and the destination table type is SNAPSHOT."""
+
+    RESTORE = "RESTORE"
+    """The source table type is SNAPSHOT and the destination table type is TABLE."""
+
+
 class CopyJobConfig(_JobConfig):
     """Configuration options for copy jobs.
 
@@ -85,6 +106,23 @@ def destination_encryption_configuration(self, value):
             api_repr = value.to_api_repr()
         self._set_sub_prop("destinationEncryptionConfiguration", api_repr)
 
+    @property
+    def operation_type(self) -> str:
+        """The operation to perform with this copy job.
+
+        See
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationTableCopy.FIELDS.operation_type
+        """
+        return self._get_sub_prop(
+            "operationType", OperationType.OPERATION_TYPE_UNSPECIFIED
+        )
+
+    @operation_type.setter
+    def operation_type(self, value: Optional[str]):
+        if value is None:
+            value = OperationType.OPERATION_TYPE_UNSPECIFIED
+        self._set_sub_prop("operationType", value)
+
 
 class CopyJob(_AsyncJob):
     """Asynchronous job: copy data into a table from other tables.

google/cloud/bigquery/table.py

@@ -321,6 +321,7 @@ class Table(object):
         "range_partitioning": "rangePartitioning",
         "time_partitioning": "timePartitioning",
         "schema": "schema",
+        "snapshot_definition": "snapshotDefinition",
         "streaming_buffer": "streamingBuffer",
         "self_link": "selfLink",
         "table_id": ["tableReference", "tableId"],
@@ -910,6 +911,19 @@ def external_data_configuration(self, value):
             self._PROPERTY_TO_API_FIELD["external_data_configuration"]
         ] = api_repr
 
+    @property
+    def snapshot_definition(self) -> Optional["SnapshotDefinition"]:
+        """Information about the snapshot. This value is set via snapshot creation.
+
+        See: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table.FIELDS.snapshot_definition
+        """
+        snapshot_info = self._properties.get(
+            self._PROPERTY_TO_API_FIELD["snapshot_definition"]
+        )
+        if snapshot_info is not None:
+            snapshot_info = SnapshotDefinition(snapshot_info)
+        return snapshot_info
+
     @classmethod
     def from_string(cls, full_table_id: str) -> "Table":
         """Construct a table from fully-qualified table ID.
@@ -1274,6 +1288,30 @@ def __init__(self, resource):
             )
 
 
+class SnapshotDefinition:
+    """Information about base table and snapshot time of the snapshot.
+
+    See https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#snapshotdefinition
+
+    Args:
+        resource: Snapshot definition representation returned from the API.
+    """
+
+    def __init__(self, resource: Dict[str, Any]):
+        self.base_table_reference = None
+        if "baseTableReference" in resource:
+            self.base_table_reference = TableReference.from_api_repr(
+                resource["baseTableReference"]
+            )
+
+        self.snapshot_time = None
+        if "snapshotTime" in resource:
+            dt = google.cloud._helpers._rfc3339_to_datetime(resource["snapshotTime"])
+            # The helper returns a timezone *aware* datetime object (with UTC tzinfo),
+            # we need to make it naive.
+            self.snapshot_time = dt.replace(tzinfo=None)
+
+
 class Row(object):
     """A BigQuery row.
 

tests/unit/job/test_copy.py

@@ -30,16 +30,22 @@ def _get_target_class():
 
     def test_ctor_w_properties(self):
         from google.cloud.bigquery.job import CreateDisposition
+        from google.cloud.bigquery.job import OperationType
         from google.cloud.bigquery.job import WriteDisposition
 
         create_disposition = CreateDisposition.CREATE_NEVER
         write_disposition = WriteDisposition.WRITE_TRUNCATE
+        snapshot_operation = OperationType.SNAPSHOT
+
         config = self._get_target_class()(
-            create_disposition=create_disposition, write_disposition=write_disposition
+            create_disposition=create_disposition,
+            write_disposition=write_disposition,
+            operation_type=snapshot_operation,
         )
 
         self.assertEqual(config.create_disposition, create_disposition)
         self.assertEqual(config.write_disposition, write_disposition)
+        self.assertEqual(config.operation_type, snapshot_operation)
 
     def test_to_api_repr_with_encryption(self):
         from google.cloud.bigquery.encryption_configuration import (
@@ -70,6 +76,20 @@ def test_to_api_repr_with_encryption_none(self):
             resource, {"copy": {"destinationEncryptionConfiguration": None}}
         )
 
+    def test_operation_type_unspecified(self):
+        from google.cloud.bigquery.job import OperationType
+
+        config = self._make_one()
+        self.assertEqual(
+            config.operation_type, OperationType.OPERATION_TYPE_UNSPECIFIED
+        )
+
+        # Setting it to None is the same as setting it to OPERATION_TYPE_UNSPECIFIED.
+        config.operation_type = None
+        self.assertEqual(
+            config.operation_type, OperationType.OPERATION_TYPE_UNSPECIFIED
+        )
+
 
 class TestCopyJob(_Base):
     JOB_TYPE = "copy"

tests/unit/test_table.py

@@ -684,6 +684,39 @@ def test_props_set_by_server(self):
         self.assertEqual(table.full_table_id, TABLE_FULL_ID)
         self.assertEqual(table.table_type, "TABLE")
 
+    def test_snapshot_definition_not_set(self):
+        dataset = DatasetReference(self.PROJECT, self.DS_ID)
+        table_ref = dataset.table(self.TABLE_NAME)
+        table = self._make_one(table_ref)
+
+        assert table.snapshot_definition is None
+
+    def test_snapshot_definition_set(self):
+        from google.cloud.bigquery.table import SnapshotDefinition
+
+        dataset = DatasetReference(self.PROJECT, self.DS_ID)
+        table_ref = dataset.table(self.TABLE_NAME)
+        table = self._make_one(table_ref)
+
+        table._properties["snapshotDefinition"] = {
+            "baseTableReference": {
+                "projectId": "project_x",
+                "datasetId": "dataset_y",
+                "tableId": "table_z",
+            },
+            "snapshotTime": "2010-09-28T10:20:30.123Z",
+        }
+
+        snapshot = table.snapshot_definition
+
+        assert isinstance(snapshot, SnapshotDefinition)
+        assert snapshot.base_table_reference.path == (
+            "/projects/project_x/datasets/dataset_y/tables/table_z"
+        )
+        assert snapshot.snapshot_time == datetime.datetime(
+            2010, 9, 28, 10, 20, 30, 123000
+        )
+
     def test_description_setter_bad_value(self):
         dataset = DatasetReference(self.PROJECT, self.DS_ID)
         table_ref = dataset.table(self.TABLE_NAME)
@@ -1509,6 +1542,45 @@ def test_to_api_repr(self):
         self.assertEqual(table.to_api_repr(), resource)
 
 
+class TestSnapshotDefinition:
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.bigquery.table import SnapshotDefinition
+
+        return SnapshotDefinition
+
+    @classmethod
+    def _make_one(cls, *args, **kwargs):
+        klass = cls._get_target_class()
+        return klass(*args, **kwargs)
+
+    def test_ctor_empty_resource(self):
+        instance = self._make_one(resource={})
+        assert instance.base_table_reference is None
+        assert instance.snapshot_time is None
+
+    def test_ctor_full_resource(self):
+        from google.cloud.bigquery.table import TableReference
+
+        resource = {
+            "baseTableReference": {
+                "projectId": "my-project",
+                "datasetId": "your-dataset",
+                "tableId": "our-table",
+            },
+            "snapshotTime": "2005-06-07T19:35:02.123Z",
+        }
+        instance = self._make_one(resource)
+
+        expected_table_ref = TableReference.from_string(
+            "my-project.your-dataset.our-table"
+        )
+        assert instance.base_table_reference == expected_table_ref
+
+        expected_time = datetime.datetime(2005, 6, 7, 19, 35, 2, 123000)
+        assert instance.snapshot_time == expected_time
+
+
 class TestRow(unittest.TestCase):
     def test_row(self):
         from google.cloud.bigquery.table import Row