[superset] 01/01: chore: remove get_columns_description duplication

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch sc-72620
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 5207978f7668e74fb2ef42dfdaf34b67701b6d0b
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Wed Jul 26 13:45:00 2023 -0700

    chore: remove get_columns_description duplication
---
 superset/connectors/sqla/models.py |  2 +-
 superset/connectors/sqla/utils.py  | 19 +++++--------------
 2 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 979a53e086..e7a3f50ae1 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -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
diff --git a/superset/connectors/sqla/utils.py b/superset/connectors/sqla/utils.py
index 82d8f90f22..4bff059b37 100644
--- a/superset/connectors/sqla/utils.py
+++ b/superset/connectors/sqla/utils.py
@@ -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)