sqlalchemy-cratedb 0.41.0.dev0__py3-none-any.whl

sqlalchemy_cratedb/support/polyfill.py

@@ -0,0 +1,130 @@
+import typing as t
+
+import sqlalchemy as sa
+from sqlalchemy.event import listen
+
+from sqlalchemy_cratedb.support.util import refresh_dirty, refresh_table
+
+
+def patch_autoincrement_timestamp():
+    """
+    Configure SQLAlchemy model columns with an alternative to `autoincrement=True`.
+    Use the current timestamp instead.
+
+    This is used by CrateDB's MLflow adapter.
+
+    TODO: Maybe enable through a dialect parameter `crate_polyfill_autoincrement` or such.
+    """
+    import sqlalchemy.sql.schema as schema
+
+    init_dist = schema.Column.__init__
+
+    def __init__(self, *args, **kwargs):
+        if "autoincrement" in kwargs:
+            del kwargs["autoincrement"]
+            if "default" not in kwargs:
+                kwargs["default"] = sa.func.now()
+        init_dist(self, *args, **kwargs)
+
+    schema.Column.__init__ = __init__  # type: ignore[method-assign]
+
+
+def check_uniqueness_factory(sa_entity, *attribute_names):
+    """
+    Run a manual column value uniqueness check on a table, and raise an IntegrityError if applicable.
+
+    CrateDB does not support the UNIQUE constraint on columns. This attempts to emulate it.
+
+    https://github.com/crate/sqlalchemy-cratedb/issues/76
+
+    This is used by CrateDB's MLflow adapter.
+
+    TODO: Maybe enable through a dialect parameter `crate_polyfill_unique` or such.
+    """  # noqa: E501
+
+    # Synthesize a canonical "name" for the constraint,
+    # composed of all column names involved.
+    constraint_name: str = "-".join(attribute_names)
+
+    def check_uniqueness(mapper, connection, target):
+        from sqlalchemy.exc import IntegrityError
+
+        if isinstance(target, sa_entity):
+            # TODO: How to use `session.query(SqlExperiment)` here?
+            stmt = mapper.selectable.select()
+            for attribute_name in attribute_names:
+                stmt = stmt.filter(
+                    getattr(sa_entity, attribute_name) == getattr(target, attribute_name)
+                )
+            stmt = stmt.compile(bind=connection.engine)
+            results = connection.execute(stmt)
+            if results.rowcount > 0:
+                raise IntegrityError(
+                    statement=stmt,
+                    params=[],
+                    orig=Exception(
+                        f"DuplicateKeyException in table '{target.__tablename__}' "
+                        f"on constraint '{constraint_name}'"
+                    ),
+                )
+
+    return check_uniqueness
+
+
+def refresh_after_dml_session(session: sa.orm.Session):
+    """
+    Run `REFRESH TABLE` after each DML operation (INSERT, UPDATE, DELETE).
+
+    CrateDB is eventually consistent, i.e. write operations are not flushed to
+    disk immediately, so readers may see stale data. In a traditional OLTP-like
+    application, this is not applicable.
+
+    This SQLAlchemy extension makes sure that data is synchronized after each
+    operation manipulating data.
+
+    > `after_{insert,update,delete}` events only apply to the session flush operation
+    > and do not apply to the ORM DML operations described at ORM-Enabled INSERT,
+    > UPDATE, and DELETE statements. To intercept ORM DML events, use
+    > `SessionEvents.do_orm_execute().`
+    > -- https://docs.sqlalchemy.org/en/20/orm/events.html#sqlalchemy.orm.MapperEvents.after_insert
+
+    > Intercept statement executions that occur on behalf of an ORM Session object.
+    > -- https://docs.sqlalchemy.org/en/20/orm/events.html#sqlalchemy.orm.SessionEvents.do_orm_execute
+
+    > Execute after flush has completed, but before commit has been called.
+    > -- https://docs.sqlalchemy.org/en/20/orm/events.html#sqlalchemy.orm.SessionEvents.after_flush
+
+    This is used by CrateDB's LangChain adapter.
+
+    TODO: Maybe enable through a dialect parameter `crate_dml_refresh` or such.
+    """  # noqa: E501
+    listen(session, "after_flush", refresh_dirty)
+
+
+def refresh_after_dml_engine(engine: sa.engine.Engine):
+    """
+    Run `REFRESH TABLE` after each DML operation (INSERT, UPDATE, DELETE).
+
+    This is used by CrateDB's Singer/Meltano and `rdflib-sqlalchemy` adapters.
+    """
+
+    def receive_after_execute(
+        conn: sa.engine.Connection, clauseelement, multiparams, params, execution_options, result
+    ):
+        if isinstance(clauseelement, (sa.sql.Insert, sa.sql.Update, sa.sql.Delete)):
+            if not isinstance(clauseelement.table, sa.sql.Join):
+                refresh_table(conn, clauseelement.table)
+
+    sa.event.listen(engine, "after_execute", receive_after_execute)
+
+
+def refresh_after_dml(engine_or_session: t.Union[sa.engine.Engine, sa.orm.Session]):
+    """
+    Run `REFRESH TABLE` after each DML operation (INSERT, UPDATE, DELETE).
+    """
+    if isinstance(engine_or_session, sa.engine.Engine):
+        refresh_after_dml_engine(engine_or_session)
+    elif isinstance(engine_or_session, (sa.orm.Session, sa.orm.scoping.scoped_session)):
+        refresh_after_dml_session(engine_or_session)
+    else:
+        raise TypeError(f"Unknown type: {type(engine_or_session)}")

sqlalchemy_cratedb/support/util.py

@@ -0,0 +1,82 @@
+import itertools
+import typing as t
+
+import sqlalchemy as sa
+
+from sqlalchemy_cratedb.dialect import CrateDialect
+
+if t.TYPE_CHECKING:
+    try:
+        from sqlalchemy.orm import DeclarativeBase
+    except ImportError:
+        pass
+
+
+# An instance of the dialect used for quoting purposes.
+identifier_preparer = CrateDialect().identifier_preparer
+
+
+def refresh_table(
+    connection, target: t.Union[str, "DeclarativeBase", "sa.sql.selectable.TableClause"]
+):
+    """
+    Invoke a `REFRESH TABLE` statement.
+    """
+
+    if isinstance(target, sa.sql.selectable.TableClause):
+        full_table_name = f'"{target.name}"'
+        if target.schema is not None:
+            full_table_name = f'"{target.schema}".' + full_table_name
+    elif hasattr(target, "__tablename__"):
+        full_table_name = target.__tablename__
+    else:
+        full_table_name = target
+
+    sql = f"REFRESH TABLE {full_table_name}"
+    connection.execute(sa.text(sql))
+
+
+def refresh_dirty(session, flush_context=None):
+    """
+    Invoke a `REFRESH TABLE` statement on each table entity flagged as "dirty".
+
+    SQLAlchemy event handler for the 'after_flush' event,
+    invoking `REFRESH TABLE` on each table which has been modified.
+    """
+    dirty_entities = itertools.chain(session.new, session.dirty, session.deleted)
+    dirty_classes = {entity.__class__ for entity in dirty_entities}
+    for class_ in dirty_classes:
+        refresh_table(session, class_)
+
+
+def quote_relation_name(ident: str) -> str:
+    """
+    Quote a simple or full-qualified table/relation name, when needed.
+
+    Simple:         <table>
+    Full-qualified: <schema>.<table>
+
+    Happy path examples:
+
+        foo => foo
+        Foo => "Foo"
+        "Foo" => "Foo"
+        foo.bar => foo.bar
+        foo-bar.baz_qux => "foo-bar".baz_qux
+
+    Such input strings will not be modified:
+
+        "foo.bar" => "foo.bar"
+    """
+
+    # If a quote exists at the beginning or the end of the input string,
+    # let's consider that the relation name has been quoted already.
+    if ident.startswith('"') or ident.endswith('"'):
+        return ident
+
+    # If a dot is included, it's a full-qualified identifier like <schema>.<table>.
+    # It needs to be split, in order to apply identifier quoting properly.
+    parts = ident.split(".")
+    if len(parts) > 3:
+        raise ValueError(f"Invalid relation name, too many parts: {ident}")
+    return ".".join(map(identifier_preparer.quote, parts))

sqlalchemy_cratedb/type/__init__.py

@@ -0,0 +1,13 @@
+from .array import ObjectArray
+from .geo import Geopoint, Geoshape
+from .object import ObjectType
+from .vector import FloatVector, knn_match
+
+__all__ = [
+    Geopoint,
+    Geoshape,
+    ObjectArray,
+    ObjectType,
+    FloatVector,
+    knn_match,
+]

sqlalchemy_cratedb/type/array.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8; -*-
+#
+# Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
+# license agreements.  See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.  Crate licenses
+# this file to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.  You may
+# obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+# License for the specific language governing permissions and limitations
+# under the License.
+#
+# However, if you have executed another commercial license agreement
+# with Crate these terms will supersede the license and you may use the
+# software solely pursuant to the terms of the relevant commercial agreement.
+
+# ruff: noqa: A005  # Module `array` shadows a Python standard-library module
+
+import sqlalchemy.types as sqltypes
+from sqlalchemy.ext.mutable import Mutable
+from sqlalchemy.sql import default_comparator, expression, operators
+
+
+class MutableList(Mutable, list):
+    @classmethod
+    def coerce(cls, key, value):
+        """Convert plain list to MutableList"""
+        if not isinstance(value, MutableList):
+            if isinstance(value, list):
+                return MutableList(value)
+            elif value is None:
+                return value
+            else:
+                return MutableList([value])
+        else:
+            return value
+
+    def __init__(self, initval=None):
+        list.__init__(self, initval or [])
+
+    def __setitem__(self, key, value):
+        list.__setitem__(self, key, value)
+        self.changed()
+
+    def __eq__(self, other):
+        return list.__eq__(self, other)
+
+    def append(self, item):
+        list.append(self, item)
+        self.changed()
+
+    def insert(self, idx, item):
+        list.insert(self, idx, item)
+        self.changed()
+
+    def extend(self, iterable):
+        list.extend(self, iterable)
+        self.changed()
+
+    def pop(self, index=-1):
+        list.pop(self, index)
+        self.changed()
+
+    def remove(self, item):
+        list.remove(self, item)
+        self.changed()
+
+
+class Any(expression.ColumnElement):
+    """Represent the clause ``left operator ANY (right)``.  ``right`` must be
+    an array expression.
+
+    copied from postgresql dialect
+
+    .. seealso::
+
+        :class:`sqlalchemy.dialects.postgresql.ARRAY`
+
+        :meth:`sqlalchemy.dialects.postgresql.ARRAY.Comparator.any`
+            ARRAY-bound method
+
+    """
+
+    __visit_name__ = "any"
+    inherit_cache = True
+
+    def __init__(self, left, right, operator=operators.eq):
+        self.type = sqltypes.Boolean()
+        self.left = expression.literal(left)
+        self.right = right
+        self.operator = operator
+
+
+class _ObjectArray(sqltypes.UserDefinedType):
+    cache_ok = True
+
+    class Comparator(sqltypes.TypeEngine.Comparator):
+        def __getitem__(self, key):
+            return default_comparator._binary_operate(self.expr, operators.getitem, key)
+
+        def any(self, other, operator=operators.eq):
+            """Return ``other operator ANY (array)`` clause.
+
+            Argument places are switched, because ANY requires array
+            expression to be on the right hand-side.
+
+            E.g.::
+
+                from sqlalchemy.sql import operators
+
+                conn.execute(
+                    select([table.c.data]).where(
+                            table.c.data.any(7, operator=operators.lt)
+                        )
+                )
+
+            :param other: expression to be compared
+            :param operator: an operator object from the
+             :mod:`sqlalchemy.sql.operators`
+             package, defaults to :func:`.operators.eq`.
+
+            .. seealso::
+
+                :class:`.postgresql.Any`
+
+                :meth:`.postgresql.ARRAY.Comparator.all`
+
+            """
+            return Any(other, self.expr, operator=operator)
+
+    type = MutableList
+    comparator_factory = Comparator
+
+    def get_col_spec(self, **kws):
+        return "ARRAY(OBJECT)"
+
+
+ObjectArray = MutableList.as_mutable(_ObjectArray)

sqlalchemy_cratedb/type/geo.py

@@ -0,0 +1,43 @@
+import geojson
+from sqlalchemy import types as sqltypes
+from sqlalchemy.sql import default_comparator, operators
+
+
+class Geopoint(sqltypes.UserDefinedType):
+    cache_ok = True
+
+    class Comparator(sqltypes.TypeEngine.Comparator):
+        def __getitem__(self, key):
+            return default_comparator._binary_operate(self.expr, operators.getitem, key)
+
+    def get_col_spec(self):
+        return "GEO_POINT"
+
+    def bind_processor(self, dialect):
+        def process(value):
+            if isinstance(value, geojson.Point):
+                return value.coordinates
+            return value
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        return tuple
+
+    comparator_factory = Comparator
+
+
+class Geoshape(sqltypes.UserDefinedType):
+    cache_ok = True
+
+    class Comparator(sqltypes.TypeEngine.Comparator):
+        def __getitem__(self, key):
+            return default_comparator._binary_operate(self.expr, operators.getitem, key)
+
+    def get_col_spec(self):
+        return "GEO_SHAPE"
+
+    def result_processor(self, dialect, coltype):
+        return geojson.GeoJSON.to_instance
+
+    comparator_factory = Comparator

sqlalchemy_cratedb/type/object.py

@@ -0,0 +1,94 @@
+import warnings
+
+from sqlalchemy import types as sqltypes
+from sqlalchemy.ext.mutable import Mutable
+
+
+class MutableDict(Mutable, dict):
+    @classmethod
+    def coerce(cls, key, value):
+        "Convert plain dictionaries to MutableDict."
+
+        if not isinstance(value, MutableDict):
+            if isinstance(value, dict):
+                return MutableDict(value)
+
+            # this call will raise ValueError
+            return Mutable.coerce(key, value)
+        else:
+            return value
+
+    def __init__(self, initval=None, to_update=None, root_change_key=None):
+        initval = initval or {}
+        self._changed_keys = set()
+        self._deleted_keys = set()
+        self._overwrite_key = root_change_key
+        self.to_update = self if to_update is None else to_update
+        for k in initval:
+            initval[k] = self._convert_dict(
+                initval[k], overwrite_key=k if self._overwrite_key is None else self._overwrite_key
+            )
+        dict.__init__(self, initval)
+
+    def __setitem__(self, key, value):
+        value = self._convert_dict(
+            value, key if self._overwrite_key is None else self._overwrite_key
+        )
+        dict.__setitem__(self, key, value)
+        self.to_update.on_key_changed(key if self._overwrite_key is None else self._overwrite_key)
+
+    def __delitem__(self, key):
+        dict.__delitem__(self, key)
+        # add the key to the deleted keys if this is the root object
+        # otherwise update on root object
+        if self._overwrite_key is None:
+            self._deleted_keys.add(key)
+            self.changed()
+        else:
+            self.to_update.on_key_changed(self._overwrite_key)
+
+    def on_key_changed(self, key):
+        self._deleted_keys.discard(key)
+        self._changed_keys.add(key)
+        self.changed()
+
+    def _convert_dict(self, value, overwrite_key):
+        if isinstance(value, dict) and not isinstance(value, MutableDict):
+            return MutableDict(value, self.to_update, overwrite_key)
+        return value
+
+    def __eq__(self, other):
+        return dict.__eq__(self, other)
+
+
+class ObjectTypeImpl(sqltypes.UserDefinedType, sqltypes.JSON):
+    __visit_name__ = "OBJECT"
+
+    cache_ok = False
+    none_as_null = False
+
+
+# Designated name to refer to. `Object` is too ambiguous.
+ObjectType = MutableDict.as_mutable(ObjectTypeImpl)
+
+# Backward-compatibility aliases.
+_deprecated_Craty = ObjectType
+_deprecated_Object = ObjectType
+
+# https://www.lesinskis.com/deprecating-module-scope-variables.html
+deprecated_names = ["Craty", "Object"]
+
+
+def __getattr__(name):
+    if name in deprecated_names:
+        warnings.warn(
+            f"{name} is deprecated and will be removed in future releases. "
+            f"Please use ObjectType instead.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        return globals()[f"_deprecated_{name}"]
+    raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
+__all__ = deprecated_names + ["ObjectType"]

sqlalchemy_cratedb/type/vector.py

@@ -0,0 +1,176 @@
+"""
+## About
+SQLAlchemy data type implementation for CrateDB's `FLOAT_VECTOR` type.
+
+## References
+- https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#float-vector
+- https://cratedb.com/docs/crate/reference/en/latest/general/builtins/scalar-functions.html#scalar-knn-match
+
+## Details
+The implementation is based on SQLAlchemy's `TypeDecorator`, and also
+offers compiler support.
+
+## Notes
+CrateDB currently only supports the similarity function `VectorSimilarityFunction.EUCLIDEAN`.
+-- https://github.com/crate/crate/blob/5.5.1/server/src/main/java/io/crate/types/FloatVectorType.java#L55
+
+pgvector use a comparator to apply different similarity functions as operators,
+see `pgvector.sqlalchemy.Vector.comparator_factory`.
+
+<->: l2/euclidean_distance
+<#>: max_inner_product
+<=>: cosine_distance
+
+## Backlog
+- After dropping support for SQLAlchemy 1.3, use
+  `class FloatVector(sa.TypeDecorator[t.Sequence[float]]):`
+
+## Origin
+This module is based on the corresponding pgvector implementation
+by Andrew Kane. Thank you.
+
+The MIT License (MIT)
+Copyright (c) 2021-2023 Andrew Kane
+https://github.com/pgvector/pgvector-python
+"""
+
+import typing as t
+
+if t.TYPE_CHECKING:
+    import numpy.typing as npt  # pragma: no cover
+
+import sqlalchemy as sa
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.expression import ColumnElement, literal
+
+__all__ = [
+    "from_db",
+    "knn_match",
+    "to_db",
+    "FloatVector",
+]
+
+
+def from_db(value: t.Iterable) -> t.Optional["npt.ArrayLike"]:
+    import numpy as np
+
+    # from `pgvector.utils`
+    # could be ndarray if already cast by lower-level driver
+    if value is None or isinstance(value, np.ndarray):
+        return value
+
+    return np.array(value, dtype=np.float32)
+
+
+def to_db(value: t.Any, dim: t.Optional[int] = None) -> t.Optional[t.List]:
+    import numpy as np
+
+    # from `pgvector.utils`
+    if value is None:
+        return value
+
+    if isinstance(value, np.ndarray):
+        if value.ndim != 1:
+            raise ValueError("expected ndim to be 1")
+
+        if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(
+            value.dtype, np.floating
+        ):
+            raise ValueError("dtype must be numeric")
+
+        value = value.tolist()
+
+    if dim is not None and len(value) != dim:
+        raise ValueError("expected %d dimensions, not %d" % (dim, len(value)))
+
+    return value
+
+
+class FloatVector(sa.TypeDecorator):
+    """
+    SQLAlchemy `FloatVector` data type for CrateDB.
+    """
+
+    cache_ok = False
+
+    __visit_name__ = "FLOAT_VECTOR"
+
+    _is_array = True
+
+    zero_indexes = False
+
+    impl = sa.ARRAY
+
+    def __init__(self, dimensions: int = None):
+        super().__init__(sa.FLOAT, dimensions=dimensions)
+
+    def as_generic(self, allow_nulltype=False):
+        return sa.ARRAY(item_type=sa.FLOAT)
+
+    @property
+    def python_type(self):
+        return list
+
+    def bind_processor(self, dialect: sa.engine.Dialect) -> t.Callable:
+        def process(value: t.Iterable) -> t.Optional[t.List]:
+            return to_db(value, self.dimensions)
+
+        return process
+
+    def result_processor(self, dialect: sa.engine.Dialect, coltype: t.Any) -> t.Callable:
+        def process(value: t.Any) -> t.Optional["npt.ArrayLike"]:
+            return from_db(value)
+
+        return process
+
+
+class KnnMatch(ColumnElement):
+    """
+    Wrap CrateDB's `KNN_MATCH` function into an SQLAlchemy function.
+
+    https://cratedb.com/docs/crate/reference/en/latest/general/builtins/scalar-functions.html#scalar-knn-match
+    """
+
+    inherit_cache = True
+
+    def __init__(self, column, term, k=None):
+        super().__init__()
+        self.column = column
+        self.term = term
+        self.k = k
+
+    def compile_column(self, compiler):
+        return compiler.process(self.column)
+
+    def compile_term(self, compiler):
+        return compiler.process(literal(self.term))
+
+    def compile_k(self, compiler):
+        return compiler.process(literal(self.k))
+
+
+def knn_match(column, term, k):
+    """
+    Generate a match predicate for vector search.
+
+    :param column: A reference to a column or an index.
+
+    :param term: The term to match against. This is an array of floating point
+     values, which is compared to other vectors using a HNSW index search.
+
+    :param k: The `k` argument determines the number of nearest neighbours to
+    search in the index.
+    """
+    return KnnMatch(column, term, k)
+
+
+@compiles(KnnMatch)
+def compile_knn_match(knn_match, compiler, **kwargs):
+    """
+    Clause compiler for `KNN_MATCH`.
+    """
+    return "KNN_MATCH(%s, %s, %s)" % (
+        knn_match.compile_column(compiler),
+        knn_match.compile_term(compiler),
+        knn_match.compile_k(compiler),
+    )