iop
/
superset
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

chore: remove get_columns_description duplication (#24819)

This commit is contained in:
Beto Dealmeida 2023-07-27 14:39:15 -07:00 committed by GitHub
parent a1396d00aa
commit 7cd317fe2e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
superset/connectors/sqla/models.py
View File

@ -1003,7 +1003,7 @@ class SqlaTable(
tbl, _ = self.get_from_clause(template_processor)
qry = sa.select([sqla_column]).limit(1).select_from(tbl)
sql = self.database.compile_sqla_query(qry)
col_desc = get_columns_description(self.database, sql)
col_desc = get_columns_description(self.database, self.schema, sql)
is_dttm = col_desc[0]["is_dttm"] # type: ignore
except SupersetGenericDBErrorException as ex:
raise ColumnNotFoundException(message=str(ex)) from ex

19
superset/connectors/sqla/utils.py
View File

@ -125,28 +125,19 @@ def get_virtual_table_metadata(dataset: SqlaTable) -> list[ResultSetColumnType]:
level=ErrorLevel.ERROR,
)
)
# TODO(villebro): refactor to use same code that's used by
# sql_lab.py:execute_sql_statements
try:
with dataset.database.get_raw_connection(schema=dataset.schema) as conn:
cursor = conn.cursor()
query = dataset.database.apply_limit_to_sql(statements[0], limit=1)
db_engine_spec.execute(cursor, query)
result = db_engine_spec.fetch_data(cursor, limit=1)
result_set = SupersetResultSet(result, cursor.description, db_engine_spec)
cols = result_set.columns
except Exception as ex:
raise SupersetGenericDBErrorException(message=str(ex)) from ex
return cols
return get_columns_description(dataset.database, dataset.schema, statements[0])
def get_columns_description(
database: Database,
schema: str | None,
query: str,
) -> list[ResultSetColumnType]:
# TODO(villebro): refactor to use same code that's used by
# sql_lab.py:execute_sql_statements
db_engine_spec = database.db_engine_spec
try:
with database.get_raw_connection() as conn:
with database.get_raw_connection(schema=schema) as conn:
cursor = conn.cursor()
query = database.apply_limit_to_sql(query, limit=1)
cursor.execute(query)