From a167f9a95f0a8fbf0bdb4943d06f07c03768c132 Mon Sep 17 00:00:00 2001 From: Dmytro Karacheban Date: Thu, 11 Jan 2024 18:37:26 +0200 Subject: [PATCH] feat: Add `table_constraints` field to Table model (#1755) * feat: add `table_constraints` field to Table model * Change `raise` to `return` in __eq__ methods * Fix __eq__ for ColumnReference * Add column_references to ForeignKey __eq__ * Add missing coverage * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py --------- Co-authored-by: Chalmer Lowe Co-authored-by: Chalmer Lowe --- google/cloud/bigquery/table.py | 128 +++++++++++++++ tests/unit/test_table.py | 280 +++++++++++++++++++++++++++++++++ 2 files changed, 408 insertions(+) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index 0ae7851a1..b3be4ff90 100644 --- a/google/cloud/bigquery/table.py +++ b/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. @@ -2958,6 +2969,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 TypeError("The value provided is not a BigQuery PrimaryKey.") + 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 TypeError("The value provided is not a BigQuery ColumnReference.") + return ( + self.referencing_column == other.referencing_column + and self.referenced_column == other.referenced_column + ) + + +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 TypeError("The value provided is not a BigQuery ForeignKey.") + return ( + self.name == other.name + and self.referenced_table == other.referenced_table + and self.column_references == other.column_references + ) + + @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( + column_reference_resource["referencingColumn"], + column_reference_resource["referencedColumn"], + ) + 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. diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 4a85a0823..e4d0c66ab 100644 --- a/tests/unit/test_table.py +++ b/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) @@ -5393,6 +5409,270 @@ def test_set_expiration_w_none(self): assert time_partitioning._properties["expirationMs"] is None +class TestPrimaryKey(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigquery.table import PrimaryKey + + return PrimaryKey + + @classmethod + def _make_one(cls, *args, **kwargs): + return cls._get_target_class()(*args, **kwargs) + + def test_constructor_explicit(self): + columns = ["id", "product_id"] + primary_key = self._make_one(columns) + + self.assertEqual(primary_key.columns, columns) + + def test__eq__columns_mismatch(self): + primary_key = self._make_one(columns=["id", "product_id"]) + other_primary_key = self._make_one(columns=["id"]) + + self.assertNotEqual(primary_key, other_primary_key) + + def test__eq__other_type(self): + primary_key = self._make_one(columns=["id", "product_id"]) + with self.assertRaises(TypeError): + primary_key == "This is not a Primary Key" + + +class TestColumnReference(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigquery.table import ColumnReference + + return ColumnReference + + @classmethod + def _make_one(cls, *args, **kwargs): + return cls._get_target_class()(*args, **kwargs) + + def test_constructor_explicit(self): + referencing_column = "product_id" + referenced_column = "id" + column_reference = self._make_one(referencing_column, referenced_column) + + self.assertEqual(column_reference.referencing_column, referencing_column) + self.assertEqual(column_reference.referenced_column, referenced_column) + + def test__eq__referencing_column_mismatch(self): + column_reference = self._make_one( + referencing_column="product_id", + referenced_column="id", + ) + other_column_reference = self._make_one( + referencing_column="item_id", + referenced_column="id", + ) + + self.assertNotEqual(column_reference, other_column_reference) + + def test__eq__referenced_column_mismatch(self): + column_reference = self._make_one( + referencing_column="product_id", + referenced_column="id", + ) + other_column_reference = self._make_one( + referencing_column="product_id", + referenced_column="id_1", + ) + + self.assertNotEqual(column_reference, other_column_reference) + + def test__eq__other_type(self): + column_reference = self._make_one( + referencing_column="product_id", + referenced_column="id", + ) + with self.assertRaises(TypeError): + column_reference == "This is not a Column Reference" + + +class TestForeignKey(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigquery.table import ForeignKey + + return ForeignKey + + @classmethod + def _make_one(cls, *args, **kwargs): + return cls._get_target_class()(*args, **kwargs) + + def test_constructor_explicit(self): + name = "my_fk" + referenced_table = TableReference.from_string("my-project.mydataset.mytable") + column_references = [] + foreign_key = self._make_one(name, referenced_table, column_references) + + self.assertEqual(foreign_key.name, name) + self.assertEqual(foreign_key.referenced_table, referenced_table) + self.assertEqual(foreign_key.column_references, column_references) + + def test__eq__name_mismatch(self): + referenced_table = TableReference.from_string("my-project.mydataset.mytable") + column_references = [] + foreign_key = self._make_one( + name="my_fk", + referenced_table=referenced_table, + column_references=column_references, + ) + other_foreign_key = self._make_one( + name="my_other_fk", + referenced_table=referenced_table, + column_references=column_references, + ) + + self.assertNotEqual(foreign_key, other_foreign_key) + + def test__eq__referenced_table_mismatch(self): + name = "my_fk" + column_references = [] + foreign_key = self._make_one( + name=name, + referenced_table=TableReference.from_string("my-project.mydataset.mytable"), + column_references=column_references, + ) + other_foreign_key = self._make_one( + name=name, + referenced_table=TableReference.from_string( + "my-project.mydataset.my-other-table" + ), + column_references=column_references, + ) + + self.assertNotEqual(foreign_key, other_foreign_key) + + def test__eq__column_references_mismatch(self): + from google.cloud.bigquery.table import ColumnReference + + name = "my_fk" + referenced_table = TableReference.from_string("my-project.mydataset.mytable") + foreign_key = self._make_one( + name=name, + referenced_table=referenced_table, + column_references=[], + ) + other_foreign_key = self._make_one( + name=name, + referenced_table=referenced_table, + column_references=[ + ColumnReference( + referencing_column="product_id", referenced_column="id" + ), + ], + ) + + self.assertNotEqual(foreign_key, other_foreign_key) + + def test__eq__other_type(self): + foreign_key = self._make_one( + name="my_fk", + referenced_table=TableReference.from_string("my-project.mydataset.mytable"), + column_references=[], + ) + with self.assertRaises(TypeError): + foreign_key == "This is not a Foreign Key" + + +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`" )