diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index 38542023b..f1facf8b8 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -1022,6 +1022,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. @@ -3247,7 +3259,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. @@ -3284,6 +3309,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. diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index ff0593470..b762ad98b 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -870,6 +870,20 @@ 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 PrimaryKey, 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=PrimaryKey(columns=["id"]), foreign_keys=[]) + table.table_constraints = table_constraints + + assert table._properties["tableConstraints"] == { + "primaryKey": {"columns": ["id"]}, + } + def test_description_setter_bad_value(self): dataset = DatasetReference(self.PROJECT, self.DS_ID) table_ref = dataset.table(self.TABLE_NAME)