Skip to content

Commit 5e7110a

Browse files
BenConstable9BenConstable9
committed
Update databricks connector
1 parent 75d962b commit 5e7110a

File tree

3 files changed

+66
-35
lines changed

3 files changed

+66
-35
lines changed

‎text_2_sql/data_dictionary/databricks_data_dictionary_creator.py‎

Lines changed: 40 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import os
88

99

10-
class SnowflakeDataDictionaryCreator(DataDictionaryCreator):
10+
class DatabricksDataDictionaryCreator(DataDictionaryCreator):
1111
def __init__(
1212
self,
1313
entities: list[str] = None,
@@ -30,11 +30,11 @@ def __init__(
3030
self.catalog = os.environ["Text2Sql__Databricks__Catalog"]
3131
self.database_engine = DatabaseEngine.DATABRICKS
3232

33-
"""A class to extract data dictionary information from a Snowflake database."""
33+
"""A class to extract data dictionary information from Databricks Unity Catalog."""
3434

3535
@property
3636
def extract_table_entities_sql_query(self) -> str:
37-
"""A property to extract table entities from a Snowflake database."""
37+
"""A property to extract table entities from Databricks Unity Catalog."""
3838
return f"""SELECT
3939
t.TABLE_NAME AS Entity,
4040
t.TABLE_SCHEMA AS EntitySchema,
@@ -47,7 +47,7 @@ def extract_table_entities_sql_query(self) -> str:
4747

4848
@property
4949
def extract_view_entities_sql_query(self) -> str:
50-
"""A property to extract view entities from a Snowflake database."""
50+
"""A property to extract view entities from Databricks Unity Catalog."""
5151
return """SELECT
5252
v.TABLE_NAME AS Entity,
5353
v.TABLE_SCHEMA AS EntitySchema
@@ -58,7 +58,7 @@ def extract_view_entities_sql_query(self) -> str:
5858
v.TABLE_CATALOG = '{self.catalog}'"""
5959

6060
def extract_columns_sql_query(self, entity: EntityItem) -> str:
61-
"""A property to extract column information from a Snowflake database."""
61+
"""A property to extract column information from Databricks Unity Catalog."""
6262
return f"""SELECT
6363
COLUMN_NAME AS Name,
6464
DATA_TYPE AS Type,
@@ -72,23 +72,43 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
7272

7373
@property
7474
def extract_entity_relationships_sql_query(self) -> str:
75-
"""A property to extract entity relationships from a SQL Server database."""
76-
return """SELECT
77-
tc.table_schema AS EntitySchema,
78-
tc.table_name AS Entity,
79-
rc.unique_constraint_schema AS ForeignEntitySchema,
80-
rc.unique_constraint_name AS ForeignEntityConstraint,
81-
rc.constraint_name AS ForeignKeyConstraint
75+
"""A property to extract entity relationships from Databricks Unity Catalog."""
76+
return f"""SELECT
77+
fk_schema.TABLE_SCHEMA AS EntitySchema,
78+
fk_tab.TABLE_NAME AS Entity,
79+
pk_schema.TABLE_SCHEMA AS ForeignEntitySchema,
80+
pk_tab.TABLE_NAME AS ForeignEntity,
81+
fk_col.COLUMN_NAME AS [Column],
82+
pk_col.COLUMN_NAME AS ForeignColumn
8283
FROM
83-
information_schema.referential_constraints rc
84-
JOIN
85-
information_schema.table_constraints tc
86-
ON rc.constraint_schema = tc.constraint_schema
87-
AND rc.constraint_name = tc.constraint_name
84+
INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS fk
85+
INNER JOIN
86+
INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS fkc
87+
ON fk.constraint_name = fkc.constraint_name
88+
INNER JOIN
89+
INFORMATION_SCHEMA.TABLES AS fk_tab
90+
ON fk_tab.TABLE_NAME = fkc.TABLE_NAME AND fk_tab.TABLE_SCHEMA = fkc.TABLE_SCHEMA
91+
INNER JOIN
92+
INFORMATION_SCHEMA.SCHEMATA AS fk_schema
93+
ON fk_tab.TABLE_SCHEMA = fk_schema.TABLE_SCHEMA
94+
INNER JOIN
95+
INFORMATION_SCHEMA.TABLES AS pk_tab
96+
ON pk_tab.TABLE_NAME = fkc.referenced_TABLE_NAME AND pk_tab.TABLE_SCHEMA = fkc.referenced_TABLE_SCHEMA
97+
INNER JOIN
98+
INFORMATION_SCHEMA.SCHEMATA AS pk_schema
99+
ON pk_tab.TABLE_SCHEMA = pk_schema.TABLE_SCHEMA
100+
INNER JOIN
101+
INFORMATION_SCHEMA.COLUMNS AS fk_col
102+
ON fkc.COLUMN_NAME = fk_col.COLUMN_NAME AND fkc.TABLE_NAME = fk_col.TABLE_NAME AND fkc.TABLE_SCHEMA = fk_col.TABLE_SCHEMA
103+
INNER JOIN
104+
INFORMATION_SCHEMA.COLUMNS AS pk_col
105+
ON fkc.referenced_COLUMN_NAME = pk_col.COLUMN_NAME AND fkc.referenced_TABLE_NAME = pk_col.TABLE_NAME AND fkc.referenced_TABLE_SCHEMA = pk_col.TABLE_SCHEMA
88106
WHERE
89-
tc.constraint_type = 'FOREIGN KEY'
107+
fk.constraint_type = 'FOREIGN KEY'
108+
AND fk_tab.TABLE_CATALOG = '{self.catalog}'
109+
AND pk_tab.TABLE_CATALOG = '{self.catalog}'
90110
ORDER BY
91-
EntitySchema, Entity, ForeignEntitySchema, ForeignEntityConstraint;
111+
EntitySchema, Entity, ForeignEntitySchema, ForeignEntity;
92112
"""
93113

94114
async def query_entities(self, sql_query: str, cast_to: any = None) -> list[dict]:
@@ -143,5 +163,5 @@ async def query_entities(self, sql_query: str, cast_to: any = None) -> list[dict
143163

144164

145165
if __name__ == "__main__":
146-
data_dictionary_creator = SnowflakeDataDictionaryCreator()
166+
data_dictionary_creator = DatabricksDataDictionaryCreator()
147167
asyncio.run(data_dictionary_creator.create_data_dictionary())

‎text_2_sql/data_dictionary/snowflake_data_dictionary_creator.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
6767

6868
@property
6969
def extract_entity_relationships_sql_query(self) -> str:
70-
"""A property to extract entity relationships from a SQL Server database."""
70+
"""A property to extract entity relationships from a Snowflake database."""
7171
return """SELECT
7272
tc.table_schema AS EntitySchema,
7373
tc.table_name AS Entity,

‎text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py‎

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -86,28 +86,39 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
8686
def extract_entity_relationships_sql_query(self) -> str:
8787
"""A property to extract entity relationships from a SQL Server database."""
8888
return """SELECT
89-
fk_schema.name AS EntitySchema,
90-
fk_tab.name AS Entity,
91-
pk_schema.name AS ForeignEntitySchema,
92-
pk_tab.name AS ForeignEntity,
93-
fk_col.name AS [Column],
94-
pk_col.name AS ForeignColumn
89+
fk_schema.schema_name AS EntitySchema,
90+
fk_tab.table_name AS Entity,
91+
pk_schema.schema_name AS ForeignEntitySchema,
92+
pk_tab.table_name AS ForeignEntity,
93+
fk_col.column_name AS [Column],
94+
pk_col.column_name AS ForeignColumn
9595
FROM
96-
sys.foreign_keys AS fk
96+
INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS fk
9797
INNER JOIN
98-
sys.foreign_key_columns AS fkc ON fk.object_id = fkc.constraint_object_id
98+
INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS fkc
99+
ON fk.constraint_name = fkc.constraint_name
99100
INNER JOIN
100-
sys.tables AS fk_tab ON fk_tab.object_id = fk.parent_object_id
101+
INFORMATION_SCHEMA.TABLES AS fk_tab
102+
ON fk_tab.table_name = fkc.table_name AND fk_tab.table_schema = fkc.table_schema
101103
INNER JOIN
102-
sys.schemas AS fk_schema ON fk_tab.schema_id = fk_schema.schema_id
104+
INFORMATION_SCHEMA.SCHEMATA AS fk_schema
105+
ON fk_tab.table_schema = fk_schema.schema_name
103106
INNER JOIN
104-
sys.tables AS pk_tab ON pk_tab.object_id = fk.referenced_object_id
107+
INFORMATION_SCHEMA.TABLES AS pk_tab
108+
ON pk_tab.table_name = fkc.referenced_table_name AND pk_tab.table_schema = fkc.referenced_table_schema
105109
INNER JOIN
106-
sys.schemas AS pk_schema ON pk_tab.schema_id = pk_schema.schema_id
110+
INFORMATION_SCHEMA.SCHEMATA AS pk_schema
111+
ON pk_tab.table_schema = pk_schema.schema_name
107112
INNER JOIN
108-
sys.columns AS fk_col ON fkc.parent_object_id = fk_col.object_id AND fkc.parent_column_id = fk_col.column_id
113+
INFORMATION_SCHEMA.COLUMNS AS fk_col
114+
ON fkc.column_name = fk_col.column_name AND fkc.table_name = fk_col.table_name AND fkc.table_schema = fk_col.table_schema
109115
INNER JOIN
110-
sys.columns AS pk_col ON fkc.referenced_object_id = pk_col.object_id AND fkc.referenced_column_id = pk_col.column_id
116+
INFORMATION_SCHEMA.COLUMNS AS pk_col
117+
ON fkc.referenced_column_name = pk_col.column_name AND fkc.referenced_table_name = pk_col.table_name AND fkc.referenced_table_schema = pk_col.table_schema
118+
WHERE
119+
fk.constraint_type = 'FOREIGN KEY'
120+
AND fk_tab.table_catalog = 'your_catalog_name'
121+
AND pk_tab.table_catalog = 'your_catalog_name'
111122
ORDER BY
112123
EntitySchema, Entity, ForeignEntitySchema, ForeignEntity;
113124
"""

0 commit comments

Comments
 (0)