diff --git a/CHANGES.md b/CHANGES.md index 94d56e44..475f430e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -2,6 +2,7 @@ ## Unreleased +- Dependencies: Migrate from `crate[sqlalchemy]` to `sqlalchemy-cratedb` ## 2024/05/30 v0.0.12 - Fix InfluxDB Cloud <-> CrateDB Cloud connectivity by using diff --git a/cratedb_toolkit/sqlalchemy/patch.py b/cratedb_toolkit/sqlalchemy/patch.py index 78c89770..ec9e7d44 100644 --- a/cratedb_toolkit/sqlalchemy/patch.py +++ b/cratedb_toolkit/sqlalchemy/patch.py @@ -22,13 +22,16 @@ def get_effective_schema(engine: sa.Engine): schema_name = schema_name_raw[0] return schema_name - from crate.client.sqlalchemy.dialect import CrateDialect + try: + from sqlalchemy_cratedb import dialect + except ImportError: # pragma: nocover + from crate.client.sqlalchemy.dialect import CrateDialect as dialect - get_table_names_dist = CrateDialect.get_table_names + get_table_names_dist = dialect.get_table_names def get_table_names(self, connection: sa.Connection, schema: t.Optional[str] = None, **kw: t.Any) -> t.List[str]: if schema is None: schema = get_effective_schema(connection.engine) return get_table_names_dist(self, connection=connection, schema=schema, **kw) - CrateDialect.get_table_names = get_table_names # type: ignore + dialect.get_table_names = get_table_names # type: ignore diff --git a/cratedb_toolkit/util/common.py b/cratedb_toolkit/util/common.py index 1aad601c..6769c713 100644 --- a/cratedb_toolkit/util/common.py +++ b/cratedb_toolkit/util/common.py @@ -20,6 +20,7 @@ def setup_logging(level=logging.INFO, verbose: bool = False): logging.getLogger("sqlalchemy").setLevel(level) logging.getLogger("crate.client").setLevel(level) + logging.getLogger("sqlalchemy_cratedb").setLevel(level) logging.getLogger("urllib3.connectionpool").setLevel(level) # logging.getLogger("docker.auth").setLevel(logging.INFO) # noqa: ERA001 diff --git a/cratedb_toolkit/util/database.py b/cratedb_toolkit/util/database.py index 1e284009..2667bcc9 100644 --- a/cratedb_toolkit/util/database.py +++ b/cratedb_toolkit/util/database.py @@ -244,7 +244,11 @@ def import_csv_pandas( Import CSV data using pandas. """ import pandas as pd - from crate.client.sqlalchemy.support import insert_bulk + + try: + from sqlalchemy_cratedb.support import insert_bulk + except ImportError: # pragma: nocover + from crate.client.sqlalchemy.support import insert_bulk df = pd.read_csv(filepath) with self.engine.connect() as connection: @@ -267,7 +271,11 @@ def import_csv_dask( """ import dask.dataframe as dd import pandas as pd - from crate.client.sqlalchemy.support import insert_bulk + + try: + from sqlalchemy_cratedb.support import insert_bulk + except ImportError: # pragma: nocover + from crate.client.sqlalchemy.support import insert_bulk # Set a few defaults. npartitions = npartitions or os.cpu_count() diff --git a/pyproject.toml b/pyproject.toml index 7dabd8fd..db1cb12f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,11 +90,11 @@ dependencies = [ "colorama<1", "colorlog", "crash", - "crate[sqlalchemy]>=0.34", + "crate==1.0.0dev0", 'importlib-metadata; python_version <= "3.7"', "python-dotenv<2", "python-slugify<9", - "sqlalchemy", + "sqlalchemy-cratedb", "sqlparse<0.6", 'typing-extensions<5; python_version <= "3.7"', ]