fix: adding property setter for table constraints, #1990

google/cloud/bigquery/table.py

@@ -1105,11 +1105,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.
@@ -1138,6 +1136,17 @@ def foreign_type_info(self, value: Union[_schema.ForeignTypeInfo, dict, None]):
             self._properties, self._PROPERTY_TO_API_FIELD["foreign_type_info"], value
         )
 
+    @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
+
     @classmethod
     def from_string(cls, full_table_id: str) -> "Table":
         """Construct a table from fully-qualified table ID.
@@ -3398,6 +3407,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.
@@ -3419,6 +3442,17 @@ 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."""
@@ -3434,6 +3468,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 = {}
+        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.

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)