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

sqlalchemy_cratedb/dialect.py

@@ -0,0 +1,414 @@
+# -*- 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.
+
+import logging
+from datetime import date, datetime
+
+from sqlalchemy import types as sqltypes
+from sqlalchemy.engine import default, reflection
+from sqlalchemy.sql import functions
+from sqlalchemy.util import asbool, to_list
+
+from .compiler import (
+    CrateDDLCompiler,
+    CrateIdentifierPreparer,
+    CrateTypeCompiler,
+)
+from .sa_version import SA_1_4, SA_2_0, SA_VERSION
+from .type import FloatVector, ObjectArray, ObjectType
+
+TYPES_MAP = {
+    "boolean": sqltypes.Boolean,
+    "short": sqltypes.SmallInteger,
+    "smallint": sqltypes.SmallInteger,
+    "timestamp": sqltypes.TIMESTAMP(timezone=False),
+    "timestamp with time zone": sqltypes.TIMESTAMP(timezone=True),
+    "object": ObjectType,
+    "integer": sqltypes.Integer,
+    "long": sqltypes.NUMERIC,
+    "bigint": sqltypes.NUMERIC,
+    "double": sqltypes.DECIMAL,
+    "double precision": sqltypes.DECIMAL,
+    "object_array": ObjectArray,
+    "float": sqltypes.Float,
+    "real": sqltypes.Float,
+    "string": sqltypes.String,
+    "text": sqltypes.String,
+    "float_vector": FloatVector,
+}
+
+# Needed for SQLAlchemy >= 1.1.
+# TODO: Dissolve.
+try:
+    from sqlalchemy.types import ARRAY
+
+    TYPES_MAP["integer_array"] = ARRAY(sqltypes.Integer)
+    TYPES_MAP["boolean_array"] = ARRAY(sqltypes.Boolean)
+    TYPES_MAP["short_array"] = ARRAY(sqltypes.SmallInteger)
+    TYPES_MAP["smallint_array"] = ARRAY(sqltypes.SmallInteger)
+    TYPES_MAP["timestamp_array"] = ARRAY(sqltypes.TIMESTAMP(timezone=False))
+    TYPES_MAP["timestamp with time zone_array"] = ARRAY(sqltypes.TIMESTAMP(timezone=True))
+    TYPES_MAP["long_array"] = ARRAY(sqltypes.NUMERIC)
+    TYPES_MAP["bigint_array"] = ARRAY(sqltypes.NUMERIC)
+    TYPES_MAP["double_array"] = ARRAY(sqltypes.DECIMAL)
+    TYPES_MAP["double precision_array"] = ARRAY(sqltypes.DECIMAL)
+    TYPES_MAP["float_array"] = ARRAY(sqltypes.Float)
+    TYPES_MAP["real_array"] = ARRAY(sqltypes.Float)
+    TYPES_MAP["string_array"] = ARRAY(sqltypes.String)
+    TYPES_MAP["text_array"] = ARRAY(sqltypes.String)
+except Exception:  # noqa: S110
+    pass
+
+
+log = logging.getLogger(__name__)
+
+
+class Date(sqltypes.Date):
+    def bind_processor(self, dialect):
+        def process(value):
+            if value is not None:
+                assert isinstance(value, date)  # noqa: S101
+                return value.strftime("%Y-%m-%d")
+            return None
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if not value:
+                return None
+            try:
+                return datetime.utcfromtimestamp(value / 1e3).date()
+            except TypeError:
+                pass
+
+            # Crate doesn't really have datetime or date types but a
+            # timestamp type. The "date" mapping (conversion to long)
+            # is only applied if the schema definition for the column exists
+            # and if the sql insert statement was used.
+            # In case of dynamic mapping or using the rest indexing endpoint
+            # the date will be returned in the format it was inserted.
+            log.warning(
+                "Received timestamp isn't a long value."
+                "Trying to parse as date string and then as datetime string"
+            )
+            try:
+                return datetime.strptime(value, "%Y-%m-%d").date()
+            except ValueError:
+                return datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%fZ").date()
+
+        return process
+
+
+class DateTime(sqltypes.DateTime):
+    def bind_processor(self, dialect):
+        def process(value):
+            if isinstance(value, (datetime, date)):
+                return value.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
+            return value
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if not value:
+                return None
+            try:
+                return datetime.utcfromtimestamp(value / 1e3)
+            except TypeError:
+                pass
+
+            # Crate doesn't really have datetime or date types but a
+            # timestamp type. The "date" mapping (conversion to long)
+            # is only applied if the schema definition for the column exists
+            # and if the sql insert statement was used.
+            # In case of dynamic mapping or using the rest indexing endpoint
+            # the date will be returned in the format it was inserted.
+            log.warning(
+                "Received timestamp isn't a long value."
+                "Trying to parse as datetime string and then as date string"
+            )
+            try:
+                return datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%fZ")
+            except ValueError:
+                return datetime.strptime(value, "%Y-%m-%d")
+
+        return process
+
+
+colspecs = {
+    sqltypes.Date: Date,
+    sqltypes.DateTime: DateTime,
+    sqltypes.TIMESTAMP: DateTime,
+}
+
+
+if SA_VERSION >= SA_2_0:
+    from .compat.core20 import CrateCompilerSA20
+
+    statement_compiler = CrateCompilerSA20
+elif SA_VERSION >= SA_1_4:
+    from .compat.core14 import CrateCompilerSA14
+
+    statement_compiler = CrateCompilerSA14
+else:
+    from .compat.core10 import CrateCompilerSA10
+
+    statement_compiler = CrateCompilerSA10
+
+
+class CrateDialect(default.DefaultDialect):
+    name = "crate"
+    driver = "crate-python"
+    default_paramstyle = "qmark"
+    statement_compiler = statement_compiler
+    ddl_compiler = CrateDDLCompiler
+    type_compiler = CrateTypeCompiler
+    preparer = CrateIdentifierPreparer
+    use_insertmanyvalues = True
+    use_insertmanyvalues_wo_returning = True
+    supports_multivalues_insert = True
+    supports_native_boolean = True
+    supports_statement_cache = True
+    colspecs = colspecs
+    implicit_returning = True
+    insert_returning = True
+    update_returning = True
+
+    def __init__(self, **kwargs):
+        default.DefaultDialect.__init__(self, **kwargs)
+
+        # CrateDB does not need `OBJECT` types to be serialized as JSON.
+        # Corresponding data is forwarded 1:1, and will get marshalled
+        # by the low-level driver.
+        self._json_deserializer = lambda x: x
+        self._json_serializer = lambda x: x
+
+        # Currently, our SQL parser doesn't support unquoted column names that
+        # start with _. Adding it here causes sqlalchemy to quote such columns.
+        self.identifier_preparer.illegal_initial_characters.add("_")
+
+    def initialize(self, connection):
+        # get lowest server version
+        self.server_version_info = self._get_server_version_info(connection)
+        # get default schema name
+        self.default_schema_name = self._get_default_schema_name(connection)
+
+    def do_rollback(self, connection):
+        # if any exception is raised by the dbapi, sqlalchemy by default
+        # attempts to do a rollback crate doesn't support rollbacks.
+        # implementing this as noop seems to cause sqlalchemy to propagate the
+        # original exception to the user
+        pass
+
+    def connect(self, host=None, port=None, *args, **kwargs):
+        server = None
+        if host:
+            server = "{0}:{1}".format(host, port or "4200")
+        if "servers" in kwargs:
+            server = kwargs.pop("servers")
+        servers = to_list(server)
+        if servers:
+            use_ssl = asbool(kwargs.pop("ssl", False))
+            if use_ssl:
+                servers = ["https://" + server for server in servers]
+            return self.dbapi.connect(servers=servers, **kwargs)
+        return self.dbapi.connect(**kwargs)
+
+    def do_execute(self, cursor, statement, parameters, context=None):
+        """
+        Slightly amended to store its response into the request context instance.
+        """
+        result = cursor.execute(statement, parameters)
+        if context is not None:
+            context.last_result = result
+
+    def do_execute_no_params(self, cursor, statement, context=None):
+        """
+        Slightly amended to store its response into the request context instance.
+        """
+        result = cursor.execute(statement)
+        if context is not None:
+            context.last_result = result
+
+    def do_executemany(self, cursor, statement, parameters, context=None):
+        """
+        Slightly amended to store its response into the request context instance.
+        """
+        result = cursor.executemany(statement, parameters)
+        if context is not None:
+            context.last_result = result
+
+    def _get_default_schema_name(self, connection):
+        return "doc"
+
+    def _get_effective_schema_name(self, connection):
+        schema_name_raw = connection.engine.url.query.get("schema")
+        schema_name = None
+        if isinstance(schema_name_raw, str):
+            schema_name = schema_name_raw
+        elif isinstance(schema_name_raw, tuple):
+            schema_name = schema_name_raw[0]
+        return schema_name
+
+    def _get_server_version_info(self, connection):
+        return tuple(connection.connection.lowest_server_version.version)
+
+    @classmethod
+    def import_dbapi(cls):
+        from crate import client
+
+        return client
+
+    @classmethod
+    def dbapi(cls):
+        return cls.import_dbapi()
+
+    def has_schema(self, connection, schema, **kw):
+        return schema in self.get_schema_names(connection, **kw)
+
+    def has_table(self, connection, table_name, schema=None, **kw):
+        return table_name in self.get_table_names(connection, schema=schema, **kw)
+
+    @reflection.cache
+    def get_schema_names(self, connection, **kw):
+        cursor = connection.exec_driver_sql(
+            "select schema_name from information_schema.schemata order by schema_name asc"
+        )
+        return [row[0] for row in cursor.fetchall()]
+
+    @reflection.cache
+    def get_table_names(self, connection, schema=None, **kw):
+        if schema is None:
+            schema = self._get_effective_schema_name(connection)
+        cursor = connection.exec_driver_sql(
+            "SELECT table_name FROM information_schema.tables "
+            "WHERE {0} = ? "
+            "AND table_type = 'BASE TABLE' "
+            "ORDER BY table_name ASC, {0} ASC".format(self.schema_column),
+            (schema or self.default_schema_name,),
+        )
+        return [row[0] for row in cursor.fetchall()]
+
+    @reflection.cache
+    def get_view_names(self, connection, schema=None, **kw):
+        cursor = connection.exec_driver_sql(
+            "SELECT table_name FROM information_schema.views "
+            "ORDER BY table_name ASC, {0} ASC".format(self.schema_column),
+            (schema or self.default_schema_name,),
+        )
+        return [row[0] for row in cursor.fetchall()]
+
+    @reflection.cache
+    def get_columns(self, connection, table_name, schema=None, **kw):
+        query = (
+            "SELECT column_name, data_type "
+            "FROM information_schema.columns "
+            "WHERE table_name = ? AND {0} = ? "
+            "AND column_name !~ ?".format(self.schema_column)
+        )
+        cursor = connection.exec_driver_sql(
+            query,
+            (
+                table_name,
+                schema or self.default_schema_name,
+                r"(.*)\[\'(.*)\'\]",
+            ),  # regex to filter subscript
+        )
+        return [self._create_column_info(row) for row in cursor.fetchall()]
+
+    @reflection.cache
+    def get_pk_constraint(self, engine, table_name, schema=None, **kw):
+        if self.server_version_info >= (3, 0, 0):
+            query = """SELECT column_name
+                    FROM information_schema.key_column_usage
+                    WHERE table_name = ? AND table_schema = ?"""
+
+            def result_fun(result):
+                rows = result.fetchall()
+                return set(map(lambda el: el[0], rows))
+
+        elif self.server_version_info >= (2, 3, 0):
+            query = """SELECT column_name
+                    FROM information_schema.key_column_usage
+                    WHERE table_name = ? AND table_catalog = ?"""
+
+            def result_fun(result):
+                rows = result.fetchall()
+                return set(map(lambda el: el[0], rows))
+
+        else:
+            query = """SELECT constraint_name
+                   FROM information_schema.table_constraints
+                   WHERE table_name = ? AND {schema_col} = ?
+                   AND constraint_type='PRIMARY_KEY'
+                   """.format(schema_col=self.schema_column)
+
+            def result_fun(result):
+                rows = result.fetchone()
+                return set(rows[0] if rows else [])
+
+        pk_result = engine.exec_driver_sql(query, (table_name, schema or self.default_schema_name))
+        pks = result_fun(pk_result)
+        return {"constrained_columns": sorted(pks), "name": "PRIMARY KEY"}
+
+    @reflection.cache
+    def get_foreign_keys(
+        self, connection, table_name, schema=None, postgresql_ignore_search_path=False, **kw
+    ):
+        # Crate doesn't support Foreign Keys, so this stays empty
+        return []
+
+    @reflection.cache
+    def get_indexes(self, connection, table_name, schema, **kw):
+        return []
+
+    @property
+    def schema_column(self):
+        return "table_schema"
+
+    def _create_column_info(self, row):
+        return {
+            "name": row[0],
+            "type": self._resolve_type(row[1]),
+            # In Crate every column is nullable except PK
+            # Primary Key Constraints are not nullable anyway, no matter what
+            # we return here, so it's fine to return always `True`
+            "nullable": True,
+        }
+
+    def _resolve_type(self, type_):
+        return TYPES_MAP.get(type_, sqltypes.UserDefinedType)
+
+    def has_ilike_operator(self):
+        """
+        Only CrateDB 4.1.0 and higher implements the `ILIKE` operator.
+        """
+        server_version_info = self.server_version_info
+        return server_version_info is not None and server_version_info >= (4, 1, 0)
+
+
+class DateTrunc(functions.GenericFunction):
+    name = "date_trunc"
+    type = sqltypes.TIMESTAMP
+
+
+dialect = CrateDialect

sqlalchemy_cratedb/predicate.py

@@ -0,0 +1,96 @@
+# -*- 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.
+
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.expression import ColumnElement, literal
+
+
+class Match(ColumnElement):
+    inherit_cache = True
+
+    def __init__(self, column, term, match_type=None, options=None):
+        super(Match, self).__init__()
+        self.column = column
+        self.term = term
+        self.match_type = match_type
+        self.options = options
+
+    def compile_column(self, compiler):
+        if isinstance(self.column, dict):
+            column = ", ".join(
+                sorted(["{0} {1}".format(compiler.process(k), v) for k, v in self.column.items()])
+            )
+            return "({0})".format(column)
+        else:
+            return "{0}".format(compiler.process(self.column))
+
+    def compile_term(self, compiler):
+        return compiler.process(literal(self.term))
+
+    def compile_using(self, compiler):
+        if self.match_type:
+            using = "using {0}".format(self.match_type)
+            with_clause = self.with_clause()
+            if with_clause:
+                using = " ".join([using, with_clause])
+            return using
+        if self.options:
+            raise ValueError(
+                "missing match_type. "
+                + "It's not allowed to specify options "
+                + "without match_type"
+            )
+        return None
+
+    def with_clause(self):
+        if self.options:
+            options = ", ".join(sorted(["{0}={1}".format(k, v) for k, v in self.options.items()]))
+
+            return "with ({0})".format(options)
+        return None
+
+
+def match(column, term, match_type=None, options=None):
+    """Generates match predicate for fulltext search
+
+    :param column: A reference to a column or an index, or a subcolumn, or a
+     dictionary of subcolumns with boost values.
+
+    :param term: The term to match against. This string is analyzed and the
+     resulting tokens are compared to the index.
+
+    :param match_type (optional): The match type. Determine how the term is
+     applied and the score calculated.
+
+    :param options (optional): The match options. Specify match type behaviour.
+     (Not possible without a specified match type.) Match options must be
+     supplied as a dictionary.
+    """
+    return Match(column, term, match_type, options)
+
+
+@compiles(Match)
+def compile_match(match, compiler, **kwargs):
+    func = "match(%s, %s)" % (match.compile_column(compiler), match.compile_term(compiler))
+    using = match.compile_using(compiler)
+    if using:
+        func = " ".join([func, using])
+    return func

sqlalchemy_cratedb/sa_version.py

@@ -0,0 +1,28 @@
+# -*- 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.
+
+import sqlalchemy as sa
+from verlib2 import Version
+
+SA_VERSION = Version(sa.__version__)
+
+SA_1_4 = Version("1.4.0b1")
+SA_2_0 = Version("2.0.0")

sqlalchemy_cratedb/support/__init__.py

@@ -0,0 +1,18 @@
+from sqlalchemy_cratedb.support.pandas import insert_bulk, table_kwargs
+from sqlalchemy_cratedb.support.polyfill import (
+    check_uniqueness_factory,
+    patch_autoincrement_timestamp,
+    refresh_after_dml,
+)
+from sqlalchemy_cratedb.support.util import quote_relation_name, refresh_dirty, refresh_table
+
+__all__ = [
+    check_uniqueness_factory,
+    insert_bulk,
+    patch_autoincrement_timestamp,
+    quote_relation_name,
+    refresh_after_dml,
+    refresh_dirty,
+    refresh_table,
+    table_kwargs,
+]

sqlalchemy_cratedb/support/pandas.py

@@ -0,0 +1,110 @@
+# -*- 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.
+import logging
+from contextlib import contextmanager
+from typing import Any
+from unittest.mock import patch
+
+import sqlalchemy as sa
+
+from sqlalchemy_cratedb.sa_version import SA_2_0, SA_VERSION
+
+logger = logging.getLogger(__name__)
+
+
+def insert_bulk(pd_table, conn, keys, data_iter):
+    """
+    Use CrateDB's "bulk operations" endpoint as a fast path for pandas' and Dask's `to_sql()` [1] method.
+
+    The idea is to break out of SQLAlchemy, compile the insert statement, and use the raw
+    DBAPI connection client, in order to invoke a request using `bulk_parameters` [2]::
+
+        cursor.execute(sql=sql, bulk_parameters=data)
+
+    The vanilla implementation, used by SQLAlchemy, is::
+
+        data = [dict(zip(keys, row)) for row in data_iter]
+        conn.execute(pd_table.table.insert(), data)
+
+    Batch chunking will happen outside of this function, for example [3] demonstrates
+    the relevant code in `pandas.io.sql`.
+
+    [1] https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.to_sql.html
+    [2] https://cratedb.com/docs/crate/reference/en/latest/interfaces/http.html#bulk-operations
+    [3] https://github.com/pandas-dev/pandas/blob/v2.0.1/pandas/io/sql.py#L1011-L1027
+    """  # noqa: E501
+
+    # Compile SQL statement and materialize batch.
+    sql = str(pd_table.table.insert().compile(bind=conn))
+    data = list(data_iter)
+
+    # For debugging and tracing the batches running through this method.
+    if logger.level == logging.DEBUG:
+        logger.debug(f"Bulk SQL:     {sql}")
+        logger.debug(f"Bulk records: {len(data)}")
+        # logger.debug(f"Bulk data:    {data}")  # noqa: ERA001
+
+    # Invoke bulk insert operation.
+    cursor = conn._dbapi_connection.cursor()
+    cursor.execute(sql=sql, bulk_parameters=data)
+    cursor.close()
+
+
+@contextmanager
+def table_kwargs(**kwargs):
+    """
+    Context manager for adding SQLAlchemy dialect-specific table options at runtime.
+
+    In certain cases where SQLAlchemy orchestration is implemented within a
+    framework, like at this spot [1] in pandas' `SQLTable._create_table_setup`,
+    it is not easily possible to forward SQLAlchemy dialect options at table
+    creation time.
+
+    In order to augment the SQL DDL statement to make it honor database-specific
+    dialect options, the only way to work around the unfortunate situation is by
+    monkey-patching the call to `sa.Table()` at runtime, relaying additional
+    dialect options through corresponding keyword arguments in their original
+    `<dialect>_<kwarg>` format [2].
+
+    [1] https://github.com/pandas-dev/pandas/blob/v2.2.2/pandas/io/sql.py#L1282-L1285
+    [2] https://docs.sqlalchemy.org/en/20/core/foundation.html#sqlalchemy.sql.base.DialectKWArgs.dialect_kwargs
+    """
+
+    if SA_VERSION < SA_2_0:
+        _init_dist = sa.sql.schema.Table._init
+
+        def _init(self, name, metadata, *args, **kwargs_effective):
+            kwargs_effective.update(kwargs)
+            return _init_dist(self, name, metadata, *args, **kwargs_effective)
+
+        with patch("sqlalchemy.sql.schema.Table._init", _init):
+            yield
+
+    else:
+        new_dist = sa.sql.schema.Table._new
+
+        def _new(cls, *args: Any, **kw: Any) -> Any:
+            kw.update(kwargs)
+            table = new_dist(cls, *args, **kw)
+            return table
+
+        with patch("sqlalchemy.sql.schema.Table._new", _new):
+            yield