feat: Add table_constraints field to Table model (#1755)

* feat: add `table_constraints` field to Table model * Change `raise` to `return` in __eq__ methods * Fix __eq__ for ColumnReference * Add column_references to ForeignKey __eq__ * Add missing coverage * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update google/cloud/bigquery/table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py * Update tests/unit/test_table.py --------- Co-authored-by: Chalmer Lowe <chalmerlowe@google.com> Co-authored-by: Chalmer Lowe <chalmer.lowe@gmail.com>

Author: dimkonko

google/cloud/bigquery/table.py

@@ -390,6 +390,7 @@ class Table(_TableBase):
  "view_use_legacy_sql": "view",
  "view_query": "view",
  "require_partition_filter": "requirePartitionFilter",
+ "table_constraints": "tableConstraints",
  }
 
  def __init__(self, table_ref, schema=None) -> None:
@@ -973,6 +974,16 @@ def clone_definition(self) -> Optional["CloneDefinition"]:
  clone_info = CloneDefinition(clone_info)
  return clone_info
 
+ @property
+ def table_constraints(self) -> Optional["TableConstraints"]:
+ """Tables Primary Key and Foreign Key information."""
+ table_constraints = self._properties.get(
+ self._PROPERTY_TO_API_FIELD["table_constraints"]
+ )
+ if table_constraints is not None:
+ table_constraints = TableConstraints.from_api_repr(table_constraints)
+ return table_constraints
+
  @classmethod
  def from_string(cls, full_table_id: str) -> "Table":
  """Construct a table from fully-qualified table ID.
@@ -2958,6 +2969,123 @@ def __repr__(self):
  return "TimePartitioning({})".format(",".join(key_vals))
 
 
+class PrimaryKey:
+ """Represents the primary key constraint on a table's columns.
+
+ Args:
+ columns: The columns that are composed of the primary key constraint.
+ """
+
+ def __init__(self, columns: List[str]):
+ self.columns = columns
+
+ def __eq__(self, other):
+ if not isinstance(other, PrimaryKey):
+ raise TypeError("The value provided is not a BigQuery PrimaryKey.")
+ return self.columns == other.columns
+
+
+class ColumnReference:
+ """The pair of the foreign key column and primary key column.
+
+ Args:
+ referencing_column: The column that composes the foreign key.
+ referenced_column: The column in the primary key that are referenced by the referencingColumn.
+ """
+
+ def __init__(self, referencing_column: str, referenced_column: str):
+ self.referencing_column = referencing_column
+ self.referenced_column = referenced_column
+
+ def __eq__(self, other):
+ if not isinstance(other, ColumnReference):
+ raise TypeError("The value provided is not a BigQuery ColumnReference.")
+ return (
+ self.referencing_column == other.referencing_column
+ and self.referenced_column == other.referenced_column
+ )
+
+
+class ForeignKey:
+ """Represents a foreign key constraint on a table's columns.
+
+ Args:
+ name: Set only if the foreign key constraint is named.
+ referenced_table: The table that holds the primary key and is referenced by this foreign key.
+ column_references: The columns that compose the foreign key.
+ """
+
+ def __init__(
+ self,
+ name: str,
+ referenced_table: TableReference,
+ column_references: List[ColumnReference],
+ ):
+ self.name = name
+ self.referenced_table = referenced_table
+ self.column_references = column_references
+
+ def __eq__(self, other):
+ if not isinstance(other, ForeignKey):
+ raise TypeError("The value provided is not a BigQuery ForeignKey.")
+ return (
+ self.name == other.name
+ and self.referenced_table == other.referenced_table
+ and self.column_references == other.column_references
+ )
+
+ @classmethod
+ def from_api_repr(cls, api_repr: Dict[str, Any]) -> "ForeignKey":
+ """Create an instance from API representation."""
+ return cls(
+ name=api_repr["name"],
+ referenced_table=TableReference.from_api_repr(api_repr["referencedTable"]),
+ column_references=[
+ ColumnReference(
+ column_reference_resource["referencingColumn"],
+ column_reference_resource["referencedColumn"],
+ )
+ for column_reference_resource in api_repr["columnReferences"]
+ ],
+ )
+
+
+class TableConstraints:
+ """The TableConstraints defines the primary key and foreign key.
+
+ Args:
+ primary_key:
+ Represents a primary key constraint on a table's columns. Present only if the table
+ has a primary key. The primary key is not enforced.
+ foreign_keys:
+ Present only if the table has a foreign key. The foreign key is not enforced.
+
+ """
+
+ def __init__(
+ self,
+ primary_key: Optional[PrimaryKey],
+ foreign_keys: Optional[List[ForeignKey]],
+ ):
+ self.primary_key = primary_key
+ self.foreign_keys = foreign_keys
+
+ @classmethod
+ def from_api_repr(cls, resource: Dict[str, Any]) -> "TableConstraints":
+ """Create an instance from API representation."""
+ primary_key = None
+ if "primaryKey" in resource:
+ primary_key = PrimaryKey(resource["primaryKey"]["columns"])
+
+ foreign_keys = None
+ if "foreignKeys" in resource:
+ foreign_keys = [
+ ForeignKey.from_api_repr(foreign_key_resource)
+ for foreign_key_resource in resource["foreignKeys"]
+ ]
+ return cls(primary_key, foreign_keys)
+
+
 def _item_to_row(iterator, resource):
  """Convert a JSON row to the native object.