Skip to content

Commit

Permalink
refactor: Use sql_to_jsonschema_converter
Browse files Browse the repository at this point in the history
  • Loading branch information
edgarrmondragon committed Nov 28, 2024
1 parent bea709e commit d94eb87
Showing 1 changed file with 12 additions and 10 deletions.
22 changes: 12 additions & 10 deletions tap_postgres/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,20 @@
class PostgresSQLToJSONSchema(SQLToJSONSchema):
"""Custom SQL to JSON Schema conversion for Postgres."""

def __init__(self, dates_as_string: bool, json_as_object: bool, *args, **kwargs):
def __init__(self, *, dates_as_string: bool, json_as_object: bool, **kwargs):
"""Initialize the SQL to JSON Schema converter."""
super().__init__(*args, **kwargs)
super().__init__(**kwargs)
self.dates_as_string = dates_as_string
self.json_as_object = json_as_object

@classmethod
def from_config(cls, config: dict) -> PostgresSQLToJSONSchema:
"""Instantiate the SQL to JSON Schema converter from a config dictionary."""
return cls(
dates_as_string=config["dates_as_string"],
json_as_object=config["json_as_object"],
)

@functools.singledispatchmethod
def to_jsonschema(self, column_type: t.Any) -> dict:
"""Customize the JSON Schema for Postgres types."""
Expand Down Expand Up @@ -132,6 +140,8 @@ def patched_conform(elem: t.Any, property_schema: dict) -> t.Any:
class PostgresConnector(SQLConnector):
"""Connects to the Postgres SQL source."""

sql_to_jsonschema_converter = PostgresSQLToJSONSchema

def __init__(
self,
config: dict | None = None,
Expand Down Expand Up @@ -160,14 +170,6 @@ def __init__(

super().__init__(config=config, sqlalchemy_url=sqlalchemy_url)

@functools.cached_property
def sql_to_jsonschema(self):
"""Return a mapping of SQL types to JSON Schema types."""
return PostgresSQLToJSONSchema(
dates_as_string=self.config["dates_as_string"],
json_as_object=self.config["json_as_object"],
)

def get_schema_names(self, engine: Engine, inspected: Inspector) -> list[str]:
"""Return a list of schema names in DB, or overrides with user-provided values.
Expand Down

0 comments on commit d94eb87

Please sign in to comment.