Skip to content

Commit 6d2094d

Browse files
committed
Support multi-schemas for refresh endpoint
1 parent f6d617e commit 6d2094d

File tree

3 files changed

+43
-27
lines changed

3 files changed

+43
-27
lines changed

Diff for: dataherald/api/fastapi.py

+18-6
Original file line numberDiff line numberDiff line change
@@ -195,18 +195,30 @@ def refresh_table_description(
195195
db_connection = db_connection_repository.find_by_id(
196196
refresh_table_description.db_connection_id
197197
)
198-
198+
scanner = self.system.instance(Scanner)
199+
database_connection_service = DatabaseConnectionService(scanner, self.storage)
199200
try:
200-
sql_database = SQLDatabase.get_sql_engine(db_connection, True)
201-
tables = sql_database.get_tables_and_views()
201+
data = {}
202+
if db_connection.schemas:
203+
for schema in db_connection.schemas:
204+
sql_database = database_connection_service.get_sql_database(
205+
db_connection, schema
206+
)
207+
if schema not in data.keys():
208+
data[schema] = []
209+
data[schema] = sql_database.get_tables_and_views()
210+
else:
211+
sql_database = database_connection_service.get_sql_database(
212+
db_connection
213+
)
214+
data[None] = sql_database.get_tables_and_views()
202215

203-
# Get tables and views and create missing table-descriptions as NOT_SCANNED and update DEPRECATED
204216
scanner_repository = TableDescriptionRepository(self.storage)
205-
scanner = self.system.instance(Scanner)
217+
206218
return [
207219
TableDescriptionResponse(**record.dict())
208220
for record in scanner.refresh_tables(
209-
tables, str(db_connection.id), scanner_repository
221+
data, str(db_connection.id), scanner_repository
210222
)
211223
]
212224
except Exception as e:

Diff for: dataherald/db_scanner/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def create_tables(
4242
@abstractmethod
4343
def refresh_tables(
4444
self,
45-
tables: list[str],
45+
schemas_and_tables: dict[str, list],
4646
db_connection_id: str,
4747
repository: TableDescriptionRepository,
4848
metadata: dict = None,

Diff for: dataherald/db_scanner/sqlalchemy.py

+24-20
Original file line numberDiff line numberDiff line change
@@ -62,34 +62,38 @@ def create_tables(
6262
@override
6363
def refresh_tables(
6464
self,
65-
tables: list[str],
65+
schemas_and_tables: dict[str, list],
6666
db_connection_id: str,
6767
repository: TableDescriptionRepository,
6868
metadata: dict = None,
6969
) -> list[TableDescription]:
70-
stored_tables = repository.find_by({"db_connection_id": str(db_connection_id)})
71-
stored_tables_list = [table.table_name for table in stored_tables]
72-
7370
rows = []
74-
for table_description in stored_tables:
75-
if table_description.table_name not in tables:
76-
table_description.status = TableDescriptionStatus.DEPRECATED.value
77-
rows.append(repository.save_table_info(table_description))
78-
else:
79-
rows.append(TableDescription(**table_description.dict()))
71+
for schema, tables in schemas_and_tables.items():
72+
stored_tables = repository.find_by(
73+
{"db_connection_id": str(db_connection_id), "schema": schema}
74+
)
75+
stored_tables_list = [table.table_name for table in stored_tables]
8076

81-
for table in tables:
82-
if table not in stored_tables_list:
83-
rows.append(
84-
repository.save_table_info(
85-
TableDescription(
86-
db_connection_id=db_connection_id,
87-
table_name=table,
88-
status=TableDescriptionStatus.NOT_SCANNED.value,
89-
metadata=metadata,
77+
for table_description in stored_tables:
78+
if table_description.table_name not in tables:
79+
table_description.status = TableDescriptionStatus.DEPRECATED.value
80+
rows.append(repository.save_table_info(table_description))
81+
else:
82+
rows.append(TableDescription(**table_description.dict()))
83+
84+
for table in tables:
85+
if table not in stored_tables_list:
86+
rows.append(
87+
repository.save_table_info(
88+
TableDescription(
89+
db_connection_id=db_connection_id,
90+
table_name=table,
91+
status=TableDescriptionStatus.NOT_SCANNED.value,
92+
metadata=metadata,
93+
schema_name=schema,
94+
)
9095
)
9196
)
92-
)
9397
return rows
9498

9599
@override

0 commit comments

Comments
 (0)