Skip to content

Commit

Permalink
fix: Retry SQLAlchemy engine creation for adapters without JSON SerDe…
Browse files Browse the repository at this point in the history
… support (#1949)

* fix: Retry SQLAlchemy engine creation for adapters without JSON SerDe support

* Add test

* Fix for SQLAlchemy 1
  • Loading branch information
edgarrmondragon authored Sep 11, 2023
1 parent 0fb0c1b commit fca070e
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 7 deletions.
Empty file.
32 changes: 32 additions & 0 deletions samples/sample_custom_sql_adapter/connector.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
from __future__ import annotations

import typing as t

from sqlalchemy.engine.default import DefaultDialect

if t.TYPE_CHECKING:
from types import ModuleType


class CustomSQLDialect(DefaultDialect):
"""Custom SQLite dialect that supports JSON."""

name = "myrdbms"

def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)

@classmethod
def import_dbapi(cls):
"""Import the sqlite3 DBAPI."""
import sqlite3

return sqlite3

@classmethod
def dbapi(cls) -> ModuleType: # type: ignore[override]
"""Return the DBAPI module.
NOTE: This is a legacy method that will stop being used by SQLAlchemy at some point.
""" # noqa: E501
return cls.import_dbapi()
21 changes: 15 additions & 6 deletions singer_sdk/connectors/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -319,12 +319,21 @@ def create_engine(self) -> Engine:
Returns:
A new SQLAlchemy Engine.
"""
return sqlalchemy.create_engine(
self.sqlalchemy_url,
echo=False,
json_serializer=self.serialize_json,
json_deserializer=self.deserialize_json,
)
try:
return sqlalchemy.create_engine(
self.sqlalchemy_url,
echo=False,
json_serializer=self.serialize_json,
json_deserializer=self.deserialize_json,
)
except TypeError:
self.logger.exception(
"Retrying engine creation with fewer arguments due to TypeError.",
)
return sqlalchemy.create_engine(
self.sqlalchemy_url,
echo=False,
)

def quote(self, name: str) -> str:
"""Quote a name if it needs quoting, using '.' as a name-part delimiter.
Expand Down
21 changes: 20 additions & 1 deletion tests/core/test_connector_sql.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
from __future__ import annotations

import typing as t
from decimal import Decimal
from unittest import mock

import pytest
import sqlalchemy
from sqlalchemy.dialects import sqlite
from sqlalchemy.dialects import registry, sqlite

from singer_sdk.connectors import SQLConnector
from singer_sdk.exceptions import ConfigValidationError

if t.TYPE_CHECKING:
from sqlalchemy.engine import Engine


def stringify(in_dict):
return {k: str(v) for k, v in in_dict.items()}
Expand Down Expand Up @@ -283,3 +287,18 @@ def test_engine_json_serialization(self, connector: SQLConnector):
(1, {"x": Decimal("1.0")}),
(2, {"x": Decimal("2.0"), "y": [1, 2, 3]}),
]


def test_adapter_without_json_serde():
registry.register(
"myrdbms",
"samples.sample_custom_sql_adapter.connector",
"CustomSQLDialect",
)

class CustomConnector(SQLConnector):
def create_engine(self) -> Engine:
return super().create_engine()

connector = CustomConnector(config={"sqlalchemy_url": "myrdbms:///"})
connector.create_engine()

0 comments on commit fca070e

Please sign in to comment.