diff --git a/metadata-ingestion/src/datahub/ingestion/source/postgres.py b/metadata-ingestion/src/datahub/ingestion/source/postgres.py index 85a874373f3fa2..c8654c23a078e6 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/postgres.py +++ b/metadata-ingestion/src/datahub/ingestion/source/postgres.py @@ -1,5 +1,6 @@ # This import verifies that the dependencies are available. import psycopg2 # noqa: F401 +import sqlalchemy.dialects.postgresql as custom_types # GeoAlchemy adds support for PostGIS extensions in SQLAlchemy. In order to # activate it, we must import it so that it can hook into SQLAlchemy. While @@ -8,7 +9,18 @@ # https://geoalchemy-2.readthedocs.io/en/latest/core_tutorial.html#reflecting-tables. from geoalchemy2 import Geometry # noqa: F401 -from .sql_common import BasicSQLAlchemyConfig, SQLAlchemySource +from datahub.metadata.com.linkedin.pegasus2avro.schema import ( + ArrayTypeClass, + BytesTypeClass, + MapTypeClass, +) + +from .sql_common import BasicSQLAlchemyConfig, SQLAlchemySource, register_custom_type + +register_custom_type(custom_types.ARRAY, ArrayTypeClass) +register_custom_type(custom_types.JSON, BytesTypeClass) +register_custom_type(custom_types.JSONB, BytesTypeClass) +register_custom_type(custom_types.HSTORE, MapTypeClass) class PostgresConfig(BasicSQLAlchemyConfig):