Add table_constraints field to Table model

Author: dimkonko

google/cloud/bigquery/table.py

@@ -390,6 +390,7 @@ class Table(_TableBase):
         "view_use_legacy_sql": "view",
         "view_query": "view",
         "require_partition_filter": "requirePartitionFilter",
+        "table_constraints": "tableConstraints",
     }
 
     def __init__(self, table_ref, schema=None) -> None:
@@ -973,6 +974,16 @@ def clone_definition(self) -> Optional["CloneDefinition"]:
             clone_info = CloneDefinition(clone_info)
         return clone_info
 
+    @property
+    def table_constraints(self) -> Optional["TableConstraints"]:
+        """Tables Primary Key and Foreign Key information."""
+        table_constraints = self._properties.get(
+            self._PROPERTY_TO_API_FIELD["table_constraints"]
+        )
+        if table_constraints is not None:
+            table_constraints = TableConstraints.from_api_repr(table_constraints)
+        return table_constraints
+
     @classmethod
     def from_string(cls, full_table_id: str) -> "Table":
         """Construct a table from fully-qualified table ID.
@@ -2942,6 +2953,123 @@ def __repr__(self):
         return "TimePartitioning({})".format(",".join(key_vals))
 
 
+class PrimaryKey:
+    """Represents the primary key constraint on a table's columns.
+
+    Args:
+        columns: The columns that are composed of the primary key constraint.
+    """
+
+    def __init__(self, columns: List[str]):
+        self.columns = columns
+
+    def __eq__(self, other):
+        if not isinstance(other, PrimaryKey):
+            raise NotImplementedError
+        return self.columns == other.columns
+
+
+class ColumnReference:
+    """The pair of the foreign key column and primary key column.
+
+    Args:
+        referencing_column: The column that composes the foreign key.
+        referenced_column: The column in the primary key that are referenced by the referencingColumn.
+    """
+
+    def __init__(self, referencing_column: str, referenced_column: str):
+        self.referencing_column = referencing_column
+        self.referenced_column = referenced_column
+
+    def __eq__(self, other):
+        if not isinstance(other, ColumnReference):
+            raise NotImplementedError
+        return (
+            self.referenced_column == other.referencing_column
+            and self.referenced_column == other.referenced_column
+        )
+
+    @classmethod
+    def from_api_repr(cls, api_repr: Dict[str, Any]) -> "ColumnReference":
+        """Create an instance from API representation."""
+        return cls(api_repr["referencingColumn"], api_repr["referencedColumn"])
+
+
+class ForeignKey:
+    """Represents a foreign key constraint on a table's columns.
+
+    Args:
+        name: Set only if the foreign key constraint is named.
+        referenced_table: The table that holds the primary key and is referenced by this foreign key.
+        column_references: The columns that compose the foreign key.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        referenced_table: TableReference,
+        column_references: List[ColumnReference],
+    ):
+        self.name = name
+        self.referenced_table = referenced_table
+        self.column_references = column_references
+
+    def __eq__(self, other):
+        if not isinstance(other, ForeignKey):
+            raise NotImplementedError
+        return (
+            self.name == other.name and self.referenced_table == other.referenced_table
+        )
+
+    @classmethod
+    def from_api_repr(cls, api_repr: Dict[str, Any]) -> "ForeignKey":
+        """Create an instance from API representation."""
+        return cls(
+            name=api_repr["name"],
+            referenced_table=TableReference.from_api_repr(api_repr["referencedTable"]),
+            column_references=[
+                ColumnReference.from_api_repr(column_reference_resource)
+                for column_reference_resource in api_repr["columnReferences"]
+            ],
+        )
+
+
+class TableConstraints:
+    """The TableConstraints defines the primary key and foreign key.
+
+    Args:
+        primary_key:
+            Represents a primary key constraint on a table's columns. Present only if the table
+            has a primary key. The primary key is not enforced.
+        foreign_keys:
+            Present only if the table has a foreign key. The foreign key is not enforced.
+
+    """
+
+    def __init__(
+        self,
+        primary_key: Optional[PrimaryKey],
+        foreign_keys: Optional[List[ForeignKey]],
+    ):
+        self.primary_key = primary_key
+        self.foreign_keys = foreign_keys
+
+    @classmethod
+    def from_api_repr(cls, resource: Dict[str, Any]) -> "TableConstraints":
+        """Create an instance from API representation."""
+        primary_key = None
+        if "primaryKey" in resource:
+            primary_key = PrimaryKey(resource["primaryKey"]["columns"])
+
+        foreign_keys = None
+        if "foreignKeys" in resource:
+            foreign_keys = [
+                ForeignKey.from_api_repr(foreign_key_resource)
+                for foreign_key_resource in resource["foreignKeys"]
+            ]
+        return cls(primary_key, foreign_keys)
+
+
 def _item_to_row(iterator, resource):
     """Convert a JSON row to the native object.
 

tests/unit/test_table.py

@@ -603,6 +603,7 @@ def test_ctor(self):
         self.assertIsNone(table.encryption_configuration)
         self.assertIsNone(table.time_partitioning)
         self.assertIsNone(table.clustering_fields)
+        self.assertIsNone(table.table_constraints)
 
     def test_ctor_w_schema(self):
         from google.cloud.bigquery.schema import SchemaField
@@ -901,6 +902,21 @@ def test_clone_definition_set(self):
             2010, 9, 28, 10, 20, 30, 123000, tzinfo=UTC
         )
 
+    def test_table_constraints_property_getter(self):
+        from google.cloud.bigquery.table import PrimaryKey, TableConstraints
+
+        dataset = DatasetReference(self.PROJECT, self.DS_ID)
+        table_ref = dataset.table(self.TABLE_NAME)
+        table = self._make_one(table_ref)
+        table._properties["tableConstraints"] = {
+            "primaryKey": {"columns": ["id"]},
+        }
+
+        table_constraints = table.table_constraints
+
+        assert isinstance(table_constraints, TableConstraints)
+        assert table_constraints.primary_key == PrimaryKey(columns=["id"])
+
     def test_description_setter_bad_value(self):
         dataset = DatasetReference(self.PROJECT, self.DS_ID)
         table_ref = dataset.table(self.TABLE_NAME)
@@ -5385,6 +5401,102 @@ def test_set_expiration_w_none(self):
         assert time_partitioning._properties["expirationMs"] is None
 
 
+class TestTableConstraint(unittest.TestCase):
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.bigquery.table import TableConstraints
+
+        return TableConstraints
+
+    @classmethod
+    def _make_one(cls, *args, **kwargs):
+        return cls._get_target_class()(*args, **kwargs)
+
+    def test_constructor_defaults(self):
+        instance = self._make_one(primary_key=None, foreign_keys=None)
+        self.assertIsNone(instance.primary_key)
+        self.assertIsNone(instance.foreign_keys)
+
+    def test_from_api_repr_full_resource(self):
+        from google.cloud.bigquery.table import (
+            ColumnReference,
+            ForeignKey,
+            TableReference,
+        )
+
+        resource = {
+            "primaryKey": {
+                "columns": ["id", "product_id"],
+            },
+            "foreignKeys": [
+                {
+                    "name": "my_fk_name",
+                    "referencedTable": {
+                        "projectId": "my-project",
+                        "datasetId": "your-dataset",
+                        "tableId": "products",
+                    },
+                    "columnReferences": [
+                        {"referencingColumn": "product_id", "referencedColumn": "id"},
+                    ],
+                }
+            ],
+        }
+        instance = self._get_target_class().from_api_repr(resource)
+
+        self.assertIsNotNone(instance.primary_key)
+        self.assertEqual(instance.primary_key.columns, ["id", "product_id"])
+        self.assertEqual(
+            instance.foreign_keys,
+            [
+                ForeignKey(
+                    name="my_fk_name",
+                    referenced_table=TableReference.from_string(
+                        "my-project.your-dataset.products"
+                    ),
+                    column_references=[
+                        ColumnReference(
+                            referencing_column="product_id", referenced_column="id"
+                        ),
+                    ],
+                ),
+            ],
+        )
+
+    def test_from_api_repr_only_primary_key_resource(self):
+        resource = {
+            "primaryKey": {
+                "columns": ["id"],
+            },
+        }
+        instance = self._get_target_class().from_api_repr(resource)
+
+        self.assertIsNotNone(instance.primary_key)
+        self.assertEqual(instance.primary_key.columns, ["id"])
+        self.assertIsNone(instance.foreign_keys)
+
+    def test_from_api_repr_only_foreign_keys_resource(self):
+        resource = {
+            "foreignKeys": [
+                {
+                    "name": "my_fk_name",
+                    "referencedTable": {
+                        "projectId": "my-project",
+                        "datasetId": "your-dataset",
+                        "tableId": "products",
+                    },
+                    "columnReferences": [
+                        {"referencingColumn": "product_id", "referencedColumn": "id"},
+                    ],
+                }
+            ]
+        }
+        instance = self._get_target_class().from_api_repr(resource)
+
+        self.assertIsNone(instance.primary_key)
+        self.assertIsNotNone(instance.foreign_keys)
+
+
 @pytest.mark.skipif(
     bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`"
 )