From 00d429e51e17c213db58d3665a6c963ababfc19e Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Mon, 10 May 2021 14:19:45 -0700 Subject: [PATCH] fix(ingest): add support for custom postgres types --- .../src/datahub/ingestion/source/postgres.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) 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):