diff --git a/tap_postgres/client.py b/tap_postgres/client.py index a6ae564c..3572a0f5 100644 --- a/tap_postgres/client.py +++ b/tap_postgres/client.py @@ -10,7 +10,6 @@ import json import select import typing as t -from functools import cached_property from types import MappingProxyType from typing import TYPE_CHECKING, Any @@ -30,7 +29,6 @@ from collections.abc import Iterable, Mapping from singer_sdk.helpers.types import Context - from sqlalchemy.dialects import postgresql from sqlalchemy.engine import Engine from sqlalchemy.engine.reflection import Inspector @@ -44,7 +42,12 @@ def __init__(self, dates_as_string: bool, json_as_object: bool, *args, **kwargs) self.dates_as_string = dates_as_string self.json_as_object = json_as_object - @SQLToJSONSchema.to_jsonschema.register # type: ignore[attr-defined] + @functools.singledispatchmethod + def to_jsonschema(self, column_type: Any) -> dict: + """Customize the JSON Schema for Postgres types.""" + return super().to_jsonschema(column_type) + + @to_jsonschema.register def array_to_jsonschema(self, column_type: postgresql.ARRAY) -> dict: """Override the default mapping for NUMERIC columns. @@ -55,21 +58,21 @@ def array_to_jsonschema(self, column_type: postgresql.ARRAY) -> dict: "items": self.to_jsonschema(column_type.item_type), } - @SQLToJSONSchema.to_jsonschema.register # type: ignore[attr-defined] + @to_jsonschema.register def json_to_jsonschema(self, column_type: postgresql.JSON) -> dict: """Override the default mapping for JSON and JSONB columns.""" if self.json_as_object: return {"type": ["object", "null"]} return {"type": ["string", "number", "integer", "array", "object", "boolean"]} - @SQLToJSONSchema.to_jsonschema.register # type: ignore[attr-defined] + @to_jsonschema.register def datetime_to_jsonschema(self, column_type: sqlalchemy.types.DateTime) -> dict: """Override the default mapping for DATETIME columns.""" if self.dates_as_string: return {"type": ["string", "null"]} return super().datetime_to_jsonschema(column_type) - @SQLToJSONSchema.to_jsonschema.register # type: ignore[attr-defined] + @to_jsonschema.register def date_to_jsonschema(self, column_type: sqlalchemy.types.Date) -> dict: """Override the default mapping for DATE columns.""" if self.dates_as_string: @@ -276,7 +279,7 @@ def config(self) -> Mapping[str, Any]: """Return a read-only config dictionary.""" return MappingProxyType(self._config) - @cached_property + @functools.cached_property def schema(self) -> dict: """Override schema for log-based replication adding _sdc columns.""" schema_dict = t.cast(dict, self._singer_catalog_entry.schema.to_dict())