fix: create virtual table with exotic type (#19714)

* fix: create virtual table with exotic type * fix test
apache · Apr 14, 2022 · 2b17ce2 · 2b17ce2
1 parent 32a9265
commit 2b17ce2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 26 deletions.
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
@@ -2374,11 +2374,10 @@ def write_shadow_dataset(  # pylint: disable=too-many-locals
             referenced_tables = parsed.tables
             tables = load_or_create_tables(
                 session,
-                dataset.database_id,
+                database,
                 dataset.schema,
                 referenced_tables,
                 conditional_quote,
-                engine,
             )
 
         # create the new dataset

diff --git a/superset/connectors/sqla/utils.py b/superset/connectors/sqla/utils.py
@@ -15,13 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 from contextlib import closing
-from datetime import date, datetime, time, timedelta
-from typing import Callable, Dict, List, Optional, Set, TYPE_CHECKING
+from typing import Any, Callable, Dict, List, Optional, Set, TYPE_CHECKING
 
 import sqlparse
 from flask_babel import lazy_gettext as _
-from sqlalchemy import and_, inspect, or_
-from sqlalchemy.engine import Engine
+from sqlalchemy import and_, or_
 from sqlalchemy.exc import NoSuchTableError
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.type_api import TypeEngine
@@ -42,14 +40,11 @@
     from superset.connectors.sqla.models import SqlaTable
 
 
-TEMPORAL_TYPES = {date, datetime, time, timedelta}
-
-
 def get_physical_table_metadata(
     database: Database,
     table_name: str,
     schema_name: Optional[str] = None,
-) -> List[Dict[str, str]]:
+) -> List[Dict[str, Any]]:
     """Use SQLAlchemy inspector to get table metadata"""
     db_engine_spec = database.db_engine_spec
     db_dialect = database.get_dialect()
@@ -86,7 +81,7 @@ def get_physical_table_metadata(
             col.update(
                 {
                     "type": "UNKNOWN",
-                    "generic_type": None,
+                    "type_generic": None,
                     "is_dttm": None,
                 }
             )
@@ -173,20 +168,12 @@ def validate_adhoc_subquery(
     return ";\n".join(str(statement) for statement in statements)
 
 
-def is_column_type_temporal(column_type: TypeEngine) -> bool:
-    try:
-        return column_type.python_type in TEMPORAL_TYPES
-    except NotImplementedError:
-        return False
-
-
 def load_or_create_tables(  # pylint: disable=too-many-arguments
     session: Session,
-    database_id: int,
+    database: Database,
     default_schema: Optional[str],
     tables: Set[Table],
     conditional_quote: Callable[[str], str],
-    engine: Engine,
 ) -> List[NewTable]:
     """
     Load or create new table model instances.
@@ -206,7 +193,7 @@ def load_or_create_tables(  # pylint: disable=too-many-arguments
     predicate = or_(
         *[
             and_(
-                NewTable.database_id == database_id,
+                NewTable.database_id == database.id,
                 NewTable.schema == table.schema,
                 NewTable.name == table.table,
             )
@@ -220,9 +207,10 @@ def load_or_create_tables(  # pylint: disable=too-many-arguments
     for table in tables:
         if (table.schema, table.table) not in existing:
             try:
-                inspector = inspect(engine)
-                column_metadata = inspector.get_columns(
-                    table.table, schema=table.schema
+                column_metadata = get_physical_table_metadata(
+                    database=database,
+                    table_name=table.table,
+                    schema_name=table.schema,
                 )
             except Exception:  # pylint: disable=broad-except
                 continue
@@ -231,7 +219,7 @@ def load_or_create_tables(  # pylint: disable=too-many-arguments
                     name=column["name"],
                     type=str(column["type"]),
                     expression=conditional_quote(column["name"]),
-                    is_temporal=is_column_type_temporal(column["type"]),
+                    is_temporal=column["is_dttm"],
                     is_aggregation=False,
                     is_physical=True,
                     is_spatial=False,
@@ -245,7 +233,7 @@ def load_or_create_tables(  # pylint: disable=too-many-arguments
                     name=table.table,
                     schema=table.schema,
                     catalog=None,
-                    database_id=database_id,
+                    database_id=database.id,
                     columns=columns,
                 )
             )