Merge pull request #242 from Brooke-white/datatype-inspection

Support inspection of Redshift datatypes

CHANGES.rst

@@ -1,7 +1,8 @@
 0.8.9 (unreleased)
 ------------------
 
-- Nothing changed yet.
+- Support inspection of Redshift datatypes
+  (`Pull #242 <https://github.com/sqlalchemy-redshift/sqlalchemy-redshift/pull/242>`_)
 
 
 0.8.8 (2021-11-03)

sqlalchemy_redshift/dialect.py

@@ -11,6 +11,8 @@
 )
 from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
 from sqlalchemy.dialects.postgresql.psycopg2cffi import PGDialect_psycopg2cffi
+from sqlalchemy.engine.default import DefaultDialect
+from sqlalchemy.sql.type_api import TypeEngine
 from sqlalchemy.engine import reflection
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.expression import (
@@ -178,7 +180,18 @@ class RedshiftImpl(postgresql.PostgresqlImpl):
 ])
 
 
-class TIMESTAMPTZ(sa.dialects.postgresql.TIMESTAMP):
+class RedshiftTypeEngine(TypeEngine):
+
+    def _default_dialect(self, default=None):
+        """
+        Returns the default dialect used for TypeEngine compilation yielding String result.
+
+        :meth:`~sqlalchemy.sql.type_api.TypeEngine.compile`
+        """
+        return RedshiftDialectMixin()
+
+
+class TIMESTAMPTZ(RedshiftTypeEngine, sa.dialects.postgresql.TIMESTAMP):
     """
     Redshift defines a TIMTESTAMPTZ column type as an alias
     of TIMESTAMP WITH TIME ZONE.
@@ -192,11 +205,14 @@ class TIMESTAMPTZ(sa.dialects.postgresql.TIMESTAMP):
 
     __visit_name__ = 'TIMESTAMPTZ'
 
-    def __init__(self):
-        super(TIMESTAMPTZ, self).__init__(timezone=True)
+    def __init__(self, timezone=True, precision=None):
+        # timezone param must be present as it's provided in base class so the object
+        # can be instantiated with kwargs
+        # see :meth:`~sqlalchemy.dialects.postgresql.base.PGDialect._get_column_info`
+        super(TIMESTAMPTZ, self).__init__(timezone=True, precision=precision)
 
 
-class TIMETZ(sa.dialects.postgresql.TIME):
+class TIMETZ(RedshiftTypeEngine, sa.dialects.postgresql.TIME):
     """
     Redshift defines a TIMTETZ column type as an alias
     of TIME WITH TIME ZONE.
@@ -210,11 +226,14 @@ class TIMETZ(sa.dialects.postgresql.TIME):
 
     __visit_name__ = 'TIMETZ'
 
-    def __init__(self):
-        super(TIMETZ, self).__init__(timezone=True)
+    def __init__(self, timezone=True, precision=None):
+        # timezone param must be present as it's provided in base class so the object
+        # can be instantiated with kwargs
+        # see :meth:`~sqlalchemy.dialects.postgresql.base.PGDialect._get_column_info`
+        super(TIMETZ, self).__init__(timezone=True, precision=precision)
 
 
-class GEOMETRY(sa.dialects.postgresql.TEXT):
+class GEOMETRY(RedshiftTypeEngine, sa.dialects.postgresql.TEXT):
     """
     Redshift defines a GEOMETRY column type
     https://docs.aws.amazon.com/redshift/latest/dg/c_Supported_data_types.html
@@ -233,7 +252,7 @@ def get_dbapi_type(self, dbapi):
         return dbapi.GEOMETRY
 
 
-class SUPER(sa.dialects.postgresql.TEXT):
+class SUPER(RedshiftTypeEngine, sa.dialects.postgresql.TEXT):
     """
     Redshift defines a SUPER column type
     https://docs.aws.amazon.com/redshift/latest/dg/c_Supported_data_types.html
@@ -260,6 +279,14 @@ def process_bind_param(self, value, dialect):
             return json.dumps(value)
         return value
 
+# Mapping for database schema inspection of Amazon Redshift datatypes
+REDSHIFT_ISCHEMA_NAMES = {
+    "geometry": GEOMETRY,
+    "super": SUPER,
+    "time with time zone": TIMETZ,
+    "timestamp with time zone": TIMESTAMPTZ,
+}
+
 
 class RelationKey(namedtuple('RelationKey', ('name', 'schema'))):
     """
@@ -493,7 +520,7 @@ class RedshiftIdentifierPreparer(PGIdentifierPreparer):
     reserved_words = RESERVED_WORDS
 
 
-class RedshiftDialectMixin(object):
+class RedshiftDialectMixin(DefaultDialect):
     """
     Define Redshift-specific behavior.
 
@@ -536,6 +563,16 @@ def __init__(self, *args, **kw):
         # Redshift does not support user-created domains.
         self._domains = None
 
+    @property
+    def ischema_names(self):
+        """
+        Returns information about datatypes supported by Amazon Redshift.
+
+        Used in
+        :meth:`~sqlalchemy.engine.dialects.postgresql.base.PGDialect._get_column_info`.
+        """
+        return {**super(RedshiftDialectMixin, self).ischema_names, **REDSHIFT_ISCHEMA_NAMES}
+
     @reflection.cache
     def get_columns(self, connection, table_name, schema=None, **kw):
         """

tests/test_dialect_types.py

@@ -1,6 +1,8 @@
 import pytest
 import sqlalchemy_redshift.dialect
 import sqlalchemy
+from sqlalchemy.engine import reflection
+from sqlalchemy import MetaData
 
 
 def test_defined_types():
@@ -130,3 +132,39 @@ def test_custom_types_ddl_generation(
     create_table = sqlalchemy.schema.CreateTable(table)
     actual = compiler.process(create_table)
     assert expected == actual
+
+
+redshift_specific_datatypes = [
+    sqlalchemy_redshift.dialect.GEOMETRY,
+    sqlalchemy_redshift.dialect.SUPER,
+    sqlalchemy_redshift.dialect.TIMETZ,
+    sqlalchemy_redshift.dialect.TIMESTAMPTZ
+]
+
+
+@pytest.mark.parametrize("custom_datatype", redshift_specific_datatypes)
+def test_custom_types_reflection_inspection(
+        custom_datatype, redshift_engine
+):
+    metadata = MetaData(bind=redshift_engine)
+    sqlalchemy.Table(
+        't1',
+        metadata,
+        sqlalchemy.Column('id', sqlalchemy.INTEGER, primary_key=True),
+        sqlalchemy.Column('name', sqlalchemy.String),
+        sqlalchemy.Column('test_col', custom_datatype),
+        schema='public'
+    )
+    metadata.create_all()
+    inspect = reflection.Inspector.from_engine(redshift_engine)
+
+    actual = inspect.get_columns(table_name='t1', schema='public')
+    assert len(actual) == 3
+    assert isinstance(actual[2]['type'], custom_datatype)
+
+
+@pytest.mark.parametrize("custom_datatype", redshift_specific_datatypes)
+def test_custom_type_compilation(custom_datatype):
+    dt = custom_datatype()
+    compiled_dt = dt.compile()
+    assert compiled_dt == dt.__visit_name__