diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index aebfad0781bd5..d7d62db2a7e0e 100644 --- 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 index a2b54201d6b89..f8ed7a956704a 100644 --- 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, ) )