Switch to new SQLAlchemy dialect for CrateDB

This includes the `FloatVector` SQLAlchemy type.

Author: amotl

CHANGES.md

@@ -3,6 +3,7 @@
 ## In progress
 - Add support for container types `ARRAY`, `OBJECT`, and `FLOAT_VECTOR`.
 - Improve write operations to be closer to `target-postgres`.
+- Switch to new SQLAlchemy dialect for CrateDB.
 
 ## 2023-12-08 v0.0.1
 - Make it work. It can run the canonical Meltano GitHub -> DB example.

pyproject.toml

@@ -89,10 +89,10 @@ dynamic = [
   "version",
 ]
 dependencies = [
-  "crate[sqlalchemy]<1",
   "cratedb-toolkit",
   "importlib-resources; python_version<'3.9'",                                                                    # "meltanolabs-target-postgres==0.0.9",
   "meltanolabs-target-postgres @ git+https://github.com/singer-contrib/meltanolabs-target-postgres.git@pgvector",
+  "sqlalchemy-cratedb[vector]",
 ]
 optional-dependencies.all = [
   "meltano-target-cratedb[vector]",

target_cratedb/connector.py

@@ -7,13 +7,14 @@
 from datetime import datetime
 
 import sqlalchemy as sa
-from crate.client.sqlalchemy.types import ObjectType, ObjectTypeImpl, _ObjectArray
 from singer_sdk import typing as th
 from singer_sdk.helpers._typing import is_array_type, is_boolean_type, is_integer_type, is_number_type, is_object_type
+from sqlalchemy_cratedb.type import FloatVector, ObjectType
+from sqlalchemy_cratedb.type.array import _ObjectArray
+from sqlalchemy_cratedb.type.object import ObjectTypeImpl
 from target_postgres.connector import NOTYPE, PostgresConnector
 
 from target_cratedb.sqlalchemy.patch import polyfill_refresh_after_dml_engine
-from target_cratedb.sqlalchemy.vector import FloatVector
 
 
 class CrateDBConnector(PostgresConnector):
@@ -226,6 +227,9 @@ def _get_type_sort_key(
             if isinstance(sql_type, NOTYPE):
                 return 0, _len
 
+            if not hasattr(sql_type, "python_type"):
+                raise TypeError(f"Resolving type for sort key failed: {sql_type}")
+
             _pytype = t.cast(type, sql_type.python_type)
             if issubclass(_pytype, (str, bytes)):
                 return 900, _len

target_cratedb/sqlalchemy/patch.py

@@ -1,11 +1,11 @@
 from _decimal import Decimal
 from datetime import datetime
+from typing import Any, Union
 
+import crate.client.http
 import sqlalchemy as sa
-from crate.client.http import CrateJsonEncoder
-from crate.client.sqlalchemy.dialect import ARRAY, TYPES_MAP, DateTime
-from crate.client.sqlalchemy.types import _ObjectArray
-from sqlalchemy.sql import sqltypes
+from sqlalchemy_cratedb.dialect import TYPES_MAP, DateTime
+from sqlalchemy_cratedb.type.array import _ObjectArray
 
 
 def patch_sqlalchemy():
@@ -19,20 +19,21 @@ def patch_types():
 
     TODO: Upstream to crate-python.
     """
-    TYPES_MAP["bigint"] = sqltypes.BIGINT
-    TYPES_MAP["bigint_array"] = ARRAY(sqltypes.BIGINT)
-    TYPES_MAP["long"] = sqltypes.BIGINT
-    TYPES_MAP["long_array"] = ARRAY(sqltypes.BIGINT)
-    TYPES_MAP["real"] = sqltypes.DOUBLE
-    TYPES_MAP["real_array"] = ARRAY(sqltypes.DOUBLE)
-    TYPES_MAP["timestamp without time zone"] = sqltypes.TIMESTAMP
-    TYPES_MAP["timestamp with time zone"] = sqltypes.TIMESTAMP
+    # abc()
+    TYPES_MAP["bigint"] = sa.BIGINT
+    TYPES_MAP["bigint_array"] = sa.ARRAY(sa.BIGINT)
+    TYPES_MAP["long"] = sa.BIGINT
+    TYPES_MAP["long_array"] = sa.ARRAY(sa.BIGINT)
+    TYPES_MAP["real"] = sa.DOUBLE
+    TYPES_MAP["real_array"] = sa.ARRAY(sa.DOUBLE)
+    TYPES_MAP["timestamp without time zone"] = sa.TIMESTAMP
+    TYPES_MAP["timestamp with time zone"] = sa.TIMESTAMP
 
     # TODO: Can `ARRAY` be inherited from PostgreSQL's
     #       `ARRAY`, to make type checking work?
 
     def as_generic(self, allow_nulltype: bool = False):
-        return sqltypes.ARRAY
+        return sa.ARRAY
 
     _ObjectArray.as_generic = as_generic
 
@@ -58,14 +59,14 @@ def patch_json_encoder():
     TODO: Upstream to crate-python.
     """
 
-    json_encoder_default = CrateJsonEncoder.default
+    json_encoder_default = crate.client.http.json_encoder
 
-    def default(self, o):
-        if isinstance(o, Decimal):
-            return float(o)
-        return json_encoder_default(o)
+    def json_encoder_new(obj: Any) -> Union[int, str, float]:
+        if isinstance(obj, Decimal):
+            return float(obj)
+        return json_encoder_default(obj)
 
-    CrateJsonEncoder.default = default
+    crate.client.http.json_encoder = json_encoder_new
 
 
 def polyfill_refresh_after_dml_engine(engine: sa.Engine):

target_cratedb/sqlalchemy/vector.py

@@ -9,9 +9,10 @@
 import jsonschema
 import pytest
 import sqlalchemy as sa
-from crate.client.sqlalchemy.types import ObjectTypeImpl
 from singer_sdk.exceptions import MissingKeyPropertiesError
 from singer_sdk.testing import sync_end_to_end
+from sqlalchemy_cratedb.type import FloatVector
+from sqlalchemy_cratedb.type.object import ObjectTypeImpl
 from target_postgres.tests.samples.aapl.aapl import Fundamentals
 from target_postgres.tests.samples.sample_tap_countries.countries_tap import (
     SampleTapCountries,
@@ -21,7 +22,6 @@
 from target_cratedb.connector import CrateDBConnector
 from target_cratedb.sinks import MELTANO_CRATEDB_STRATEGY_DIRECT
 from target_cratedb.sqlalchemy.patch import polyfill_refresh_after_dml_engine
-from target_cratedb.sqlalchemy.vector import FloatVector
 from target_cratedb.target import TargetCrateDB
 
 try: