Skip to content

Commit

Permalink
refactor: Added a class method to instantiate SQLToJSONSchema from …
Browse files Browse the repository at this point in the history
…the tap configuration
  • Loading branch information
edgarrmondragon committed Nov 28, 2024
1 parent 3d8e418 commit 0d5b11f
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 10 deletions.
32 changes: 26 additions & 6 deletions docs/guides/sql-tap.md
Original file line number Diff line number Diff line change
Expand Up @@ -24,15 +24,19 @@ from my_sqlalchemy_dialect import VectorType


class CustomSQLToJSONSchema(SQLToJSONSchema):
@SQLToJSONSchema.to_jsonschema.register
@functools.singledispatchmethod
def to_jsonschema(self, column_type):
return super().to_jsonschema(column_type)

@to_jsonschema.register
def custom_number_to_jsonschema(self, column_type: Numeric):
"""Override the default mapping for NUMERIC columns.
For example, a scale of 4 translates to a multipleOf 0.0001.
"""
return {"type": ["number"], "multipleOf": 10**-column_type.scale}

@SQLToJSONSchema.to_jsonschema.register(VectorType)
@to_jsonschema.register(VectorType)
def vector_to_json_schema(self, column_type):
"""Custom vector to JSON schema."""
return th.ArrayType(th.NumberType()).to_dict()
Expand All @@ -42,7 +46,7 @@ class CustomSQLToJSONSchema(SQLToJSONSchema):
You can also use a type annotation to specify the type of the column when registering a new method:
```python
@SQLToJSONSchema.to_jsonschema.register
@to_jsonschema.register
def vector_to_json_schema(self, column_type: VectorType):
return th.ArrayType(th.NumberType()).to_dict()
```
Expand All @@ -52,7 +56,23 @@ Then, you need to use your custom type mapping in your connector:

```python
class MyConnector(SQLConnector):
@functools.cached_property
def sql_to_jsonschema(self):
return CustomSQLToJSONSchema()
sql_to_jsonschema_converter = CustomSQLToJSONSchema
```

### Adapting the type mapping based on user configuration


If your type mapping depends on some user-defined configuration, you can also override the `from_config` method to pass the configuration to your custom type mapping:

```python
class ConfiguredSQLToJSONSchema(SQLToJSONSchema):
def __init__(self, *, my_custom_setting: str, **kwargs):
super().__init__(**kwargs)
self.my_custom_setting = my_custom_setting

@classmethod
def from_config(cls, config: dict):
return cls(my_custom_setting=config.get("my_custom_setting", "default_value"))
```

Then, you can use your custom type mapping in your connector as in the previous example.
25 changes: 24 additions & 1 deletion singer_sdk/connectors/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,26 @@ class SQLToJSONSchema:
This class provides a mapping from SQLAlchemy types to JSON Schema types.
.. versionadded:: 0.41.0
.. versionchanged:: 0.43.0
Added the :meth:`singer_sdk.connectors.sql.SQLToJSONSchema.from_config` class
method.
"""

@classmethod
def from_config(cls: type[SQLToJSONSchema], config: dict) -> SQLToJSONSchema: # noqa: ARG003
"""Create a new instance from a configuration dictionary.
Override this to instantiate this converter with values from the tap's
configuration dictionary.
Args:
config: The configuration dictionary.
Returns:
A new instance of the class.
"""
return cls()

@functools.singledispatchmethod
def to_jsonschema(self, column_type: sa.types.TypeEngine) -> dict: # noqa: ARG002, D102, PLR6301
return th.StringType.type_dict # type: ignore[no-any-return]
Expand Down Expand Up @@ -453,6 +471,11 @@ class SQLConnector: # noqa: PLR0904
#: The absolute maximum length for VARCHAR columns that the database supports.
max_varchar_length: int | None = None

#: The SQL-to-JSON type mapper class for this SQL connector. Override this property
#: with a subclass of :class:`~singer_sdk.connectors.sql.SQLToJSONSchema` to provide
#: a custom mapping for your SQL dialect.
sql_to_jsonschema_converter: type[SQLToJSONSchema] = SQLToJSONSchema

def __init__(
self,
config: dict | None = None,
Expand Down Expand Up @@ -493,7 +516,7 @@ def sql_to_jsonschema(self) -> SQLToJSONSchema:
.. versionadded:: 0.41.0
"""
return SQLToJSONSchema()
return self.sql_to_jsonschema_converter.from_config(self.config)

@functools.cached_property
def jsonschema_to_sql(self) -> JSONSchemaToSQL:
Expand Down
11 changes: 8 additions & 3 deletions tests/core/test_connector_sql.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import functools
import typing as t
from decimal import Decimal
from unittest import mock
Expand Down Expand Up @@ -452,15 +453,19 @@ def test_sql_to_json_schema_map(

def test_custom_type_to_jsonschema():
class MyMap(SQLToJSONSchema):
@SQLToJSONSchema.to_jsonschema.register
def custom_number_to_jsonschema(self, column_type: sa.types.NUMERIC) -> dict:
@functools.singledispatchmethod
def to_jsonschema(self, column_type: sa.types.TypeEngine):
return super().to_jsonschema(column_type)

@to_jsonschema.register
def custom_number_to_jsonschema(self, column_type: sa.types.Numeric) -> dict:
"""Custom number to JSON schema.
For example, a scale of 4 translates to a multipleOf 0.0001.
"""
return {"type": ["number"], "multipleOf": 10**-column_type.scale}

@SQLToJSONSchema.to_jsonschema.register(MyType)
@to_jsonschema.register(MyType)
def my_type_to_jsonschema(self, column_type) -> dict: # noqa: ARG002
return {"type": ["string"], "contentEncoding": "base64"}

Expand Down

0 comments on commit 0d5b11f

Please sign in to comment.