fix: adding property setter for table constraints, #1990

google/cloud/bigquery/table.py

@@ -1023,6 +1023,18 @@ 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 isinstance(value, TableConstraints):
+            api_repr = value.to_api_repr()
+        elif value is not None:
+            raise ValueError(
+                "value must be google.cloud.bigquery.table.TableConstraints " "or 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.
@@ -3282,7 +3294,20 @@ def from_api_repr(cls, api_repr: Dict[str, Any]) -> "ForeignKey":
                 for column_reference_resource in api_repr["columnReferences"]
             ],
         )
-
+
+    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.
@@ -3319,6 +3344,16 @@ 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

@@ -93,6 +93,10 @@
         ],
     ),
 ]
+TABLE_CONSTRAINTS_SCHEMA = [
+    bigquery.SchemaField("id", "INTEGER", mode="REQUIRED"),
+    bigquery.SchemaField("product_id", "STRING", mode="REQUIRED"),
+]
 
 SOURCE_URIS_AVRO = [
     "gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/a-twitter.avro",
@@ -775,6 +779,57 @@ 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):
+        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))
+
+        table.table_constraints = {
+            "primaryKey": {
+                "columns": ["id"]
+            },
+            "foreignKeys": [
+                {
+                    "name": "fk_product_id",
+                    "referencedTable": "products",
+                    "columnReferences": [
+                        {
+                            "referencingColumn": "product_id",
+                            "referencedColumn": "id"
+                        }
+                    ]
+                }
+            ]
+        }
+        table2 = Config.CLIENT.update_table(table, ["table_constraints"])
+        self.assertEqual(table2.table_constraints, {
+            "primaryKey": {
+                "columns": ["id"]
+            },
+            "foreignKeys": [
+                {
+                    "name": "fk_product_id",
+                    "referencedTable": "products",
+                    "columnReferences": [
+                        {
+                            "referencingColumn": "product_id",
+                            "referencedColumn": "id"
+                        }
+                    ]
+                }
+            ]
+        })
+
+        table2.table_constraints = None
+        table3 = Config.CLIENT.update_table(table2, ["table_constraints"])
+        self.assertIsNone(table3.table_constraints, None)
+
     @staticmethod
     def _fetch_single_page(table, selected_fields=None):
         iterator = Config.CLIENT.list_rows(table, selected_fields=selected_fields)

tests/unit/test_table.py

@@ -19,6 +19,7 @@
 from sys import version_info
 import time
 import types
+from typing import Type
 import unittest
 from unittest import mock
 import warnings
@@ -892,6 +893,97 @@ def test_table_constraints_property_getter(self):
         assert isinstance(table_constraints, TableConstraints)
         assert table_constraints.primary_key == PrimaryKey(columns=["id"])
 
+    def test_table_constraints_property_setter(self):
+        from google.cloud.bigquery.table import ColumnReference, ForeignKey, PrimaryKey, TableConstraints, TableReference
+
+        dataset = DatasetReference(self.PROJECT, self.DS_ID)
+        table_ref = dataset.table(self.TABLE_NAME)
+        table = self._make_one(table_ref)
+
+        primary_key = PrimaryKey(columns=["id"])
+        foreign_keys = [ForeignKey(
+            name="fk_name",
+            referenced_table=TableReference.from_string("my_project.my_dataset.table"),
+            column_references=[ColumnReference(referenced_column="product_id", referencing_column="id")]
+        )]
+        table_constraints = TableConstraints(primary_key=primary_key, foreign_keys=foreign_keys)
+        table.table_constraints = table_constraints
+
+        assert table._properties["tableConstraints"] == {
+            "primaryKey": {"columns": ["id"]},
+            "foreignKeys": [ {
+                    "name": "fk_name",
+                    "referencedTable": {
+                        "projectId": "my_project",
+                        "datasetId": "my_dataset",
+                        "tableId": "table"
+                    },
+                    "columnReferences": [
+                        {"referencedColumn": "product_id", "referencingColumn": "id"}
+                    ]
+                }
+            ]
+        }
+
+    def test_table_constraints_property_setter_only_primary_key(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)
+
+        primary_key = PrimaryKey(columns=["id"])
+
+        table_constraints = TableConstraints(primary_key=primary_key, foreign_keys=None)
+        table.table_constraints = table_constraints
+
+        assert table._properties["tableConstraints"] == {
+            "primaryKey": {"columns": ["id"]}
+        }
+
+    def test_table_constraints_property_setter_only_foriegn_keys(self):
+        from google.cloud.bigquery.table import ColumnReference, ForeignKey, TableConstraints, TableReference
+
+        dataset = DatasetReference(self.PROJECT, self.DS_ID)
+        table_ref = dataset.table(self.TABLE_NAME)
+        table = self._make_one(table_ref)
+
+        foreign_keys = [ForeignKey(
+            name="fk_name",
+            referenced_table=TableReference.from_string("my_project.my_dataset.table"),
+            column_references=[ColumnReference(referenced_column="product_id", referencing_column="id")]
+        )]
+        table_constraints = TableConstraints(primary_key=None, foreign_keys=foreign_keys)
+        table.table_constraints = table_constraints
+
+        assert table._properties["tableConstraints"] == {
+            "foreignKeys": [ {
+                    "name": "fk_name",
+                    "referencedTable": {
+                        "projectId": "my_project",
+                        "datasetId": "my_dataset",
+                        "tableId": "table"
+                    },
+                    "columnReferences": [
+                        {"referencedColumn": "product_id", "referencingColumn": "id"}
+                    ]
+                }
+            ]
+        }
+
+    def test_table_constraints_property_setter_no_constraints(self):
+        from google.cloud.bigquery.table import TableConstraints
+
+        dataset = DatasetReference(self.PROJECT, self.DS_ID)
+        table_ref = dataset.table(self.TABLE_NAME)
+        table = self._make_one(table_ref)
+
+
+        table_constraints = TableConstraints(primary_key=None, foreign_keys=None)
+        table.table_constraints = table_constraints
+
+        assert table._properties["tableConstraints"] == {}    
+
     def test_description_setter_bad_value(self):
         dataset = DatasetReference(self.PROJECT, self.DS_ID)
         table_ref = dataset.table(self.TABLE_NAME)
@@ -5783,6 +5875,47 @@ def test__eq__other_type(self):
             foreign_key == "This is not a Foreign Key"
 
 
+    def test_to_api_repr(self):
+        from google.cloud.bigquery.table import ColumnReference, TableReference
+
+        name = "my_fk"
+        referenced_table = TableReference.from_string("my-project.mydataset.mytable")
+        column_references = [
+            ColumnReference(referencing_column="product_id", referenced_column="id")
+        ]
+        foreign_key = self._make_one(name, referenced_table, column_references)
+
+        expected = {
+            "name": name,
+            "referencedTable": {
+                "projectId": "my-project",
+                "datasetId": "mydataset",
+                "tableId": "mytable",
+            },
+            "columnReferences": [
+                {"referencingColumn": "product_id", "referencedColumn": "id"}
+            ],
+        }
+        self.assertEqual(foreign_key.to_api_repr(), expected)
+
+    def test_to_api_repr_empty_column_references(self):
+        from google.cloud.bigquery.table import TableReference
+        name = "my_fk"
+        referenced_table = TableReference.from_string("my-project.mydataset.mytable")
+        column_references = []
+        foreign_key = self._make_one(name, referenced_table, column_references)
+
+        expected = {
+            "name": name,
+            "referencedTable": {
+                "projectId": "my-project",
+                "datasetId": "mydataset",
+                "tableId": "mytable",
+            },
+            "columnReferences": [],
+        }
+        self.assertEqual(foreign_key.to_api_repr(), expected)
+
 class TestTableConstraint(unittest.TestCase):
     @staticmethod
     def _get_target_class():
@@ -5879,6 +6012,87 @@ def test_from_api_repr_only_foreign_keys_resource(self):
         self.assertIsNotNone(instance.foreign_keys)
 
 
+    def test_to_api_repr(self):
+        from google.cloud.bigquery.table import ColumnReference, ForeignKey, PrimaryKey
+
+        primary_key = PrimaryKey(columns=["id", "product_id"])
+        foreign_keys = [
+            ForeignKey(
+                name="my_fk_name",
+                referenced_table=TableReference.from_string("my-project.my-dataset.products"),
+                column_references=[
+                    ColumnReference(referencing_column="product_id", referenced_column="id"),
+                ],
+            )
+        ]
+        instance = self._make_one(primary_key=primary_key, foreign_keys=foreign_keys)
+
+        expected = {
+            "primaryKey": {
+                "columns": ["id", "product_id"],
+            },
+            "foreignKeys": [
+                {
+                    "name": "my_fk_name",
+                    "referencedTable": {
+                        "projectId": "my-project",
+                        "datasetId": "my-dataset",
+                        "tableId": "products",
+                    },
+                    "columnReferences": [
+                        {"referencingColumn": "product_id", "referencedColumn": "id"},
+                    ],
+                }
+            ],
+        }
+        self.assertEqual(instance.to_api_repr(), expected)
+
+    def test_to_api_repr_only_primary_key(self):
+        from google.cloud.bigquery.table import PrimaryKey
+        primary_key = PrimaryKey(columns=["id", "product_id"])
+        instance = self._make_one(primary_key=primary_key, foreign_keys=None)
+        expected = {
+            "primaryKey": {
+                "columns": ["id", "product_id"],
+            },
+        }
+        self.assertEqual(instance.to_api_repr(), expected)
+
+    def test_to_api_repr_only_foreign_keys(self):
+        from google.cloud.bigquery.table import ColumnReference, ForeignKey
+        foreign_keys = [
+            ForeignKey(
+                name="my_fk_name",
+                referenced_table=TableReference.from_string("my-project.my-dataset.products"),
+                column_references=[
+                    ColumnReference(referencing_column="product_id", referenced_column="id"),
+                ],
+            )
+        ]
+        instance = self._make_one(primary_key=None, foreign_keys=foreign_keys)
+        expected = {
+            "foreignKeys": [
+                {
+                    "name": "my_fk_name",
+                    "referencedTable": {
+                        "projectId": "my-project",
+                        "datasetId": "my-dataset",
+                        "tableId": "products",
+                    },
+                    "columnReferences": [
+                        {"referencingColumn": "product_id", "referencedColumn": "id"},
+                    ],
+                }
+            ],
+        }
+        self.assertEqual(instance.to_api_repr(), expected)
+
+    def test_to_api_repr_no_constraints(self):
+        instance = self._make_one(primary_key=None, foreign_keys=None)
+        expected = {}
+        self.assertEqual(instance.to_api_repr(), expected)
+
+
 @pytest.mark.parametrize(
     "table_path",
     (