Skip to content

Commit a167f9a

Browse files
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 <[email protected]> Co-authored-by: Chalmer Lowe <[email protected]>
1 parent 08483fb commit a167f9a

File tree

2 files changed

+408
-0
lines changed

2 files changed

+408
-0
lines changed

google/cloud/bigquery/table.py

+128
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,7 @@ class Table(_TableBase):
390390
"view_use_legacy_sql": "view",
391391
"view_query": "view",
392392
"require_partition_filter": "requirePartitionFilter",
393+
"table_constraints": "tableConstraints",
393394
}
394395

395396
def __init__(self, table_ref, schema=None) -> None:
@@ -973,6 +974,16 @@ def clone_definition(self) -> Optional["CloneDefinition"]:
973974
clone_info = CloneDefinition(clone_info)
974975
return clone_info
975976

977+
@property
978+
def table_constraints(self) -> Optional["TableConstraints"]:
979+
"""Tables Primary Key and Foreign Key information."""
980+
table_constraints = self._properties.get(
981+
self._PROPERTY_TO_API_FIELD["table_constraints"]
982+
)
983+
if table_constraints is not None:
984+
table_constraints = TableConstraints.from_api_repr(table_constraints)
985+
return table_constraints
986+
976987
@classmethod
977988
def from_string(cls, full_table_id: str) -> "Table":
978989
"""Construct a table from fully-qualified table ID.
@@ -2958,6 +2969,123 @@ def __repr__(self):
29582969
return "TimePartitioning({})".format(",".join(key_vals))
29592970

29602971

2972+
class PrimaryKey:
2973+
"""Represents the primary key constraint on a table's columns.
2974+
2975+
Args:
2976+
columns: The columns that are composed of the primary key constraint.
2977+
"""
2978+
2979+
def __init__(self, columns: List[str]):
2980+
self.columns = columns
2981+
2982+
def __eq__(self, other):
2983+
if not isinstance(other, PrimaryKey):
2984+
raise TypeError("The value provided is not a BigQuery PrimaryKey.")
2985+
return self.columns == other.columns
2986+
2987+
2988+
class ColumnReference:
2989+
"""The pair of the foreign key column and primary key column.
2990+
2991+
Args:
2992+
referencing_column: The column that composes the foreign key.
2993+
referenced_column: The column in the primary key that are referenced by the referencingColumn.
2994+
"""
2995+
2996+
def __init__(self, referencing_column: str, referenced_column: str):
2997+
self.referencing_column = referencing_column
2998+
self.referenced_column = referenced_column
2999+
3000+
def __eq__(self, other):
3001+
if not isinstance(other, ColumnReference):
3002+
raise TypeError("The value provided is not a BigQuery ColumnReference.")
3003+
return (
3004+
self.referencing_column == other.referencing_column
3005+
and self.referenced_column == other.referenced_column
3006+
)
3007+
3008+
3009+
class ForeignKey:
3010+
"""Represents a foreign key constraint on a table's columns.
3011+
3012+
Args:
3013+
name: Set only if the foreign key constraint is named.
3014+
referenced_table: The table that holds the primary key and is referenced by this foreign key.
3015+
column_references: The columns that compose the foreign key.
3016+
"""
3017+
3018+
def __init__(
3019+
self,
3020+
name: str,
3021+
referenced_table: TableReference,
3022+
column_references: List[ColumnReference],
3023+
):
3024+
self.name = name
3025+
self.referenced_table = referenced_table
3026+
self.column_references = column_references
3027+
3028+
def __eq__(self, other):
3029+
if not isinstance(other, ForeignKey):
3030+
raise TypeError("The value provided is not a BigQuery ForeignKey.")
3031+
return (
3032+
self.name == other.name
3033+
and self.referenced_table == other.referenced_table
3034+
and self.column_references == other.column_references
3035+
)
3036+
3037+
@classmethod
3038+
def from_api_repr(cls, api_repr: Dict[str, Any]) -> "ForeignKey":
3039+
"""Create an instance from API representation."""
3040+
return cls(
3041+
name=api_repr["name"],
3042+
referenced_table=TableReference.from_api_repr(api_repr["referencedTable"]),
3043+
column_references=[
3044+
ColumnReference(
3045+
column_reference_resource["referencingColumn"],
3046+
column_reference_resource["referencedColumn"],
3047+
)
3048+
for column_reference_resource in api_repr["columnReferences"]
3049+
],
3050+
)
3051+
3052+
3053+
class TableConstraints:
3054+
"""The TableConstraints defines the primary key and foreign key.
3055+
3056+
Args:
3057+
primary_key:
3058+
Represents a primary key constraint on a table's columns. Present only if the table
3059+
has a primary key. The primary key is not enforced.
3060+
foreign_keys:
3061+
Present only if the table has a foreign key. The foreign key is not enforced.
3062+
3063+
"""
3064+
3065+
def __init__(
3066+
self,
3067+
primary_key: Optional[PrimaryKey],
3068+
foreign_keys: Optional[List[ForeignKey]],
3069+
):
3070+
self.primary_key = primary_key
3071+
self.foreign_keys = foreign_keys
3072+
3073+
@classmethod
3074+
def from_api_repr(cls, resource: Dict[str, Any]) -> "TableConstraints":
3075+
"""Create an instance from API representation."""
3076+
primary_key = None
3077+
if "primaryKey" in resource:
3078+
primary_key = PrimaryKey(resource["primaryKey"]["columns"])
3079+
3080+
foreign_keys = None
3081+
if "foreignKeys" in resource:
3082+
foreign_keys = [
3083+
ForeignKey.from_api_repr(foreign_key_resource)
3084+
for foreign_key_resource in resource["foreignKeys"]
3085+
]
3086+
return cls(primary_key, foreign_keys)
3087+
3088+
29613089
def _item_to_row(iterator, resource):
29623090
"""Convert a JSON row to the native object.
29633091

0 commit comments

Comments
 (0)