Skip to content

Commit

Permalink
feat: Add table_constraints field to Table model (#1755)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
3 people committed Jan 11, 2024
1 parent 08483fb commit a167f9a
Show file tree
Hide file tree
Showing 2 changed files with 408 additions and 0 deletions.
128 changes: 128 additions & 0 deletions google/cloud/bigquery/table.py
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit a167f9a

Please sign in to comment.