Skip to content

Commit a3b00fa

Browse files
author
James Robinson
authored
Fix \introspect for Snowflake with foreign key foreign table schemas … (#87)
* Fix \introspect for Snowflake with foreign key foreign table schemas being described as None. Err on side of the the referencing table's schema then.
1 parent 003065b commit a3b00fa

File tree

1 file changed

+15
-43
lines changed

1 file changed

+15
-43
lines changed

noteable_magics/sql/meta_commands.py

Lines changed: 15 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -526,7 +526,7 @@ def run(self, invoked_as: str, args: List[str]) -> None:
526526
"""
527527

528528
try:
529-
ds_id = self.get_datasource_id()
529+
ds_id: UUID = self.get_datasource_id()
530530
except ValueError:
531531
# Tried to introspect a datasource whose kernel-side handle in the connections dict
532532
# wasn't coercable back into a UUID. This could be the case for the pre-datasource
@@ -563,7 +563,6 @@ def delta() -> float:
563563
with ThreadPoolExecutor(
564564
max_workers=self.MAX_INTROSPECTION_THREADS
565565
) as executor, RelationStructureMessager(ds_id) as messenger:
566-
567566
future_to_relation = {
568567
executor.submit(
569568
self.fully_introspect, inspector, schema_name, relation_name, kind
@@ -622,9 +621,8 @@ def all_table_and_views(self, inspector) -> List[Tuple[str, str, str]]:
622621
return results
623622

624623
def fully_introspect(
625-
self, inspector: 'SchemaStrippingInspector', schema_name: str, relation_name: str, kind: str
624+
self, inspector: SchemaStrippingInspector, schema_name: str, relation_name: str, kind: str
626625
) -> RelationStructureDescription:
627-
628626
"""Drive introspecting into this single relation, making all the necessary Introspector API
629627
calls to learn all of the relation's sub-structures.
630628
@@ -676,19 +674,28 @@ def fully_introspect(
676674
)
677675

678676
def introspect_foreign_keys(
679-
self, inspector, schema_name, relation_name
677+
self, inspector: SchemaStrippingInspector, schema_name: str, relation_name: str
680678
) -> List[ForeignKeysModel]:
681679
"""Introspect all foreign keys for a table, describing the results as a List[ForeignKeysModel]"""
682680

683681
fkeys: List[ForeignKeysModel] = []
684682

685683
fkey_dicts = inspector.get_foreign_keys(relation_name, schema_name)
686684

685+
# Convert from SQLA foreign key dicts to our ForeignKeysModel. But beware,
686+
# Snowflake driver reports FKs with None for the target table's schema
687+
# (at least sometimes?) so in case then err on the referencing table's schema, because
688+
# ForeignKeysModel shared between us and Gate hate None for referred_schema.
689+
687690
for fkey in sorted(fkey_dicts, key=lambda d: d['name']):
688691
fkeys.append(
689692
ForeignKeysModel(
690693
name=fkey['name'],
691-
referenced_schema=fkey['referred_schema'],
694+
referenced_schema=(
695+
fkey['referred_schema']
696+
if fkey['referred_schema'] is not None
697+
else schema_name
698+
),
692699
referenced_relation=fkey['referred_table'],
693700
columns=fkey['constrained_columns'],
694701
referenced_columns=fkey['referred_columns'],
@@ -740,7 +747,6 @@ def introspect_indexes(self, inspector, schema_name, relation_name) -> List[Inde
740747
index_dicts = inspector.get_indexes(relation_name, schema_name)
741748

742749
for index_dict in sorted(index_dicts, key=lambda d: d['name']):
743-
744750
indexes.append(
745751
IndexModel(
746752
name=index_dict['name'],
@@ -752,7 +758,7 @@ def introspect_indexes(self, inspector, schema_name, relation_name) -> List[Inde
752758
return indexes
753759

754760
def introspect_primary_key(
755-
self, inspector, relation_name, schema_name
761+
self, inspector: SchemaStrippingInspector, relation_name: str, schema_name: str
756762
) -> Tuple[Optional[str], List[str]]:
757763
"""Introspect the primary key of a table, returning the pkey name and list of columns in the primary key (if any).
758764
@@ -770,9 +776,8 @@ def introspect_primary_key(
770776
return None, []
771777

772778
def introspect_columns(
773-
self, inspector: 'SchemaStrippingInspector', schema_name: str, relation_name: str
779+
self, inspector: SchemaStrippingInspector, schema_name: str, relation_name: str
774780
) -> List[ColumnModel]:
775-
776781
column_dicts = inspector.get_columns(relation_name, schema=schema_name)
777782

778783
retlist = []
@@ -798,38 +803,6 @@ def introspect_columns(
798803

799804
return retlist
800805

801-
def inform_gate_relation(
802-
self,
803-
session: requests.Session,
804-
datasource_id,
805-
relation_description: RelationStructureDescription,
806-
):
807-
"""POST this `relation_description` up to Gate for storage, relating to `datasource_id`"""
808-
809-
resp = session.post(
810-
f"http://gate.default/api/v1/datasources/{datasource_id}/schema/relation",
811-
json=relation_description.dict(),
812-
)
813-
814-
if resp.status_code == 204:
815-
print(
816-
f'Stored structure of {relation_description.schema_name}.{relation_description.relation_name}'
817-
)
818-
else:
819-
print(
820-
f'Failed storing structure of {relation_description.schema_name}.{relation_description.relation_name}: {resp.status_code}, {resp.text}'
821-
)
822-
823-
def inform_gate_completed(
824-
self, session: requests.Session, datasource_id: UUID, started_at: datetime
825-
):
826-
"""Tell gate to forget about any structures known for this datasource older than when we
827-
started this introspection run."""
828-
829-
session.delete(
830-
f"http://gate.default/api/v1/datasources/{datasource_id}/schema/relations?older_than={started_at.isoformat()}"
831-
)
832-
833806
def get_datasource_id(self) -> UUID:
834807
"""Convert a noteable_magics.sql.connection.Connection's name to the original
835808
UUID Gate knew it as.
@@ -894,7 +867,6 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
894867
self._message_gate(completed_introspection=True)
895868

896869
def _message_gate(self, completed_introspection: bool = False):
897-
898870
base_url = f"http://gate.default/api/v1/datasources/{self._datasource_id}/schema/relations"
899871

900872
if self._relations:

0 commit comments

Comments
 (0)