fix: adding property setter for table constraints, #1990 (#2092)

lkhagvadorj-amp · web-flow · commit f8572dd86595 · 2025-03-04T10:50:05.000-08:00
* fix: adding property setter for table constraints, #1990 * fix: adding unit test for to_api_repr() method * fix: adding system test for bigquery table update * fix: adding more test case for table constraints property * fix: adding more test case for table constraints property * fix: linting code * fix: adding unit tests for test table, table constraint and foreign key * fix: linting based on noxfile * fix: linting based on noxfile * fix: adding unit tests and system test * fix: clearing lint error * fix: adding table constraint eq unit tests * fix: adding type to to_api_repr resource
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
@@ -1058,6 +1058,17 @@ def table_constraints(self) -> Optional["TableConstraints"]:
             table_constraints = TableConstraints.from_api_repr(table_constraints)
         return table_constraints
 
+    @table_constraints.setter
+    def table_constraints(self, value):
+        """Tables Primary Key and Foreign Key information."""
+        api_repr = value
+        if not isinstance(value, TableConstraints) and value is not None:
+            raise ValueError(
+                "value must be google.cloud.bigquery.table.TableConstraints or None"
+            )
+        api_repr = value.to_api_repr() if value else None
+        self._properties[self._PROPERTY_TO_API_FIELD["table_constraints"]] = api_repr
+
     @property
     def resource_tags(self):
         """Dict[str, str]: Resource tags for the table.
@@ -1111,11 +1122,9 @@ def external_catalog_table_options(
     def foreign_type_info(self) -> Optional[_schema.ForeignTypeInfo]:
         """Optional. Specifies metadata of the foreign data type definition in
         field schema (TableFieldSchema.foreign_type_definition).
-
         Returns:
             Optional[schema.ForeignTypeInfo]:
                 Foreign type information, or :data:`None` if not set.
-
         .. Note::
             foreign_type_info is only required if you are referencing an
             external catalog such as a Hive table.
@@ -3404,6 +3413,20 @@ def from_api_repr(cls, api_repr: Dict[str, Any]) -> "ForeignKey":
             ],
         )
 
+    def to_api_repr(self) -> Dict[str, Any]:
+        """Return a dictionary representing this object."""
+        return {
+            "name": self.name,
+            "referencedTable": self.referenced_table.to_api_repr(),
+            "columnReferences": [
+                {
+                    "referencingColumn": column_reference.referencing_column,
+                    "referencedColumn": column_reference.referenced_column,
+                }
+                for column_reference in self.column_references
+            ],
+        }
+
 
 class TableConstraints:
     """The TableConstraints defines the primary key and foreign key.
@@ -3425,6 +3448,13 @@ def __init__(
         self.primary_key = primary_key
         self.foreign_keys = foreign_keys
 
+    def __eq__(self, other):
+        if not isinstance(other, TableConstraints) and other is not None:
+            raise TypeError("The value provided is not a BigQuery TableConstraints.")
+        return (
+            self.primary_key == other.primary_key if other.primary_key else None
+        ) and (self.foreign_keys == other.foreign_keys if other.foreign_keys else None)
+
     @classmethod
     def from_api_repr(cls, resource: Dict[str, Any]) -> "TableConstraints":
         """Create an instance from API representation."""
@@ -3440,6 +3470,17 @@ def from_api_repr(cls, resource: Dict[str, Any]) -> "TableConstraints":
             ]
         return cls(primary_key, foreign_keys)
 
+    def to_api_repr(self) -> Dict[str, Any]:
+        """Return a dictionary representing this object."""
+        resource: Dict[str, Any] = {}
+        if self.primary_key:
+            resource["primaryKey"] = {"columns": self.primary_key.columns}
+        if self.foreign_keys:
+            resource["foreignKeys"] = [
+                foreign_key.to_api_repr() for foreign_key in self.foreign_keys
+            ]
+        return resource
+
 
 def _item_to_row(iterator, resource):
     """Convert a JSON row to the native object.
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
@@ -97,6 +97,10 @@
         ],
     ),
 ]
+TABLE_CONSTRAINTS_SCHEMA = [
+    bigquery.SchemaField("id", "INTEGER", mode="REQUIRED"),
+    bigquery.SchemaField("fk_id", "STRING", mode="REQUIRED"),
+]
 
 SOURCE_URIS_AVRO = [
     "gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/a-twitter.avro",
@@ -901,6 +905,79 @@ def test_update_table_clustering_configuration(self):
         table3 = Config.CLIENT.update_table(table2, ["clustering_fields"])
         self.assertIsNone(table3.clustering_fields, None)
 
+    def test_update_table_constraints(self):
+        from google.cloud.bigquery.table import TableConstraints
+        from google.cloud.bigquery.table import (
+            PrimaryKey,
+            ForeignKey,
+            TableReference,
+            ColumnReference,
+        )
+
+        dataset = self.temp_dataset(_make_dataset_id("update_table"))
+
+        TABLE_NAME = "test_table"
+        table_arg = Table(dataset.table(TABLE_NAME), schema=TABLE_CONSTRAINTS_SCHEMA)
+        self.assertFalse(_table_exists(table_arg))
+
+        table = helpers.retry_403(Config.CLIENT.create_table)(table_arg)
+        self.to_delete.insert(0, table)
+        self.assertTrue(_table_exists(table))
+
+        REFERENCE_TABLE_NAME = "test_table2"
+        reference_table_arg = Table(
+            dataset.table(REFERENCE_TABLE_NAME),
+            schema=[
+                bigquery.SchemaField("id", "INTEGER", mode="REQUIRED"),
+            ],
+        )
+        reference_table = helpers.retry_403(Config.CLIENT.create_table)(
+            reference_table_arg
+        )
+        self.to_delete.insert(0, reference_table)
+        self.assertTrue(_table_exists(reference_table))
+
+        reference_table.table_constraints = TableConstraints(
+            primary_key=PrimaryKey(columns=["id"]), foreign_keys=None
+        )
+        reference_table2 = Config.CLIENT.update_table(
+            reference_table, ["table_constraints"]
+        )
+        self.assertEqual(
+            reference_table2.table_constraints.primary_key,
+            reference_table.table_constraints.primary_key,
+        )
+
+        table_constraints = TableConstraints(
+            primary_key=PrimaryKey(columns=["id"]),
+            foreign_keys=[
+                ForeignKey(
+                    name="fk_id",
+                    referenced_table=TableReference(dataset, "test_table2"),
+                    column_references=[
+                        ColumnReference(referencing_column="id", referenced_column="id")
+                    ],
+                ),
+            ],
+        )
+
+        table.table_constraints = table_constraints
+        table2 = Config.CLIENT.update_table(table, ["table_constraints"])
+        self.assertEqual(
+            table2.table_constraints,
+            table_constraints,
+        )
+
+        table2.table_constraints = None
+        table3 = Config.CLIENT.update_table(table2, ["table_constraints"])
+        self.assertIsNone(table3.table_constraints, None)
+
+        reference_table2.table_constraints = None
+        reference_table3 = Config.CLIENT.update_table(
+            reference_table2, ["table_constraints"]
+        )
+        self.assertIsNone(reference_table3.table_constraints, None)
+
     @staticmethod
     def _fetch_single_page(table, selected_fields=None):
         iterator = Config.CLIENT.list_rows(table, selected_fields=selected_fields)
diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py
