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

* 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

Author: lkhagvadorj-amp

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.

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)