From 74c5479926d89cebe5bad193123d8ecaff65f360 Mon Sep 17 00:00:00 2001 From: Smart-Codi Date: Thu, 2 Jun 2022 05:43:11 +1000 Subject: [PATCH] fix: datatype tracking issue on virtual dataset (#20088) * Fix datatype tracking issue on virtual dataset * fix pytest issue on postgresql --- superset/db_engine_specs/postgres.py | 9 +++++++++ tests/integration_tests/sqla_models_tests.py | 3 +-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index f81e2f7b3e4a6..c5ffc79ee7100 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -21,6 +21,7 @@ from typing import Any, Dict, List, Optional, Pattern, Tuple, TYPE_CHECKING from flask_babel import gettext as __ +from psycopg2.extensions import binary_types, string_types from sqlalchemy.dialects.postgresql import ARRAY, DOUBLE_PRECISION, ENUM, JSON from sqlalchemy.dialects.postgresql.base import PGInspector from sqlalchemy.types import String @@ -287,6 +288,14 @@ def get_column_spec( native_type, column_type_mappings=column_type_mappings ) + @classmethod + def get_datatype(cls, type_code: Any) -> Optional[str]: + types = binary_types.copy() + types.update(string_types) + if type_code in types: + return types[type_code].name + return None + @classmethod def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: """ diff --git a/tests/integration_tests/sqla_models_tests.py b/tests/integration_tests/sqla_models_tests.py index d23b95f53cd3d..6c5b6736d1a15 100644 --- a/tests/integration_tests/sqla_models_tests.py +++ b/tests/integration_tests/sqla_models_tests.py @@ -52,11 +52,10 @@ from .base_tests import SupersetTestCase - VIRTUAL_TABLE_INT_TYPES: Dict[str, Pattern[str]] = { "hive": re.compile(r"^INT_TYPE$"), "mysql": re.compile("^LONGLONG$"), - "postgresql": re.compile(r"^INT$"), + "postgresql": re.compile(r"^INTEGER$"), "presto": re.compile(r"^INTEGER$"), "sqlite": re.compile(r"^INT$"), }