From f8cdcacb34ab2789df8600b8c69ecf5bf8af7fab Mon Sep 17 00:00:00 2001 From: Andreas Motl Date: Wed, 25 Mar 2026 04:42:22 +0100 Subject: [PATCH 1/2] QA: Migrate type checker from `mypy` to `ty` -- Human edits --- cratedb_toolkit/__init__.py | 2 +- cratedb_toolkit/adapter/pymongo/api.py | 2 +- cratedb_toolkit/adapter/pymongo/collection.py | 8 ++-- cratedb_toolkit/adapter/pymongo/cursor.py | 5 +-- cratedb_toolkit/adapter/pymongo/reactor.py | 13 ++++--- cratedb_toolkit/adapter/pymongo/util.py | 2 +- cratedb_toolkit/adapter/rockset/cli.py | 2 +- .../adapter/rockset/server/api/query.py | 2 +- cratedb_toolkit/cfr/cli.py | 4 +- cratedb_toolkit/cfr/marimo.py | 1 - cratedb_toolkit/cfr/systable.py | 8 ++-- cratedb_toolkit/cli.py | 2 +- cratedb_toolkit/cluster/cli.py | 2 +- cratedb_toolkit/cluster/core.py | 24 +++++++----- cratedb_toolkit/cluster/croud.py | 8 ++-- cratedb_toolkit/cmd/tail/main.py | 8 ++-- cratedb_toolkit/datasets/model.py | 2 +- cratedb_toolkit/datasets/util.py | 3 +- cratedb_toolkit/dms/cli.py | 2 +- cratedb_toolkit/dms/table_mapping.py | 4 +- cratedb_toolkit/docs/functions.py | 8 ++-- cratedb_toolkit/exception.py | 6 ++- cratedb_toolkit/info/http.py | 2 +- cratedb_toolkit/io/awslambda/kinesis.py | 4 +- cratedb_toolkit/io/cratedb/bulk.py | 13 +++++-- cratedb_toolkit/io/kinesis/adapter.py | 8 +++- cratedb_toolkit/io/kinesis/relay.py | 11 ++++-- cratedb_toolkit/io/mongodb/adapter.py | 4 ++ cratedb_toolkit/io/mongodb/cdc.py | 3 +- cratedb_toolkit/io/mongodb/copy.py | 5 ++- cratedb_toolkit/io/mongodb/core.py | 2 +- cratedb_toolkit/io/mongodb/export.py | 8 +++- cratedb_toolkit/io/mongodb/transform.py | 9 +++-- cratedb_toolkit/io/mongodb/translate.py | 3 +- cratedb_toolkit/model.py | 2 +- cratedb_toolkit/query/cli.py | 2 +- cratedb_toolkit/query/mcp/cli.py | 2 +- cratedb_toolkit/query/mcp/pg_mcp.py | 2 +- cratedb_toolkit/retention/cli.py | 2 +- cratedb_toolkit/retention/store.py | 12 +++--- .../testing/testcontainers/azurite.py | 2 +- .../testing/testcontainers/cratedb.py | 9 +++-- .../testing/testcontainers/influxdb2.py | 4 +- .../testing/testcontainers/localstack.py | 2 + .../testing/testcontainers/mongodb.py | 5 +++ .../testing/testcontainers/util.py | 17 ++++----- cratedb_toolkit/util/app.py | 2 +- cratedb_toolkit/util/cli.py | 4 +- cratedb_toolkit/util/client.py | 13 ++++--- cratedb_toolkit/util/crash.py | 10 ++--- cratedb_toolkit/util/croud.py | 6 +-- cratedb_toolkit/util/database.py | 18 +++++---- cratedb_toolkit/util/pandas.py | 9 +++-- cratedb_toolkit/util/platform.py | 2 +- cratedb_toolkit/util/runtime.py | 2 +- cratedb_toolkit/util/setting.py | 7 +++- cratedb_toolkit/util/sqlalchemy.py | 2 +- pyproject.toml | 37 +++++++++++++------ tests/adapter/test_pymongo.py | 2 +- tests/cfr/test_systable.py | 4 +- tests/cluster/test_core.py | 1 + tests/conftest.py | 15 ++++---- tests/datasets/test_loading.py | 2 +- tests/examples/test_python.py | 2 +- tests/examples/test_shell.py | 2 +- tests/io/dynamodb/conftest.py | 11 ++++-- tests/io/dynamodb/test_adapter.py | 2 +- tests/io/dynamodb/test_relay.py | 1 + tests/io/influxdb/conftest.py | 6 ++- tests/io/kinesis/conftest.py | 4 +- tests/io/kinesis/manager.py | 4 +- tests/io/kinesis/test_relay.py | 5 ++- tests/io/mongodb/conftest.py | 13 ++++--- tests/io/mongodb/test_integration.py | 3 +- tests/io/test_iceberg.py | 2 +- tests/shell/test_cli.py | 4 +- tests/util/test_run_sql.py | 4 +- 77 files changed, 268 insertions(+), 182 deletions(-) diff --git a/cratedb_toolkit/__init__.py b/cratedb_toolkit/__init__.py index e1e1ea57..db9662ae 100644 --- a/cratedb_toolkit/__init__.py +++ b/cratedb_toolkit/__init__.py @@ -2,7 +2,7 @@ try: from importlib.metadata import PackageNotFoundError, version except (ImportError, ModuleNotFoundError): # pragma:nocover - from importlib_metadata import PackageNotFoundError, version # type: ignore[assignment,no-redef,unused-ignore] + from importlib_metadata import PackageNotFoundError, version __appname__ = "cratedb-toolkit" diff --git a/cratedb_toolkit/adapter/pymongo/api.py b/cratedb_toolkit/adapter/pymongo/api.py index ab17726d..4a74f70d 100644 --- a/cratedb_toolkit/adapter/pymongo/api.py +++ b/cratedb_toolkit/adapter/pymongo/api.py @@ -17,7 +17,7 @@ def __init__(self, dburi: str): self.cratedb = DatabaseAdapter(dburi=dburi) self.collection_backup = pymongo.collection.Collection - collection_patched = collection_factory(cratedb=self.cratedb) # type: ignore[misc] + collection_patched = collection_factory(cratedb=self.cratedb) self.patches = [ # Patch PyMongo's `Collection` implementation. patch("pymongo.collection.Collection", collection_patched), diff --git a/cratedb_toolkit/adapter/pymongo/collection.py b/cratedb_toolkit/adapter/pymongo/collection.py index 9257ea1b..a6c4bf2d 100644 --- a/cratedb_toolkit/adapter/pymongo/collection.py +++ b/cratedb_toolkit/adapter/pymongo/collection.py @@ -52,7 +52,7 @@ def get_df_info(df: pd.DataFrame) -> str: def insert_one( self: Collection, document: Union[_DocumentType, RawBSONDocument], - bypass_document_validation: bool = False, + bypass_document_validation: Optional[bool] = None, session: Optional[ClientSession] = None, comment: Optional[Any] = None, ) -> InsertOneResult: @@ -112,7 +112,7 @@ def insert_many( self, documents: Iterable[Union[_DocumentType, RawBSONDocument]], ordered: bool = True, - bypass_document_validation: bool = False, + bypass_document_validation: Optional[bool] = None, session: Optional[ClientSession] = None, comment: Optional[Any] = None, ) -> InsertManyResult: @@ -127,8 +127,8 @@ def gen() -> Iterator[Mapping[str, Any]]: if "_id" in document: identifier = ObjectId(document["_id"]) if isinstance(document, RawBSONDocument): - document = document.decode() - del document["_id"] + document = document.decode() # ty: ignore[unresolved-attribute] + del document["_id"] # ty: ignore[not-subscriptable] else: identifier = ObjectId() inserted_ids.append(identifier) diff --git a/cratedb_toolkit/adapter/pymongo/cursor.py b/cratedb_toolkit/adapter/pymongo/cursor.py index 04c5825a..3118d084 100644 --- a/cratedb_toolkit/adapter/pymongo/cursor.py +++ b/cratedb_toolkit/adapter/pymongo/cursor.py @@ -1,5 +1,4 @@ # Compansate pymongo<>4.9 woes. -# mypy: disable-error-code="arg-type,attr-defined,call-arg,misc" # Make Python 3.7 and 3.8 support generic types like `dict` instead of `typing.Dict`. from __future__ import annotations @@ -250,7 +249,7 @@ def _refresh(self) -> int: self.__query_spec(), self.__projection, self.__codec_options, - self._read_preference(), + self._read_preference(), # ty: ignore[call-non-callable,invalid-argument-type,missing-argument] self.__limit, self.__batch_size, self.__read_concern, @@ -275,7 +274,7 @@ def _refresh(self) -> int: limit, self.__id, self.__codec_options, - self._read_preference(), + self._read_preference(), # ty: ignore[call-non-callable,invalid-argument-type,missing-argument] self.__session, self.__collection.database.client, self.__max_await_time_ms, diff --git a/cratedb_toolkit/adapter/pymongo/reactor.py b/cratedb_toolkit/adapter/pymongo/reactor.py index 0a8a6cc6..863bfe59 100644 --- a/cratedb_toolkit/adapter/pymongo/reactor.py +++ b/cratedb_toolkit/adapter/pymongo/reactor.py @@ -5,22 +5,23 @@ from jessiql import Query, QueryObject, QueryObjectDict from jessiql.exc import InvalidColumnError from jessiql.typing import SARowDict +from sqlalchemy.orm import Mapper, registry -def table_to_model(table: sa.Table) -> t.Type[sa.orm.Mapper]: +def table_to_model(table: sa.Table) -> t.Type[Mapper]: """ Create SQLAlchemy model class from Table object. - https://docs.sqlalchemy.org/en/14/orm/mapping_styles.html#imperative-mapping - https://sparrigan.github.io/sql/sqla/2016/01/03/dynamic-tables.html """ - mapper_registry = sa.orm.registry(metadata=table.metadata) + mapper_registry = registry(metadata=table.metadata) Surrogate = type("Surrogate", (), {}) mapper_registry.map_imperatively(Surrogate, table) - return Surrogate + return Surrogate # ty: ignore[invalid-return-type] -def reflect_model(engine: t.Any, metadata: sa.MetaData, table_name: str) -> t.Type[sa.orm.Mapper]: +def reflect_model(engine: t.Any, metadata: sa.MetaData, table_name: str) -> t.Type[Mapper]: """ Create SQLAlchemy model class by reflecting a database table. """ @@ -29,7 +30,7 @@ def reflect_model(engine: t.Any, metadata: sa.MetaData, table_name: str) -> t.Ty def mongodb_query( - model: t.Type[sa.orm.Mapper], + model: t.Type[Mapper], select: t.Union[t.List, None] = None, filter: t.Union[t.Dict[str, t.Any], None] = None, # noqa: A002 sort: t.Union[t.List[str], None] = None, @@ -38,7 +39,7 @@ def mongodb_query( Create a JessiQL Query object from an SQLAlchemy model class and typical MongoDB query parameters. """ - select = select or list(model._sa_class_manager.keys()) # type: ignore[attr-defined] + select = select or list(model._sa_class_manager.keys()) filter = filter or {} # noqa: A001 sort = sort or [] diff --git a/cratedb_toolkit/adapter/pymongo/util.py b/cratedb_toolkit/adapter/pymongo/util.py index 171068e5..6ec88892 100644 --- a/cratedb_toolkit/adapter/pymongo/util.py +++ b/cratedb_toolkit/adapter/pymongo/util.py @@ -26,7 +26,7 @@ def from_str(cls, oid: str): return cls(bytes(oid, "ascii")) def __str__(self) -> str: - return binascii.hexlify(self.__id).decode() # type: ignore[arg-type] + return binascii.hexlify(self.__id).decode() # ty: ignore[invalid-argument-type] def __repr__(self) -> str: return f"ObjectId('{self!s}')" diff --git a/cratedb_toolkit/adapter/rockset/cli.py b/cratedb_toolkit/adapter/rockset/cli.py index 7669381f..a4589509 100644 --- a/cratedb_toolkit/adapter/rockset/cli.py +++ b/cratedb_toolkit/adapter/rockset/cli.py @@ -25,7 +25,7 @@ def help_serve(): """ # noqa: E501 -@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup) @option_cluster_url @click.option("--verbose", is_flag=True, required=False, help="Turn on logging") @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level") diff --git a/cratedb_toolkit/adapter/rockset/server/api/query.py b/cratedb_toolkit/adapter/rockset/server/api/query.py index c2e5d84e..ba6cafb6 100644 --- a/cratedb_toolkit/adapter/rockset/server/api/query.py +++ b/cratedb_toolkit/adapter/rockset/server/api/query.py @@ -52,7 +52,7 @@ async def execute(request: Request, adapter: t.Annotated[DatabaseAdapter, Depend results = adapter.run_sql( sql=sql, - parameters=parameters, + parameters=parameters or {}, records=True, ) time_duration = time.time_ns() - time_start diff --git a/cratedb_toolkit/cfr/cli.py b/cratedb_toolkit/cfr/cli.py index d8bceddc..953d871b 100644 --- a/cratedb_toolkit/cfr/cli.py +++ b/cratedb_toolkit/cfr/cli.py @@ -85,7 +85,7 @@ def help_statistics(): """ # noqa: E501 -@click.group(cls=ClickAliasedGroup, help=docstring_format_verbatim(help_statistics.__doc__)) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup, help=docstring_format_verbatim(help_statistics.__doc__)) @click.pass_context def job_statistics(ctx: click.Context): """ @@ -221,7 +221,7 @@ def job_statistics_ui(ctx: click.Context): uvicorn.run(app, host="localhost", port=7777, log_level="info") -@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup) @click.pass_context def info(ctx: click.Context): """ diff --git a/cratedb_toolkit/cfr/marimo.py b/cratedb_toolkit/cfr/marimo.py index e196a399..1cee2d87 100644 --- a/cratedb_toolkit/cfr/marimo.py +++ b/cratedb_toolkit/cfr/marimo.py @@ -1,5 +1,4 @@ # ruff: noqa: B018, ERA001, S608, T201 -# mypy: disable-error-code="arg-type" # Source: # https://docs.marimo.io/guides/working_with_data/dataframes.html diff --git a/cratedb_toolkit/cfr/systable.py b/cratedb_toolkit/cfr/systable.py index 7b4743c7..c6ff079f 100644 --- a/cratedb_toolkit/cfr/systable.py +++ b/cratedb_toolkit/cfr/systable.py @@ -83,7 +83,9 @@ def __init__(self, dburi: str): def table_names(self): return self.inspector.get_table_names(schema=SystemTableKnowledge.SYS_SCHEMA) - def ddl(self, tablename_in: str, tablename_out: str, out_schema: str = None, with_drop_table: bool = False) -> str: + def ddl( + self, tablename_in: str, tablename_out: str, out_schema: t.Optional[str] = None, with_drop_table: bool = False + ) -> str: meta = sa.MetaData(schema=SystemTableKnowledge.SYS_SCHEMA) table = sa.Table(tablename_in, meta, autoload_with=self.engine) table.schema = out_schema @@ -147,9 +149,9 @@ def dump_table(self, frame: "pl.DataFrame", file: t.Union[t.TextIO, None] = None # return df.write_csv() # noqa: ERA001 return frame.to_pandas().to_csv(file) elif self.data_format in ["jsonl", "ndjson"]: - return frame.write_ndjson(file and file.buffer) # type: ignore[arg-type] + return frame.write_ndjson(file and file.buffer) elif self.data_format in ["parquet", "pq"]: - return frame.write_parquet(file and file.buffer) # type: ignore[arg-type] + return frame.write_parquet(file and file.buffer) # ty: ignore[invalid-argument-type] else: raise NotImplementedError(f"Output format not implemented: {self.data_format}") diff --git a/cratedb_toolkit/cli.py b/cratedb_toolkit/cli.py index af273a1c..9c750d5b 100644 --- a/cratedb_toolkit/cli.py +++ b/cratedb_toolkit/cli.py @@ -18,7 +18,7 @@ from .util.setting import init_dotenv -@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup) @click.option("--verbose", is_flag=True, required=False, help="Turn on logging") @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level") @click.version_option() diff --git a/cratedb_toolkit/cluster/cli.py b/cratedb_toolkit/cluster/cli.py index 202c603f..fb661d1b 100644 --- a/cratedb_toolkit/cluster/cli.py +++ b/cratedb_toolkit/cluster/cli.py @@ -57,7 +57,7 @@ def help_list_jobs(): """ -@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup) @click.option("--verbose", is_flag=True, required=False, help="Turn on logging") @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level") @click.version_option() diff --git a/cratedb_toolkit/cluster/core.py b/cratedb_toolkit/cluster/core.py index 489a9b2a..f89e77e3 100644 --- a/cratedb_toolkit/cluster/core.py +++ b/cratedb_toolkit/cluster/core.py @@ -9,6 +9,7 @@ from pathlib import Path import click +from crate.client.connection import Connection from cratedb_toolkit.cluster.croud import CloudClusterServices, CloudRootServices from cratedb_toolkit.cluster.guide import DataImportGuide @@ -110,11 +111,11 @@ class ManagedCluster(ClusterBase): def __init__( self, - cluster_id: str = None, - cluster_name: str = None, - settings: ManagedClusterSettings = None, - address: DatabaseAddress = None, - info: ClusterInformation = None, + cluster_id: t.Optional[str] = None, + cluster_name: t.Optional[str] = None, + settings: t.Optional[ManagedClusterSettings] = None, + address: t.Optional[DatabaseAddress] = None, + info: t.Optional[ClusterInformation] = None, stop_on_exit: bool = False, ): super().__init__() @@ -408,7 +409,7 @@ def adapter(self) -> DatabaseAdapter: raise DatabaseAddressMissingError() return DatabaseAdapter(dburi=self.address.dburi, jwt=self.info.jwt) - def get_client_bundle(self, username: str = None, password: str = None) -> ClientBundle: + def get_client_bundle(self, username: t.Optional[str] = None, password: t.Optional[str] = None) -> ClientBundle: """ Return a bundle of client handles to the CrateDB Cloud cluster database. @@ -431,7 +432,7 @@ def get_client_bundle(self, username: str = None, password: str = None) -> Clien adapter = DatabaseAdapter(address.dburi) self._client_bundle = ClientBundle( adapter=adapter, - dbapi=adapter.connection.connection.dbapi_connection, + dbapi=t.cast(Connection, adapter.connection.connection.dbapi_connection), sqlalchemy=adapter.engine, ) return self._client_bundle @@ -492,7 +493,7 @@ def adapter(self) -> DatabaseAdapter: """ return DatabaseAdapter(dburi=self.address.dburi) - def get_client_bundle(self, username: str = None, password: str = None) -> ClientBundle: + def get_client_bundle(self, username: t.Optional[str] = None, password: t.Optional[str] = None) -> ClientBundle: """ Return a bundle of client handles to the CrateDB Cloud cluster database. @@ -512,7 +513,7 @@ def get_client_bundle(self, username: str = None, password: str = None) -> Clien adapter = DatabaseAdapter(address.dburi) self._client_bundle = ClientBundle( adapter=adapter, - dbapi=adapter.connection.connection.dbapi_connection, + dbapi=t.cast(Connection, adapter.connection.connection.dbapi_connection), sqlalchemy=adapter.engine, ) return self._client_bundle @@ -592,7 +593,10 @@ def from_options(cls, options: ClusterAddressOptions) -> t.Union[ManagedCluster, @classmethod def create( - cls, cluster_id: str = None, cluster_name: str = None, cluster_url: str = None + cls, + cluster_id: t.Optional[str] = None, + cluster_name: t.Optional[str] = None, + cluster_url: t.Optional[str] = None, ) -> t.Union[ManagedCluster, StandaloneCluster]: """ Create the cluster instance based on the provided parameters. diff --git a/cratedb_toolkit/cluster/croud.py b/cratedb_toolkit/cluster/croud.py index eb57a9c8..dead4b8e 100644 --- a/cratedb_toolkit/cluster/croud.py +++ b/cratedb_toolkit/cluster/croud.py @@ -164,7 +164,7 @@ def list_projects(self): item["backup_location"] = _transform_backup_location(item["backup_location"]) return data - def create_project(self, name: str, organization_id: str = None): + def create_project(self, name: str, organization_id: t.Optional[str] = None): """ Create project. @@ -184,7 +184,7 @@ def create_project(self, name: str, organization_id: str = None): call = CroudCall( fun=project_create, - specs=command_tree["projects"]["commands"]["create"]["extra_args"], + specs=command_tree["projects"]["commands"]["create"]["extra_args"], # ty: ignore[invalid-argument-type,not-subscriptable] arguments=arguments, ) @@ -216,7 +216,7 @@ def get_or_create_project(self, name: str) -> str: return project_id - def deploy_cluster(self, name: str, project_id: str, subscription_id: str = None): + def deploy_cluster(self, name: str, project_id: str, subscription_id: t.Optional[str] = None): """ Deploy cluster. @@ -273,7 +273,7 @@ def deploy_cluster(self, name: str, project_id: str, subscription_id: str = None call = CroudCall( fun=clusters_deploy, - specs=command_tree["clusters"]["commands"]["deploy"]["extra_args"], + specs=command_tree["clusters"]["commands"]["deploy"]["extra_args"], # ty: ignore[invalid-argument-type,not-subscriptable] arguments=[ "--subscription-id", subscription_id, diff --git a/cratedb_toolkit/cmd/tail/main.py b/cratedb_toolkit/cmd/tail/main.py index d85e9311..cca77f72 100644 --- a/cratedb_toolkit/cmd/tail/main.py +++ b/cratedb_toolkit/cmd/tail/main.py @@ -5,10 +5,10 @@ import typing as t import attr -import colorlog import orjson import sqlparse import yaml +from colorlog.escape_codes import escape_codes from cratedb_toolkit.model import TableAddress from cratedb_toolkit.util.database import DatabaseAdapter @@ -37,9 +37,9 @@ def template(self) -> str: @property def label(self): - red = colorlog.escape_codes.escape_codes["red"] - green = colorlog.escape_codes.escape_codes["green"] - reset = colorlog.escape_codes.escape_codes["reset"] + red = escape_codes["red"] + green = escape_codes["green"] + reset = escape_codes["reset"] if self.error: return f"{red}ERROR{reset}" else: diff --git a/cratedb_toolkit/datasets/model.py b/cratedb_toolkit/datasets/model.py index 0ff5f936..408368fb 100644 --- a/cratedb_toolkit/datasets/model.py +++ b/cratedb_toolkit/datasets/model.py @@ -9,7 +9,7 @@ try: from typing import Literal except ImportError: - from typing_extensions import Literal # type: ignore[assignment] + from typing_extensions import Literal @dataclasses.dataclass diff --git a/cratedb_toolkit/datasets/util.py b/cratedb_toolkit/datasets/util.py index 1dae50d2..8c227862 100644 --- a/cratedb_toolkit/datasets/util.py +++ b/cratedb_toolkit/datasets/util.py @@ -1,9 +1,10 @@ import os +import typing as t import zipfile from unittest.mock import patch -def load_dataset_kaggle(dataset: str, path: str, file_name: str = None): +def load_dataset_kaggle(dataset: str, path: str, file_name: t.Optional[str] = None): """ Download complete dataset or individual files from Kaggle. diff --git a/cratedb_toolkit/dms/cli.py b/cratedb_toolkit/dms/cli.py index 9a147667..108b8bae 100644 --- a/cratedb_toolkit/dms/cli.py +++ b/cratedb_toolkit/dms/cli.py @@ -17,7 +17,7 @@ def help_table_mappings(): """ -@click.group(cls=ClickAliasedGroup, help="AWS DMS utilities") # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup, help="AWS DMS utilities") @click.pass_context def cli(ctx: click.Context): pass diff --git a/cratedb_toolkit/dms/table_mapping.py b/cratedb_toolkit/dms/table_mapping.py index 95969770..7dc4d17e 100644 --- a/cratedb_toolkit/dms/table_mapping.py +++ b/cratedb_toolkit/dms/table_mapping.py @@ -42,8 +42,8 @@ def add_rule( action: str, locator: Dict[str, str], id: Union[str, None] = None, # noqa: A002 - filters: List[Any] = None, - mapping_parameters: Dict[str, str] = None, + filters: Union[List[Any], None] = None, + mapping_parameters: Union[Dict[str, str], None] = None, ) -> "TableMappingBuilder": if id is None: id = str(len(self.rules) + 1) # noqa: A001 diff --git a/cratedb_toolkit/docs/functions.py b/cratedb_toolkit/docs/functions.py index 79e311d0..879a364c 100644 --- a/cratedb_toolkit/docs/functions.py +++ b/cratedb_toolkit/docs/functions.py @@ -8,7 +8,7 @@ from docutils.examples import internals from docutils.parsers.rst.directives import register_directive from docutils.parsers.rst.directives.admonitions import Note -from docutils.parsers.rst.roles import normalized_role_options, register_canonical_role # type: ignore[attr-defined] +from docutils.parsers.rst.roles import normalized_role_options, register_canonical_role from cratedb_toolkit.docs.model import DocsItem from cratedb_toolkit.docs.util import GenericProcessor @@ -73,8 +73,10 @@ def to_dict(self) -> Dict[str, Any]: def sphinx_ref_role(role, rawtext, text=None, lineno=None, inliner=None, options=None, content=None): + if inliner is None: + raise ValueError("inliner must be provided") options = normalized_role_options(options) - text = nodes.unescape(text, True) # type: ignore[attr-defined] + text = nodes.unescape(text, True) # ty: ignore[unresolved-attribute] label = text.split(" ", 1)[0] node = nodes.raw(rawtext, label, **options) node.source, node.line = inliner.reporter.get_source_and_line(lineno) @@ -114,7 +116,7 @@ def acquire(self): for item in document: if item.tagname == "section": category_title = item.children[0].astext() - for function in item.children: # type: ignore[assignment] + for function in item.children: if function.tagname == "section": function_title = function.children[0].astext() function_body = function.children[1].astext() diff --git a/cratedb_toolkit/exception.py b/cratedb_toolkit/exception.py index 7eb7bf0b..c4baa120 100644 --- a/cratedb_toolkit/exception.py +++ b/cratedb_toolkit/exception.py @@ -1,3 +1,5 @@ +from typing import Optional + from click import ClickException @@ -26,7 +28,7 @@ class DatabaseAddressMissingError(ClickException): "environment variables." ) - def __init__(self, message: str = None): + def __init__(self, message: Optional[str] = None): if not message: message = self.EXTENDED_MESSAGE super().__init__(message) @@ -37,7 +39,7 @@ class DatabaseAddressDuplicateError(ClickException): "Duplicate database address, please specify only one of: cluster id, cluster name, or database URL" ) - def __init__(self, message: str = None): + def __init__(self, message: Optional[str] = None): if not message: message = self.STANDARD_MESSAGE super().__init__(message) diff --git a/cratedb_toolkit/info/http.py b/cratedb_toolkit/info/http.py index c887da52..29a6affc 100644 --- a/cratedb_toolkit/info/http.py +++ b/cratedb_toolkit/info/http.py @@ -30,7 +30,7 @@ def read_root(): @app.get("/info/{category}") -def info(category: str, adapter: t.Annotated[DatabaseAdapter, Depends(database_adapter)], scrub: bool = False): # type: ignore[name-defined] +def info(category: str, adapter: t.Annotated[DatabaseAdapter, Depends(database_adapter)], scrub: bool = False): if category != "all": raise HTTPException(status_code=404, detail="Info category not found") sample = InfoContainer(adapter=adapter, scrub=scrub) diff --git a/cratedb_toolkit/io/awslambda/kinesis.py b/cratedb_toolkit/io/awslambda/kinesis.py index 26e28337..f48286ec 100644 --- a/cratedb_toolkit/io/awslambda/kinesis.py +++ b/cratedb_toolkit/io/awslambda/kinesis.py @@ -80,9 +80,9 @@ # TODO: Propagate mapping definitions and other settings. cdc: t.Union[DMSTranslatorCrateDB, DynamoDBCDCTranslator] if MESSAGE_FORMAT == "dms": - cdc = DMSTranslatorCrateDB(column_types=column_types) + cdc = DMSTranslatorCrateDB(column_types=column_types) # ty: ignore[invalid-argument-type] elif MESSAGE_FORMAT == "dynamodb": - cdc = DynamoDBCDCTranslator(table_name=CRATEDB_TABLE) + cdc = DynamoDBCDCTranslator(table_name=CRATEDB_TABLE) # ty: ignore[invalid-argument-type] # Create the database connection outside the handler to allow # connections to be re-used by subsequent function invocations. diff --git a/cratedb_toolkit/io/cratedb/bulk.py b/cratedb_toolkit/io/cratedb/bulk.py index ee5e0ac1..91e3490b 100644 --- a/cratedb_toolkit/io/cratedb/bulk.py +++ b/cratedb_toolkit/io/cratedb/bulk.py @@ -17,6 +17,8 @@ from sqlalchemy.exc import ProgrammingError from tqdm import tqdm +from cratedb_toolkit.util.cli import to_list + logger = logging.getLogger(__name__) @@ -39,8 +41,8 @@ class BulkResponse: TODO: Think about refactoring this to `sqlalchemy_cratedb.support`. """ - parameters: t.Union[t.List[t.Dict[str, t.Any]], None] - cratedb_bulk_result: t.Union[t.List[BulkResultItem], None] + parameters: t.Union[t.List[t.Dict[str, t.Any]], None] = None + cratedb_bulk_result: t.Union[t.List[BulkResultItem], None] = None @cached_property def failed_records(self) -> t.List[t.Dict[str, t.Any]]: @@ -126,6 +128,9 @@ def log_level(self): return logger.warning def start(self) -> BulkMetrics: + # Sanity checks. + if not self.batch_to_operation: + raise ValueError("Callback `batch_to_operation` not defined") # Acquire batches of documents, convert to SQL operations, and submit to CrateDB. batch_count = 0 for batch in self.data: @@ -151,7 +156,7 @@ def start(self) -> BulkMetrics: self.connection.commit() if cursor.rowcount > 0: cratedb_bulk_result = getattr(cursor.context, "last_result", None) - bulk_response = BulkResponse(operation.parameters, cratedb_bulk_result) + bulk_response = BulkResponse(to_list(operation.parameters, []), cratedb_bulk_result) failed_records = bulk_response.failed_records count_success_local = bulk_response.success_count self._metrics.count_success_total += bulk_response.success_count @@ -176,7 +181,7 @@ def start(self) -> BulkMetrics: ) for record in failed_records: try: - cursor = self.connection.execute(statement=statement, parameters=record) + cursor = self.connection.execute(statement=statement, parameters=record) # ty: ignore[no-matching-overload] self.connection.commit() if cursor.rowcount != 1: raise IOError("Record has not been processed") diff --git a/cratedb_toolkit/io/kinesis/adapter.py b/cratedb_toolkit/io/kinesis/adapter.py index 2df5a165..f9435c86 100644 --- a/cratedb_toolkit/io/kinesis/adapter.py +++ b/cratedb_toolkit/io/kinesis/adapter.py @@ -184,7 +184,7 @@ async def _produce(self, data: t.Dict[str, t.Any]): create_stream_shards=self.create_shards, describe_timeout=self.describe_timeout, ) as producer: - await producer.put(data) + await producer.put(data) # ty: ignore[unresolved-attribute] class KinesisFileAdapter(KinesisAdapterBase): @@ -209,3 +209,9 @@ def consume(self, handler: t.Callable): def stop(self): pass + + def wait_until_ready(self, timeout: float = 30) -> bool: + return True + + def produce(self, data: t.Dict[str, t.Any]): + raise NotImplementedError("Unable to produce to file") diff --git a/cratedb_toolkit/io/kinesis/relay.py b/cratedb_toolkit/io/kinesis/relay.py index 05bcc042..22fc7b0a 100644 --- a/cratedb_toolkit/io/kinesis/relay.py +++ b/cratedb_toolkit/io/kinesis/relay.py @@ -7,6 +7,7 @@ from commons_codec.model import SkipOperation from commons_codec.transform.aws_dms import DMSTranslatorCrateDB from commons_codec.transform.dynamodb import DynamoDBCDCTranslator +from sqlalchemy.exc import OperationalError, ProgrammingError from tqdm import tqdm from tqdm.contrib.logging import logging_redirect_tqdm from yarl import URL @@ -50,7 +51,10 @@ def __init__( if self.recipe: pks, cms, mapping_strategy, ignore_ddl = self.recipe.codec_options() self.translator = DMSTranslatorCrateDB( - primary_keys=pks, column_types=cms, mapping_strategy=mapping_strategy, ignore_ddl=ignore_ddl + primary_keys=pks, # ty: ignore[invalid-argument-type] + column_types=cms, # ty: ignore[invalid-argument-type] + mapping_strategy=mapping_strategy, # ty: ignore[invalid-argument-type] + ignore_ddl=ignore_ddl, # ty: ignore[invalid-argument-type] ) else: raise SkipAdapterException(f"Not processing {self.kinesis_url} here") @@ -67,7 +71,8 @@ def start(self, once: bool = False): try: if self.cratedb_table is not None: if not self.cratedb_adapter.table_exists(self.cratedb_table): - self.connection.execute(sa.text(self.translator.sql_ddl)) + assert self.translator and hasattr(self.translator, "sql_ddl") # noqa: S101 + self.connection.execute(sa.text(t.cast(str, self.translator.sql_ddl))) self.connection.commit() records_target = self.cratedb_adapter.count_records(self.cratedb_table) logger.info(f"Target: CrateDB table={self.cratedb_table} count={records_target}") @@ -131,7 +136,7 @@ def process_event(self, event): self.connection.execute(sa.text(f"REFRESH TABLE {self.cratedb_table}")) self.connection.commit() - except (sa.exc.ProgrammingError, sa.exc.OperationalError): + except (ProgrammingError, OperationalError): logger.exception("Executing query failed") raise else: diff --git a/cratedb_toolkit/io/mongodb/adapter.py b/cratedb_toolkit/io/mongodb/adapter.py index 2cd4c03b..bc15c731 100644 --- a/cratedb_toolkit/io/mongodb/adapter.py +++ b/cratedb_toolkit/io/mongodb/adapter.py @@ -99,6 +99,10 @@ def record_count(self, filter_=None) -> int: def query(self): raise NotImplementedError() + @abstractmethod + def create_collection(self): + raise NotImplementedError() + @abstractmethod def subscribe_cdc(self, resume_after: t.Optional[DocumentDict] = None): raise NotImplementedError() diff --git a/cratedb_toolkit/io/mongodb/cdc.py b/cratedb_toolkit/io/mongodb/cdc.py index 3e754b98..70dc3919 100644 --- a/cratedb_toolkit/io/mongodb/cdc.py +++ b/cratedb_toolkit/io/mongodb/cdc.py @@ -61,7 +61,8 @@ def __init__( transformation = None if tm: address = CollectionAddress( - container=self.mongodb_adapter.database_name, name=self.mongodb_adapter.collection_name + container=self.mongodb_adapter.database_name, + name=t.cast(str, self.mongodb_adapter.collection_name), ) try: transformation = tm.project.get(address=address) diff --git a/cratedb_toolkit/io/mongodb/copy.py b/cratedb_toolkit/io/mongodb/copy.py index 8b7e2e97..ea014348 100644 --- a/cratedb_toolkit/io/mongodb/copy.py +++ b/cratedb_toolkit/io/mongodb/copy.py @@ -50,7 +50,8 @@ def __init__( transformation = None if tm: address = CollectionAddress( - container=self.mongodb_adapter.database_name, name=self.mongodb_adapter.collection_name + container=self.mongodb_adapter.database_name, + name=t.cast(str, self.mongodb_adapter.collection_name), ) try: transformation = tm.project.get(address=address) @@ -90,7 +91,7 @@ def start(self): processor = BulkProcessor( connection=connection, data=self.mongodb_adapter.query(), - batch_to_operation=self.translator.to_sql, + batch_to_operation=self.translator.to_sql, # ty: ignore[invalid-argument-type] progress_bar=progress_bar, on_error=self.on_error, debug=self.debug, diff --git a/cratedb_toolkit/io/mongodb/core.py b/cratedb_toolkit/io/mongodb/core.py index 2d63605e..47ad1345 100644 --- a/cratedb_toolkit/io/mongodb/core.py +++ b/cratedb_toolkit/io/mongodb/core.py @@ -101,7 +101,7 @@ def extract(args) -> t.Dict[str, t.Any]: return schemas -def translate(schemas, schemaname: str = None) -> t.Dict[str, str]: +def translate(schemas, schemaname: t.Optional[str] = None) -> t.Dict[str, str]: """ Translate a given schema into SQL DDL statements compatible with CrateDB. """ diff --git a/cratedb_toolkit/io/mongodb/export.py b/cratedb_toolkit/io/mongodb/export.py index cfa17e01..5c23c19e 100644 --- a/cratedb_toolkit/io/mongodb/export.py +++ b/cratedb_toolkit/io/mongodb/export.py @@ -45,12 +45,15 @@ def convert(d): converter = MongoDBCrateDBConverter() newdict = {} for k, v in sanitize_field_names(d).items(): - newdict[k] = converter.convert(v) + newdict[k] = converter.decode_document(v) return newdict def collection_to_json( - collection: pymongo.collection.Collection, fp: t.IO[t.Any], tm: TransformationManager = None, limit: int = 0 + collection: pymongo.collection.Collection, + fp: t.IO[t.Any], + tm: t.Optional[TransformationManager] = None, + limit: t.Optional[int] = None, ): """ Export a MongoDB collection's documents to standard JSON. @@ -62,6 +65,7 @@ def collection_to_json( file a file-like object (stream). """ + limit = limit or 0 for document in collection.find().limit(limit): bson_json = bsonjs.dumps(document.raw) json_object = json.loads(bson_json) diff --git a/cratedb_toolkit/io/mongodb/transform.py b/cratedb_toolkit/io/mongodb/transform.py index 7171ae65..39c48a70 100644 --- a/cratedb_toolkit/io/mongodb/transform.py +++ b/cratedb_toolkit/io/mongodb/transform.py @@ -49,10 +49,11 @@ def apply_type_overrides(self, database_name: str, collection_name: str, collect # TODO: Also support addressing nested elements. # Hint: Implementation already exists on another machine, # where it has not been added to the repository. Sigh. - for rule in transformation.schema.rules: - pointer = JsonPointer(f"/document{rule.pointer}/types") - type_stats = pointer.resolve(collection_schema) - type_stats[rule.type] = {"count": int(9e10)} + if transformation.schema: + for rule in transformation.schema.rules: + pointer = JsonPointer(f"/document{rule.pointer}/types") + type_stats = pointer.resolve(collection_schema) + type_stats[rule.type] = {"count": int(9e10)} def apply_transformations(self, database_name: str, collection_name: str, data: t.Dict[str, t.Any]): if not self.active: diff --git a/cratedb_toolkit/io/mongodb/translate.py b/cratedb_toolkit/io/mongodb/translate.py index 9e1e6eb4..3fbc47a6 100644 --- a/cratedb_toolkit/io/mongodb/translate.py +++ b/cratedb_toolkit/io/mongodb/translate.py @@ -33,6 +33,7 @@ import logging from functools import reduce +from typing import Optional from cratedb_toolkit.io.mongodb.util import sanitize_field_names @@ -164,7 +165,7 @@ def indent_sql(query: str) -> str: return "\n".join(lines) -def translate(schemas, schemaname: str = None): +def translate(schemas, schemaname: Optional[str] = None): """ Translate a schema definition for a set of MongoDB collection schemas. diff --git a/cratedb_toolkit/model.py b/cratedb_toolkit/model.py index b0f599a2..0361e560 100644 --- a/cratedb_toolkit/model.py +++ b/cratedb_toolkit/model.py @@ -89,7 +89,7 @@ def from_http_uri(cls, url: str) -> "DatabaseAddress": uri.scheme = "crate" return cls(uri=uri) - def with_credentials(self, username: str = None, password: str = None): + def with_credentials(self, username: t.Optional[str] = None, password: t.Optional[str] = None): """ Add credentials, with in-place modification. """ diff --git a/cratedb_toolkit/query/cli.py b/cratedb_toolkit/query/cli.py index 9dc649a5..208dd09c 100644 --- a/cratedb_toolkit/query/cli.py +++ b/cratedb_toolkit/query/cli.py @@ -10,7 +10,7 @@ logger = logging.getLogger(__name__) -@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup) @click.option("--verbose", is_flag=True, required=False, help="Turn on logging") @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level") @click.version_option() diff --git a/cratedb_toolkit/query/mcp/cli.py b/cratedb_toolkit/query/mcp/cli.py index e708d69c..afacb13b 100644 --- a/cratedb_toolkit/query/mcp/cli.py +++ b/cratedb_toolkit/query/mcp/cli.py @@ -27,7 +27,7 @@ def get_format_option(default="markdown"): format_option_json = get_format_option(default="json") -@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup) @click.option("--server-name", type=str, required=False, help="Select MCP server name") @click.option("--verbose", is_flag=True, required=False, help="Turn on logging") @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level") diff --git a/cratedb_toolkit/query/mcp/pg_mcp.py b/cratedb_toolkit/query/mcp/pg_mcp.py index dbf905d7..3e471f1e 100644 --- a/cratedb_toolkit/query/mcp/pg_mcp.py +++ b/cratedb_toolkit/query/mcp/pg_mcp.py @@ -1,7 +1,7 @@ if __name__ == "__main__": # FIXME: Improve invocation after packaging has been improved. # https://github.com/stuzero/pg-mcp/issues/10 - from server.app import logger, mcp + from server.app import logger, mcp # ty: ignore[unresolved-import] # TODO: Bring flexible invocation (sse vs. stdio) to mainline. logger.info("Starting MCP server with STDIO transport") diff --git a/cratedb_toolkit/retention/cli.py b/cratedb_toolkit/retention/cli.py index ad86404f..00a2bc9c 100644 --- a/cratedb_toolkit/retention/cli.py +++ b/cratedb_toolkit/retention/cli.py @@ -136,7 +136,7 @@ def help_run(): ) -@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] +@click.group(cls=ClickAliasedGroup) @click.option("--verbose", is_flag=True, required=False, help="Turn on logging") @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level") @click.version_option() diff --git a/cratedb_toolkit/retention/store.py b/cratedb_toolkit/retention/store.py index 9481b202..97d66305 100644 --- a/cratedb_toolkit/retention/store.py +++ b/cratedb_toolkit/retention/store.py @@ -35,14 +35,14 @@ def get_tags_constraints(self, tags: t.Union[t.List[str], t.Set[str]]): """ Return list of SQL WHERE constraint clauses from given tags. """ - from sqlalchemy.sql.selectable import NamedFromClause # type: ignore[attr-defined] + from sqlalchemy.sql.selectable import NamedFromClause - table: NamedFromClause = self.table # type: ignore[attr-defined] + table: NamedFromClause = self.table # ty: ignore[unresolved-attribute] constraints = [] for tag in tags: if not tag: continue - constraint = table.c[self.tag_column][tag] != sa.Null() # type: ignore[attr-defined] + constraint = table.c[self.tag_column][tag] != sa.Null() constraints.append(constraint) return sa.and_(sa.true(), *constraints) @@ -55,7 +55,7 @@ def tags_exist(self, tags: t.Union[t.List[str], t.Set[str]]): TODO: Create corresponding issue at crate/crate. """ - table = self.table # type: ignore[attr-defined] + table = self.table # ty: ignore[unresolved-attribute] where_clause = self.get_tags_constraints(tags) if sa_is_empty(where_clause): return False @@ -87,12 +87,12 @@ def delete_by_all_tags(self, tags: t.Union[t.List[str], t.Set[str]]): logger.warning(f"No retention policies found with tags: {tags}") return 0 - table = self.table # type: ignore[attr-defined] + table = self.table # ty: ignore[unresolved-attribute] where_clause = self.get_tags_constraints(tags) if sa_is_empty(where_clause): logger.warning("Unable to compute criteria for deletion") return 0 - deletable = sa.delete(table).where(where_clause) # type: ignore[arg-type] + deletable = sa.delete(table).where(where_clause) result = self.execute(deletable) self.synchronize() return result.rowcount diff --git a/cratedb_toolkit/testing/testcontainers/azurite.py b/cratedb_toolkit/testing/testcontainers/azurite.py index 403100b0..49dbe0a6 100644 --- a/cratedb_toolkit/testing/testcontainers/azurite.py +++ b/cratedb_toolkit/testing/testcontainers/azurite.py @@ -43,7 +43,7 @@ def get_real_host_address(self) -> str: def get_container_endpoint(self, container_name: str): container = self.get_container(container_name) hostname = self.get_real_host_address() - return container._format_url(hostname=hostname) + return container._format_url(hostname=hostname) # ty: ignore[unresolved-attribute] def get_blob_service_client(self) -> BlobServiceClient: """ diff --git a/cratedb_toolkit/testing/testcontainers/cratedb.py b/cratedb_toolkit/testing/testcontainers/cratedb.py index b1de6ae9..999d913d 100644 --- a/cratedb_toolkit/testing/testcontainers/cratedb.py +++ b/cratedb_toolkit/testing/testcontainers/cratedb.py @@ -19,7 +19,6 @@ import re from typing import Optional -from testcontainers.core.config import MAX_TRIES from testcontainers.core.generic import DbContainer from testcontainers.core.wait_strategies import LogMessageWaitStrategy from testcontainers.core.waiting_utils import wait_for_logs @@ -152,7 +151,9 @@ def get_connection_url(self, dialect: str = "crate", host: Optional[str] = None) ) def _connect(self): - wait_for_logs(self, predicate=self._wait_strategy, timeout=MAX_TRIES) + if not self._wait_strategy: + raise ValueError("No wait strategy defined") + wait_for_logs(self, predicate=self._wait_strategy) def _configure_wait_condition(self): """Wait for CrateDB node to be fully started.""" @@ -169,8 +170,8 @@ class CrateDBTestAdapter: """ def __init__(self, crate_version: str = "nightly", **kwargs): - self.cratedb: Optional[CrateDBContainer] = None - self.database: Optional[DatabaseAdapter] = None + self.cratedb: CrateDBContainer + self.database: DatabaseAdapter self.image: str = "crate/crate:{}".format(crate_version) def start(self, **kwargs): diff --git a/cratedb_toolkit/testing/testcontainers/influxdb2.py b/cratedb_toolkit/testing/testcontainers/influxdb2.py index 7a37edb2..dccdec0d 100644 --- a/cratedb_toolkit/testing/testcontainers/influxdb2.py +++ b/cratedb_toolkit/testing/testcontainers/influxdb2.py @@ -76,7 +76,9 @@ def get_connection_url(self, host=None) -> str: port=self.port_to_expose, ) - def _connect(self) -> InfluxDBClient: + def _connect(self) -> InfluxDBClient: # ty: ignore[invalid-method-override] + if not self._wait_strategy: + raise ValueError("No wait strategy defined") wait_for_logs(self, self._wait_strategy) return InfluxDBClient(url=self.get_connection_url(), org=self.ORGANIZATION, token=self.TOKEN, debug=self.debug) diff --git a/cratedb_toolkit/testing/testcontainers/localstack.py b/cratedb_toolkit/testing/testcontainers/localstack.py index 2c68dddb..02a50001 100644 --- a/cratedb_toolkit/testing/testcontainers/localstack.py +++ b/cratedb_toolkit/testing/testcontainers/localstack.py @@ -55,4 +55,6 @@ def _connect(self): the "Ready" log message. Without this, Kinesis and other service APIs receive requests before LocalStack is ready. """ + if not self._wait_strategy: + raise ValueError("No wait strategy defined") wait_for_logs(self, predicate=self._wait_strategy, timeout=60) diff --git a/cratedb_toolkit/testing/testcontainers/mongodb.py b/cratedb_toolkit/testing/testcontainers/mongodb.py index 6b532c57..82184be5 100644 --- a/cratedb_toolkit/testing/testcontainers/mongodb.py +++ b/cratedb_toolkit/testing/testcontainers/mongodb.py @@ -80,6 +80,8 @@ def _configure(self) -> None: def _create_connection_url( self, dialect: str, + username: str, + password: str, host: t.Optional[str] = None, port: t.Optional[int] = None, dbname: t.Optional[str] = None, @@ -92,6 +94,7 @@ def _create_connection_url( if self._container is None: raise ContainerStartException("container has not been started") host = host or self.get_container_host_ip() + assert port is not None # noqa: S101 port = self.get_exposed_port(port) url = f"{dialect}://{host}:{port}" if dbname: @@ -101,6 +104,8 @@ def _create_connection_url( def get_connection_url(self) -> str: return self._create_connection_url( dialect="mongodb", + username="admin", + password="", port=self.port, ) diff --git a/cratedb_toolkit/testing/testcontainers/util.py b/cratedb_toolkit/testing/testcontainers/util.py index 465925fe..b947d075 100644 --- a/cratedb_toolkit/testing/testcontainers/util.py +++ b/cratedb_toolkit/testing/testcontainers/util.py @@ -38,6 +38,7 @@ def get_real_host_ip(self) -> str: To let containers talk to each other, explicitly provide the real IP address of the container. In corresponding jargon, it appears to be the "bridge IP". """ + assert self._container is not None # noqa: S101 return self.get_docker_client().bridge_ip(self._container.id) def get_real_host_address(self) -> str: @@ -108,7 +109,7 @@ def start(self): command=self._command, detach=True, environment=self.env, - ports=self.ports, + ports=self.ports or {}, # ty: ignore[invalid-argument-type] name=self._name, volumes=self.volumes, **self._kwargs, @@ -124,11 +125,11 @@ def start(self): logger.info(f"Starting container: {container_id} ({container_name})") self._container.start() - if hasattr(self, "_connect"): - self._connect() + if hasattr(self, "_connect") and callable(self._connect): + self._connect() # ty: ignore[call-top-callable] return self - def stop(self, **kwargs): + def stop(self, force: bool = True, delete_volume: bool = True) -> None: """ Shut down container again, unless "keepalive" is enabled. """ @@ -156,7 +157,7 @@ def __init__(self, *args, **kwargs): except DockerException as ex: if any(token in str(ex) for token in ("Connection aborted", "Error while fetching server API version")): # TODO: Make this configurable through some `pytest_` variable. - raise pytest.skip(reason="Skipping test because Docker is not running", allow_module_level=True) from ex + pytest.skip("Skipping test because Docker is not running", allow_module_level=True) # ty: ignore[invalid-argument-type,too-many-positional-arguments] else: # noqa: RET506 raise @@ -175,7 +176,7 @@ class PytestTestcontainerAdapter: """ def __init__(self): - self.container: DockerContainer = None + self.container: DockerContainer self.run_setup() @abstractmethod @@ -196,9 +197,7 @@ def run_setup(self): except DockerException as ex: if any(token in str(ex) for token in ("Connection aborted", "Error while fetching server API version")): # TODO: Make this configurable through some `pytest_` variable. - raise pytest.skip( - reason="Skipping test because Docker daemon is not available", allow_module_level=True - ) from ex + pytest.skip("Skipping test because Docker daemon is not available", allow_module_level=True) # ty: ignore[invalid-argument-type,too-many-positional-arguments] else: # noqa: RET506 raise diff --git a/cratedb_toolkit/util/app.py b/cratedb_toolkit/util/app.py index 0596e8e6..a8ada8bd 100644 --- a/cratedb_toolkit/util/app.py +++ b/cratedb_toolkit/util/app.py @@ -9,7 +9,7 @@ def make_cli(): - @click.group(cls=ClickAliasedGroup) # type: ignore[arg-type] + @click.group(cls=ClickAliasedGroup) @option_cluster_id @option_cluster_name @option_cluster_url diff --git a/cratedb_toolkit/util/cli.py b/cratedb_toolkit/util/cli.py index 17d27b00..6cc2d413 100644 --- a/cratedb_toolkit/util/cli.py +++ b/cratedb_toolkit/util/cli.py @@ -33,7 +33,7 @@ def split_list(value: str, delimiter: str = ",") -> t.List[str]: return [c.strip() for c in value.split(delimiter)] -def to_list(x: t.Any, default: t.List[t.Any] = None) -> t.List[t.Any]: +def to_list(x: t.Any, default: t.Optional[t.List[t.Any]] = None) -> t.List[t.Any]: if not isinstance(default, t.List): raise ValueError("Default value is not a list") if x is None: @@ -118,7 +118,7 @@ def error_level_by_debug(debug: bool): def running_with_debug(ctx: click.Context) -> bool: - return ( + return bool( (ctx.parent and ctx.parent.params.get("debug", False)) or (ctx.parent and ctx.parent.parent and ctx.parent.parent.params.get("debug", False)) or False diff --git a/cratedb_toolkit/util/client.py b/cratedb_toolkit/util/client.py index d2e53b16..2e00ae9a 100644 --- a/cratedb_toolkit/util/client.py +++ b/cratedb_toolkit/util/client.py @@ -1,4 +1,5 @@ import contextlib +from typing import Optional from unittest.mock import patch import crate.client.http @@ -11,20 +12,22 @@ class jwt_token_patch(contextlib.ContextDecorator): Patch the `Client._request` method to add the Authorization header for JWT token-based authentication. """ - def __init__(self, jwt_token: str = None): + def __init__(self, jwt_token: Optional[str] = None): self.jwt_token = jwt_token def __enter__(self): - self.patcher = patch.object(crate.client.http.Client, "_request", _mk_crate_client_request(self.jwt_token)) - self.patcher.start() + if self.jwt_token: + self.patcher = patch.object(crate.client.http.Client, "_request", _mk_crate_client_request(self.jwt_token)) + self.patcher.start() return self def __exit__(self, type, value, tb): # noqa: A002 - self.patcher.stop() + if self.jwt_token: + self.patcher.stop() return self -def _mk_crate_client_request(jwt_token: str = None): +def _mk_crate_client_request(jwt_token: Optional[str] = None): """ Create a monkey patched `Client._request` method to add the Authorization header for JWT token-based authentication. """ diff --git a/cratedb_toolkit/util/crash.py b/cratedb_toolkit/util/crash.py index 1e7ecc07..29fcfb3f 100644 --- a/cratedb_toolkit/util/crash.py +++ b/cratedb_toolkit/util/crash.py @@ -12,11 +12,11 @@ def run_crash( hosts: str, command: str, - output_format: str = None, - schema: str = None, - username: str = None, - password: str = None, - jwt_token: str = None, + output_format: t.Optional[str] = None, + schema: t.Optional[str] = None, + username: t.Optional[str] = None, + password: t.Optional[str] = None, + jwt_token: t.Optional[str] = None, ): """ Run the interactive CrateDB database shell using `crash`. diff --git a/cratedb_toolkit/util/croud.py b/cratedb_toolkit/util/croud.py index 2cf0aa85..bb2d02ab 100644 --- a/cratedb_toolkit/util/croud.py +++ b/cratedb_toolkit/util/croud.py @@ -35,7 +35,7 @@ class CroudCall: class CroudWrapper: - def __init__(self, call: CroudCall, output_format: str = None, decode_output: bool = True): + def __init__(self, call: CroudCall, output_format: t.Optional[str] = None, decode_output: t.Optional[bool] = True): """ format: One of table,wide,json,yaml """ @@ -257,7 +257,7 @@ def get_headless_config(cls) -> Configuration: return config -croud.api.Client = CroudClient +croud.api.Client = CroudClient # ty: ignore[invalid-assignment] def get_sane_log_level(src) -> int: @@ -270,7 +270,7 @@ def get_sane_log_level(src) -> int: return level -def get_croud_output_formats() -> t.List[str]: +def get_croud_output_formats() -> t.Iterable[str]: """ Inquire the output formats `croud` understands. """ diff --git a/cratedb_toolkit/util/database.py b/cratedb_toolkit/util/database.py index 2a2f313c..9eefd0d2 100644 --- a/cratedb_toolkit/util/database.py +++ b/cratedb_toolkit/util/database.py @@ -24,7 +24,7 @@ try: from typing import Literal except ImportError: - from typing_extensions import Literal # type: ignore[assignment] + from typing_extensions import Literal if t.TYPE_CHECKING: from cratedb_toolkit.cluster.model import JwtResponse @@ -47,7 +47,7 @@ class DatabaseAdapter: internal_tag = " -- ctk" - def __init__(self, dburi: str, echo: bool = False, internal: bool = False, jwt: "JwtResponse" = None): + def __init__(self, dburi: str, echo: bool = False, internal: bool = False, jwt: t.Optional["JwtResponse"] = None): if not dburi: raise ValueError("Database URI must be specified") if dburi.startswith("crate://"): @@ -109,9 +109,9 @@ def quote_relation_name(ident: str) -> str: def run_sql( self, sql: t.Union[str, Path, io.IOBase], - parameters: t.Mapping[str, str] = None, - records: bool = False, - ignore: str = None, + parameters: t.Optional[t.Mapping[str, str]] = None, + records: t.Optional[bool] = False, + ignore: t.Optional[str] = None, ): """ Run SQL statement and return results, optionally ignoring exceptions. @@ -136,7 +136,9 @@ def run_sql( raise return None - def run_sql_real(self, sql: str, parameters: t.Mapping[str, str] = None, records: bool = False): + def run_sql_real( + self, sql: str, parameters: t.Optional[t.Mapping[str, str]] = None, records: t.Optional[bool] = False + ): """ Invoke an SQL statement and return results. """ @@ -365,8 +367,8 @@ def import_csv_dask( index=False, chunksize=1000, if_exists="replace", - npartitions: int = None, - progress: bool = False, + npartitions: t.Optional[int] = None, + progress: t.Optional[bool] = False, ): """ Import CSV data using Dask. diff --git a/cratedb_toolkit/util/pandas.py b/cratedb_toolkit/util/pandas.py index b5a088c7..81ed1b09 100644 --- a/cratedb_toolkit/util/pandas.py +++ b/cratedb_toolkit/util/pandas.py @@ -3,7 +3,7 @@ import sqlalchemy as sa -def patch_pandas_sqltable_with_dialect_parameters(table_kwargs: t.Dict = None): +def patch_pandas_sqltable_with_dialect_parameters(table_kwargs: t.Optional[t.Dict] = None): """ When using pandas' `to_sql` function, configure the SQLAlchemy database dialect implementation using custom dialect parameters. @@ -27,7 +27,7 @@ def _create_table_setup(self): # Activate enhancement code. import pandas.io.sql - pandas.io.sql.SQLTable = SQLTableWithDialectParameters + pandas.io.sql.SQLTable = SQLTableWithDialectParameters # ty: ignore[invalid-assignment] def patch_pandas_sqltable_with_extended_mapping(): @@ -49,12 +49,13 @@ def _get_column_names_and_types(self, dtype_mapper, errors: str = "raise"): raise ValueError("pandas indexes not supported yet") else: + assert self.frame is not None # noqa: S101 values = self.frame._get_value(0, name) if isinstance(values, list): first_value = values[0] first_value_type = type(first_value) # TODO: Use `boltons.remap`. - sqlalchemy_type = ARRAY_TYPE_MAP.get(first_value_type) + sqlalchemy_type = ARRAY_TYPE_MAP.get(first_value_type) # ty: ignore[invalid-argument-type] if sqlalchemy_type is None: raise TypeError(f"Data type not supported yet: List[{first_value_type}]") @@ -73,7 +74,7 @@ def _execute_create(self) -> None: # Activate enhancement code. import pandas.io.sql - pandas.io.sql.SQLTable = SQLTableWithDtypeMappersForCrateDB + pandas.io.sql.SQLTable = SQLTableWithDtypeMappersForCrateDB # ty: ignore[invalid-assignment] ARRAY_TYPE_MAP = { diff --git a/cratedb_toolkit/util/platform.py b/cratedb_toolkit/util/platform.py index 0c588dfd..ef630a63 100644 --- a/cratedb_toolkit/util/platform.py +++ b/cratedb_toolkit/util/platform.py @@ -30,7 +30,7 @@ def libraries(): data["sqlalchemy"] = { "dialects_builtin": list(sqlalchemy.dialects.registry.impls.keys()), - "dialects_3rdparty": [dialect.name for dialect in entry_points(group="sqlalchemy.dialects")], # type: ignore[attr-defined,call-arg] + "dialects_3rdparty": [dialect.name for dialect in entry_points(group="sqlalchemy.dialects")], # ty: ignore[unknown-argument,unresolved-attribute] "plugins": list(sqlalchemy.dialects.plugins.impls.keys()), } except Exception: # noqa: S110 diff --git a/cratedb_toolkit/util/runtime.py b/cratedb_toolkit/util/runtime.py index 8d620190..4beea052 100644 --- a/cratedb_toolkit/util/runtime.py +++ b/cratedb_toolkit/util/runtime.py @@ -8,7 +8,7 @@ logger = logging.getLogger() -def flexfun(domain: t.Literal["runtime", "settings"] = None): +def flexfun(domain: t.Optional[t.Literal["runtime", "settings"]] = None): """ Function decorator, which honors toolkit environment settings wrt. error handling. diff --git a/cratedb_toolkit/util/setting.py b/cratedb_toolkit/util/setting.py index d1099818..476f976d 100644 --- a/cratedb_toolkit/util/setting.py +++ b/cratedb_toolkit/util/setting.py @@ -37,7 +37,7 @@ def init_dotenv(): load_dotenv(dotenv_file) -def obtain_settings(specs: t.List[Setting], prog_name: str = None) -> t.Dict[str, str]: +def obtain_settings(specs: t.List[Setting], prog_name: t.Optional[str] = None) -> t.Dict[str, str]: """ Employ command-line parsing at runtime, using the `click` parser. @@ -74,7 +74,10 @@ def obtain_settings(specs: t.List[Setting], prog_name: str = None) -> t.Dict[str def check_mutual_exclusiveness( - specs: t.List[Setting], settings: t.Dict[str, str], message_none: str = None, message_multiple: str = None + specs: t.List[Setting], + settings: t.Dict[str, str], + message_none: t.Optional[str] = None, + message_multiple: t.Optional[str] = None, ): """ Check settings for mutual exclusiveness. diff --git a/cratedb_toolkit/util/sqlalchemy.py b/cratedb_toolkit/util/sqlalchemy.py index 7838106d..0f1d6932 100644 --- a/cratedb_toolkit/util/sqlalchemy.py +++ b/cratedb_toolkit/util/sqlalchemy.py @@ -13,4 +13,4 @@ def patch_types_map(): Register missing timestamp data type. """ # TODO: Submit patch to `crate-python`. - TYPES_MAP["timestamp without time zone"] = sqltypes.TIMESTAMP + TYPES_MAP["timestamp without time zone"] = sqltypes.TIMESTAMP # ty: ignore[invalid-assignment] diff --git a/pyproject.toml b/pyproject.toml index cd841055..2994c56c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -138,10 +138,10 @@ optional-dependencies.deltalake = [ ] optional-dependencies.develop = [ "black[jupyter]<27", - "mypy<1.20", "poethepoet<1", "pyproject-fmt<3", "ruff<0.16", + "ty==0.0.25", "validate-pyproject<1", ] optional-dependencies.docs = [ @@ -338,16 +338,29 @@ lint.per-file-ignores."doc/conf.py" = [ "A001", "ERA001" ] lint.per-file-ignores."examples/*" = [ "ERA001", "F401", "T201", "T203" ] # Allow `print` and `pprint` lint.per-file-ignores."tests/*" = [ "S101" ] # Allow use of `assert`, and `print`. -[tool.mypy] -packages = [ "cratedb_toolkit" ] -exclude = [ - "cratedb_toolkit/adapter/pymongo/", -] -check_untyped_defs = true -ignore_missing_imports = true -implicit_optional = true -install_types = true -non_interactive = true +[tool.ty] +src.include = [ "cratedb_toolkit", "tests" ] +rules.unused-ignore-comment = "ignore" +analysis.allowed-unresolved-imports = [ + "bson.**", + "bsonjs.**", + "datasets.**", + "dlt.**", + "dlt_cratedb.**", + "fastapi.**", + "importlib_metadata.**", + "importlib_resources.**", + "ingestr.**", + "jessiql.**", + "kaggle.**", + "kinesis", + "lorrystream.**", + "mcp.**", + "pymongo.**", + "sqlalchemy.**", + "tikray.**", + "undatum.**", +] [tool.pytest] ini_options.addopts = "-rfEXs -p pytester --strict-markers --verbosity=3 --cov=. --cov-report=term-missing --cov-report=xml" @@ -417,7 +430,7 @@ tasks.lint = [ { cmd = "ruff format --check" }, { cmd = "ruff check" }, { cmd = "validate-pyproject pyproject.toml" }, - { cmd = "mypy" }, + { cmd = "ty check" }, ] tasks.release = [ { cmd = "python -m build" }, diff --git a/tests/adapter/test_pymongo.py b/tests/adapter/test_pymongo.py index aed2acec..4685882a 100644 --- a/tests/adapter/test_pymongo.py +++ b/tests/adapter/test_pymongo.py @@ -14,7 +14,7 @@ check_sqlalchemy1(allow_module_level=True) if Version(pymongo.version) >= Version("4.9"): - raise pytest.skip("This feature or subsystem needs PyMongo 4.8", allow_module_level=True) + pytest.skip("This feature or subsystem needs PyMongo 4.8", allow_module_level=True) # ty: ignore[invalid-argument-type,too-many-positional-arguments] from cratedb_toolkit.adapter.pymongo import PyMongoCrateDBAdapter from cratedb_toolkit.adapter.pymongo.util import AmendedObjectId diff --git a/tests/cfr/test_systable.py b/tests/cfr/test_systable.py index aa99bdcf..3652a425 100644 --- a/tests/cfr/test_systable.py +++ b/tests/cfr/test_systable.py @@ -1,5 +1,5 @@ # ruff: noqa: E402 -import importlib +import importlib.metadata import json import os.path import re @@ -10,6 +10,8 @@ import pytest from verlib2 import Version +import tests.cfr + pymongo = pytest.importorskip("polars", reason="Skipping tests because polars is not installed") import tests diff --git a/tests/cluster/test_core.py b/tests/cluster/test_core.py index 5481962a..4f8ddd51 100644 --- a/tests/cluster/test_core.py +++ b/tests/cluster/test_core.py @@ -34,6 +34,7 @@ def test_cluster_universal_managed(cloud_environment, cloud_cluster_name): Validate the universal `DatabaseCluster` factory for "managed" CrateDB Cloud clusters. """ mc = DatabaseCluster.create(cluster_name=cloud_cluster_name).probe() + assert mc.info assert uuid.UUID(mc.info.cloud_id).time >= 774933237065562038 assert mc.info.cloud_name == cloud_cluster_name diff --git a/tests/conftest.py b/tests/conftest.py index a7538423..128e8a24 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,7 +9,7 @@ import cratedb_toolkit from cratedb_toolkit.cluster.core import ManagedClusterSettings -from cratedb_toolkit.testing.testcontainers.cratedb import CrateDBTestAdapter +from cratedb_toolkit.testing.testcontainers.cratedb import CrateDBContainer, CrateDBTestAdapter from cratedb_toolkit.testing.testcontainers.util import PytestTestcontainerAdapter from cratedb_toolkit.util.common import setup_logging from cratedb_toolkit.util.database import DatabaseAdapter @@ -41,8 +41,8 @@ class CrateDBFixture(PytestTestcontainerAdapter): """ def __init__(self): - self.container = None - self.database: DatabaseAdapter = None + self.container: CrateDBContainer + self.database: DatabaseAdapter super().__init__() def setup(self): @@ -161,7 +161,7 @@ def cloud_environment(mocker, cloud_cluster_name) -> t.Generator[t.Dict[str, str } if any(not setting for setting in settings.values()): - raise pytest.skip("Missing environment variables for headless mode with croud") + pytest.skip("Missing environment variables for headless mode with croud") # ty: ignore[invalid-argument-type,too-many-positional-arguments] mocker.patch.dict("os.environ", settings) @@ -169,7 +169,8 @@ def cloud_environment(mocker, cloud_cluster_name) -> t.Generator[t.Dict[str, str settings_accept_env=True, ) - yield settings + # TODO: Yield expression type does not match annotation. Why? + yield settings # ty: ignore[invalid-yield] cratedb_toolkit.configure( settings_accept_env=False, @@ -193,7 +194,7 @@ def check_sqlalchemy1(**kwargs): Skip pytest test cases or modules testing subsystems which need SQLAlchemy 1.x. """ if not IS_SQLALCHEMY1: - raise pytest.skip("This feature or subsystem needs SQLAlchemy 1.x", **kwargs) + pytest.skip("This feature or subsystem needs SQLAlchemy 1.x", **kwargs) # ty: ignore[invalid-argument-type,too-many-positional-arguments] @pytest.fixture @@ -209,7 +210,7 @@ def check_sqlalchemy2(**kwargs): Skip pytest test cases or modules testing subsystems which need SQLAlchemy 2.x. """ if not IS_SQLALCHEMY2: - raise pytest.skip("This feature or subsystem needs SQLAlchemy 2.x", **kwargs) + pytest.skip("This feature or subsystem needs SQLAlchemy 2.x", **kwargs) # ty: ignore[invalid-argument-type,too-many-positional-arguments] @pytest.fixture diff --git a/tests/datasets/test_loading.py b/tests/datasets/test_loading.py index 348d3d9d..fe965325 100644 --- a/tests/datasets/test_loading.py +++ b/tests/datasets/test_loading.py @@ -34,7 +34,7 @@ def test_dataset_replace(database: DatabaseAdapter, dataset: Dataset): "KAGGLE_USERNAME" in os.environ and "KAGGLE_KEY" in os.environ ) if dataset.title and "Weather" in dataset.title and not kaggle_auth_exists: - raise pytest.skip(f"Kaggle dataset can not be tested without authentication: {dataset.reference}") + pytest.skip(f"Kaggle dataset can not be tested without authentication: {dataset.reference}") # ty: ignore[invalid-argument-type,too-many-positional-arguments] dataset = load_dataset(dataset.reference) tablename = slugify(dataset.reference, separator="_") diff --git a/tests/examples/test_python.py b/tests/examples/test_python.py index 9fcd4d81..7e350a29 100644 --- a/tests/examples/test_python.py +++ b/tests/examples/test_python.py @@ -37,6 +37,6 @@ def test_managed_import_notebook(cloud_environment): # Execute the notebook. notebook = Path("examples") / "notebook" / "cloud_import.ipynb" if not notebook.exists(): - pytest.fail(f"Notebook not found: {notebook}") + pytest.fail(f"Notebook not found: {notebook}") # ty: ignore[invalid-argument-type] with testbook(notebook, timeout=180) as tb: tb.execute() diff --git a/tests/examples/test_shell.py b/tests/examples/test_shell.py index 8ebc3b89..a9252dc0 100644 --- a/tests/examples/test_shell.py +++ b/tests/examples/test_shell.py @@ -14,5 +14,5 @@ def test_example_cloud_cluster_shell(cloud_environment): """ program = ROOT / "tests" / "examples" / "test_shell.sh" if not program.exists(): - pytest.fail(f"Test program not found: {program}") + pytest.fail(f"Test program not found: {program}") # ty: ignore[invalid-argument-type] subprocess.check_call([str(program)]) # noqa: S603 diff --git a/tests/io/dynamodb/conftest.py b/tests/io/dynamodb/conftest.py index 0ca1fa7c..d5e72c3d 100644 --- a/tests/io/dynamodb/conftest.py +++ b/tests/io/dynamodb/conftest.py @@ -1,6 +1,5 @@ # ruff: noqa: E402 import logging -import typing import uuid import pytest @@ -10,6 +9,7 @@ pytest.importorskip("kinesis", reason="Skipping DynamoDB tests because 'async-kinesis' package is not installed") import botocore +import botocore.exceptions from yarl import URL from cratedb_toolkit.io.dynamodb.adapter import DynamoDBAdapter @@ -32,9 +32,11 @@ class DynamoDBFixture: """ def __init__(self): - self.container = None + from cratedb_toolkit.testing.testcontainers.localstack import LocalStackContainerWithKeepalive + + self.container: LocalStackContainerWithKeepalive self.url = None - self.dynamodb_adapter: typing.Union[DynamoDBAdapter, None] = None + self.dynamodb_adapter: DynamoDBAdapter self._stream_name = f"demo-{uuid.uuid4().hex[:8]}" self.setup() @@ -49,7 +51,8 @@ def setup(self): self.dynamodb_adapter = DynamoDBAdapter(URL(f"{self.get_connection_url_dynamodb()}?region=us-east-1")) def finalize(self): - self.container.stop() + if self.container: + self.container.stop() def reset(self): """ diff --git a/tests/io/dynamodb/test_adapter.py b/tests/io/dynamodb/test_adapter.py index cc419619..23916c45 100644 --- a/tests/io/dynamodb/test_adapter.py +++ b/tests/io/dynamodb/test_adapter.py @@ -26,7 +26,7 @@ def test_adapter_scan_failure_consistent_read(dynamodb): adapter = DynamoDBAdapter(URL(dynamodb_url)) with pytest.raises(ParamValidationError) as ex: - next(adapter.scan("demo", consistent_read=-42, on_error="raise")) + next(adapter.scan("demo", consistent_read=-42, on_error="raise")) # ty: ignore[invalid-argument-type] assert ex.match("Parameter validation failed:\nInvalid type for parameter ConsistentRead, value: -42.*") diff --git a/tests/io/dynamodb/test_relay.py b/tests/io/dynamodb/test_relay.py index 9600aaed..df81ac03 100644 --- a/tests/io/dynamodb/test_relay.py +++ b/tests/io/dynamodb/test_relay.py @@ -94,6 +94,7 @@ def test_kinesis_latest_dynamodb_cdc_insert_update(caplog, cratedb, dynamodb): thread = threading.Thread(target=table_loader.start) thread.start() # Wait for the consumer to obtain its LATEST shard iterator before producing events. + assert table_loader.kinesis_adapter is not None assert table_loader.kinesis_adapter.wait_until_ready(timeout=30), "Consumer did not become ready in time" # Populate source database with data. diff --git a/tests/io/influxdb/conftest.py b/tests/io/influxdb/conftest.py index 5f86f68c..9f25bb1b 100644 --- a/tests/io/influxdb/conftest.py +++ b/tests/io/influxdb/conftest.py @@ -21,8 +21,10 @@ class InfluxDB2Fixture(PytestTestcontainerAdapter): def __init__(self): from influxdb_client import InfluxDBClient - self.container = None - self.client: InfluxDBClient = None + from cratedb_toolkit.testing.testcontainers.influxdb2 import InfluxDB2Container + + self.container: InfluxDB2Container + self.client: InfluxDBClient super().__init__() def setup(self): diff --git a/tests/io/kinesis/conftest.py b/tests/io/kinesis/conftest.py index fad07cec..d9233382 100644 --- a/tests/io/kinesis/conftest.py +++ b/tests/io/kinesis/conftest.py @@ -24,7 +24,9 @@ class KinesisFixture: """ def __init__(self): - self.container = None + from cratedb_toolkit.testing.testcontainers.localstack import LocalStackContainerWithKeepalive + + self.container: LocalStackContainerWithKeepalive self.url = None self.kinesis_adapter: typing.Union[KinesisStreamAdapter, None] = None self._stream_name = "testdrive" diff --git a/tests/io/kinesis/manager.py b/tests/io/kinesis/manager.py index 1f30716c..661a3a59 100644 --- a/tests/io/kinesis/manager.py +++ b/tests/io/kinesis/manager.py @@ -11,8 +11,8 @@ class KinesisTestManager: def __init__(self, url: str): - url = URL(url).with_query({"region": "us-east-1", "create": "true"}) - self.adapter = KinesisAdapterBase.factory(url) + url_obj = URL(url).with_query({"region": "us-east-1", "create": "true"}) + self.adapter = KinesisAdapterBase.factory(url_obj) def load_events(self, events: t.List[t.Dict[str, t.Any]]): for event in events: diff --git a/tests/io/kinesis/test_relay.py b/tests/io/kinesis/test_relay.py index c39cec40..15c784a3 100644 --- a/tests/io/kinesis/test_relay.py +++ b/tests/io/kinesis/test_relay.py @@ -3,6 +3,7 @@ import pytest import sqlalchemy as sa +from sqlalchemy.exc import OperationalError from cratedb_toolkit.io.kinesis.model import RecipeDefinition from cratedb_toolkit.io.kinesis.relay import KinesisRelay @@ -91,12 +92,12 @@ def test_kinesis_relay_write_failure_propagates(caplog, cratedb, kinesis): def failing_execute(self, statement, *args, **kwargs): stmt_text = statement.text if hasattr(statement, "text") else str(statement) if "INSERT" in stmt_text.upper() or "UPSERT" in stmt_text.upper(): - raise sa.exc.OperationalError("test", {}, Exception("connection lost")) + raise OperationalError("test", {}, Exception("connection lost")) return original_execute(self, statement, *args, **kwargs) # Verify that the write failure propagates instead of being swallowed. with patch.object(sa.Connection, "execute", failing_execute): - with pytest.raises(sa.exc.OperationalError, match="connection lost"): + with pytest.raises(OperationalError, match="connection lost"): table_loader.start(once=True) # Verify that stop() cleaned up the connection. diff --git a/tests/io/mongodb/conftest.py b/tests/io/mongodb/conftest.py index 8da35a63..13d211d9 100644 --- a/tests/io/mongodb/conftest.py +++ b/tests/io/mongodb/conftest.py @@ -1,5 +1,6 @@ import logging import os +import typing as t import pytest @@ -30,9 +31,14 @@ class MongoDBFixture(PytestTestcontainerAdapter): def __init__(self, container_class): from pymongo import MongoClient + from cratedb_toolkit.testing.testcontainers.mongodb import ( + MongoDbContainerWithKeepalive, + MongoDbReplicasetContainer, + ) + self.container_class = container_class - self.container = None - self.client: MongoClient = None + self.container: t.Union[MongoDbContainerWithKeepalive, MongoDbReplicasetContainer] + self.client: MongoClient super().__init__() def setup(self): @@ -63,9 +69,6 @@ def get_connection_url(self): def get_connection_client(self): return self.container.get_connection_client() - def get_connection_client_replicaset(self): - return self.container.get_connection_client_replicaset() - class MongoDBFixtureFactory: def __init__(self, container): diff --git a/tests/io/mongodb/test_integration.py b/tests/io/mongodb/test_integration.py index 57526efc..7451800c 100644 --- a/tests/io/mongodb/test_integration.py +++ b/tests/io/mongodb/test_integration.py @@ -5,6 +5,7 @@ from unittest import mock import pymongo +import pymongo.errors import pytest from tests.conftest import check_sqlalchemy2 @@ -48,7 +49,7 @@ def setUpClass(cls): logger.debug(f"MongoDB server info: {server_info}") except pymongo.errors.ServerSelectionTimeoutError as ex: if cls.SKIP_IF_NOT_RUNNING: - raise cls.skipTest(cls, reason="MongoDB server not running") from ex + raise cls.skipTest(cls, reason="MongoDB server not running") from ex # ty: ignore[invalid-argument-type] else: # noqa: RET506 raise diff --git a/tests/io/test_iceberg.py b/tests/io/test_iceberg.py index 5eac7cb0..22cbe631 100644 --- a/tests/io/test_iceberg.py +++ b/tests/io/test_iceberg.py @@ -10,7 +10,7 @@ pl = pytest.importorskip("polars", reason="Skipping Iceberg tests because 'polars' package is not installed") if not hasattr(pd.DataFrame, "to_iceberg"): - pytest.skip("Older pandas releases do not support Apache Iceberg", allow_module_level=True) + pytest.skip("Older pandas releases do not support Apache Iceberg", allow_module_level=True) # ty: ignore[invalid-argument-type,too-many-positional-arguments] @pytest.fixture diff --git a/tests/shell/test_cli.py b/tests/shell/test_cli.py index b2569c35..ab94f3d6 100644 --- a/tests/shell/test_cli.py +++ b/tests/shell/test_cli.py @@ -36,7 +36,7 @@ def test_shell_managed_jwt(mocker, cloud_cluster_name): } if any(not setting for setting in settings.values()): - raise pytest.skip("Missing environment variables for headless mode with croud") + pytest.skip("Missing environment variables for headless mode with croud") # ty: ignore[invalid-argument-type,too-many-positional-arguments] # Synthesize a valid environment. mocker.patch.dict("os.environ", settings) @@ -65,7 +65,7 @@ def test_shell_managed_username_password(mocker, cloud_cluster_name): } if any(not setting for setting in settings.values()): - raise pytest.skip("Missing environment variables for headless mode with croud") + pytest.skip("Missing environment variables for headless mode with croud") # ty: ignore[invalid-argument-type,too-many-positional-arguments] # Synthesize a valid environment. mocker.patch.dict("os.environ", settings) diff --git a/tests/util/test_run_sql.py b/tests/util/test_run_sql.py index 33a67b17..9f08c92c 100644 --- a/tests/util/test_run_sql.py +++ b/tests/util/test_run_sql.py @@ -1,7 +1,7 @@ import io import pytest -import sqlalchemy as sa +from sqlalchemy.exc import OperationalError from cratedb_toolkit.util.database import run_sql @@ -52,7 +52,7 @@ def test_run_sql_multiple_statements(sqlcmd): def test_run_sql_invalid_host(capsys): - with pytest.raises(sa.exc.OperationalError) as ex: + with pytest.raises(OperationalError) as ex: run_sql(dburi="crate://localhost:12345", sql="SELECT 1;") assert ex.match( ".*ConnectionError.*No more Servers available.*HTTPConnectionPool.*" From d3be2ba5eb384c55dae143faf88f30cbe9fbbab1 Mon Sep 17 00:00:00 2001 From: Andreas Motl Date: Wed, 25 Mar 2026 22:14:39 +0100 Subject: [PATCH 2/2] QA: Migrate type checker from `mypy` to `ty` -- CodeRabbit improvements --- cratedb_toolkit/io/awslambda/kinesis.py | 7 +++++- cratedb_toolkit/io/mongodb/cdc.py | 1 + cratedb_toolkit/io/mongodb/copy.py | 8 ++++++- .../testing/testcontainers/cratedb.py | 24 ++++++++++++------- .../testing/testcontainers/mongodb.py | 22 ++++++++++++----- .../testing/testcontainers/util.py | 7 +++++- cratedb_toolkit/util/client.py | 11 +++++---- cratedb_toolkit/util/database.py | 8 +++---- tests/io/mongodb/conftest.py | 13 ++++++++-- tests/io/mongodb/test_integration.py | 2 +- 10 files changed, 72 insertions(+), 31 deletions(-) diff --git a/cratedb_toolkit/io/awslambda/kinesis.py b/cratedb_toolkit/io/awslambda/kinesis.py index f48286ec..d0632b01 100644 --- a/cratedb_toolkit/io/awslambda/kinesis.py +++ b/cratedb_toolkit/io/awslambda/kinesis.py @@ -69,6 +69,10 @@ message = f"Invalid value for MESSAGE_FORMAT: {MESSAGE_FORMAT}. Use one of: {message_formats}" logger.fatal(message) sys.exit(22) +if MESSAGE_FORMAT == "dynamodb" and not CRATEDB_TABLE: + message = "CRATEDB_TABLE environment variable is required when MESSAGE_FORMAT is 'dynamodb'" + logger.fatal(message) + sys.exit(22) try: column_types = ColumnTypeMapStore.from_json(COLUMN_TYPES) except Exception as ex: @@ -82,7 +86,8 @@ if MESSAGE_FORMAT == "dms": cdc = DMSTranslatorCrateDB(column_types=column_types) # ty: ignore[invalid-argument-type] elif MESSAGE_FORMAT == "dynamodb": - cdc = DynamoDBCDCTranslator(table_name=CRATEDB_TABLE) # ty: ignore[invalid-argument-type] + assert CRATEDB_TABLE is not None # Validated at module startup # noqa: S101 + cdc = DynamoDBCDCTranslator(table_name=CRATEDB_TABLE) # Create the database connection outside the handler to allow # connections to be re-used by subsequent function invocations. diff --git a/cratedb_toolkit/io/mongodb/cdc.py b/cratedb_toolkit/io/mongodb/cdc.py index 70dc3919..0d10ba70 100644 --- a/cratedb_toolkit/io/mongodb/cdc.py +++ b/cratedb_toolkit/io/mongodb/cdc.py @@ -62,6 +62,7 @@ def __init__( if tm: address = CollectionAddress( container=self.mongodb_adapter.database_name, + # TODO: Collection name must be permitted to be None? name=t.cast(str, self.mongodb_adapter.collection_name), ) try: diff --git a/cratedb_toolkit/io/mongodb/copy.py b/cratedb_toolkit/io/mongodb/copy.py index ea014348..5ff38b10 100644 --- a/cratedb_toolkit/io/mongodb/copy.py +++ b/cratedb_toolkit/io/mongodb/copy.py @@ -4,6 +4,7 @@ import sqlalchemy as sa from boltons.urlutils import URL +from commons_codec.model import SQLOperation from commons_codec.transform.mongodb import MongoDBCrateDBConverter, MongoDBFullLoadTranslator from tikray.model.collection import CollectionAddress from tqdm import tqdm @@ -13,6 +14,7 @@ from cratedb_toolkit.io.mongodb.adapter import mongodb_adapter_factory from cratedb_toolkit.io.mongodb.transform import TransformationManager from cratedb_toolkit.model import DatabaseAddress +from cratedb_toolkit.util.cli import to_list from cratedb_toolkit.util.database import DatabaseAdapter logger = logging.getLogger(__name__) @@ -51,6 +53,7 @@ def __init__( if tm: address = CollectionAddress( container=self.mongodb_adapter.database_name, + # TODO: Collection name must be permitted to be None? name=t.cast(str, self.mongodb_adapter.collection_name), ) try: @@ -88,10 +91,13 @@ def start(self): logger.info(f"Target: CrateDB table={self.cratedb_table} count={records_target}") progress_bar = tqdm(total=records_in) + def batch_to_sql_operation(batch: t.List[t.Dict[str, t.Any]]) -> SQLOperation: + return self.translator.to_sql(to_list(batch, [])) + processor = BulkProcessor( connection=connection, data=self.mongodb_adapter.query(), - batch_to_operation=self.translator.to_sql, # ty: ignore[invalid-argument-type] + batch_to_operation=batch_to_sql_operation, progress_bar=progress_bar, on_error=self.on_error, debug=self.debug, diff --git a/cratedb_toolkit/testing/testcontainers/cratedb.py b/cratedb_toolkit/testing/testcontainers/cratedb.py index 999d913d..a1e65c7e 100644 --- a/cratedb_toolkit/testing/testcontainers/cratedb.py +++ b/cratedb_toolkit/testing/testcontainers/cratedb.py @@ -153,7 +153,7 @@ def get_connection_url(self, dialect: str = "crate", host: Optional[str] = None) def _connect(self): if not self._wait_strategy: raise ValueError("No wait strategy defined") - wait_for_logs(self, predicate=self._wait_strategy) + wait_for_logs(self, predicate=self._wait_strategy, timeout=15) def _configure_wait_condition(self): """Wait for CrateDB node to be fully started.""" @@ -170,17 +170,23 @@ class CrateDBTestAdapter: """ def __init__(self, crate_version: str = "nightly", **kwargs): - self.cratedb: CrateDBContainer - self.database: DatabaseAdapter + self.cratedb: Optional[CrateDBContainer] = None + self._database: Optional[DatabaseAdapter] = None self.image: str = "crate/crate:{}".format(crate_version) + @property + def database(self) -> DatabaseAdapter: + if self._database is None: + raise ValueError("DatabaseAdapter is not initialized") + return self._database + def start(self, **kwargs): """ Start testcontainer, used for tests set up """ self.cratedb = CrateDBContainer(image=self.image, **kwargs) self.cratedb.start() - self.database = DatabaseAdapter(dburi=self.get_connection_url(), echo=False) + self._database = DatabaseAdapter(dburi=self.get_connection_url(), echo=False) def stop(self): """ @@ -193,19 +199,19 @@ def reset(self, tables: Optional[list] = None, schemas: Optional[list] = None): """ Drop tables from the given list, used for tests set up or tear down """ - if not self.database: + if not self._database: return if schemas: for reset_schema in schemas: - self.database.connection.exec_driver_sql( - f"DROP SCHEMA IF EXISTS {self.database.quote_relation_name(reset_schema)} CASCADE;" + self._database.connection.exec_driver_sql( + f"DROP SCHEMA IF EXISTS {self._database.quote_relation_name(reset_schema)} CASCADE;" ) if tables: for reset_table in tables: - self.database.connection.exec_driver_sql( - f"DROP TABLE IF EXISTS {self.database.quote_relation_name(reset_table)};" + self._database.connection.exec_driver_sql( + f"DROP TABLE IF EXISTS {self._database.quote_relation_name(reset_table)};" ) def get_connection_url(self, *args, **kwargs): diff --git a/cratedb_toolkit/testing/testcontainers/mongodb.py b/cratedb_toolkit/testing/testcontainers/mongodb.py index 82184be5..81a372fe 100644 --- a/cratedb_toolkit/testing/testcontainers/mongodb.py +++ b/cratedb_toolkit/testing/testcontainers/mongodb.py @@ -18,6 +18,7 @@ from pymongo import MongoClient from testcontainers.core.exceptions import ContainerStartException from testcontainers.mongodb import MongoDbContainer +from yarl import URL from cratedb_toolkit.testing.testcontainers.util import DockerSkippingContainer, KeepaliveContainer @@ -26,7 +27,7 @@ class MongoDbContainerWithKeepalive(DockerSkippingContainer, KeepaliveContainer, """ A Testcontainer for MongoDB with improved configurability. - It honors the `TC_KEEPALIVE` and `MONGODB_VERSION` environment variables. + It honours the `TC_KEEPALIVE` and `MONGODB_VERSION` environment variables. Defining `TC_KEEPALIVE` will set a signal not to shut down the container after running the test cases, in order to speed up subsequent invocations. @@ -87,6 +88,11 @@ def _create_connection_url( dbname: t.Optional[str] = None, **kwargs, ) -> str: + """ + TODO: Why does this method need to be overwritten? + Hint: Probably because of getting rid of credentials forwarding, + which is currently intended to reduce configuration complexity. + """ from testcontainers.core.utils import raise_for_deprecated_parameter if raise_for_deprecated_parameter(kwargs, "db_name", "dbname"): @@ -96,16 +102,20 @@ def _create_connection_url( host = host or self.get_container_host_ip() assert port is not None # noqa: S101 port = self.get_exposed_port(port) - url = f"{dialect}://{host}:{port}" + url = URL(f"{dialect}://{host}:{port}") + if username: + url = url.with_user(username) + if password: + url = url.with_password(password) if dbname: - url = f"{url}/{dbname}" - return url + url = url.with_path(dbname) + return str(url) def get_connection_url(self) -> str: return self._create_connection_url( dialect="mongodb", - username="admin", - password="", + username=None, # ty: ignore[invalid-argument-type] + password=None, # ty: ignore[invalid-argument-type] port=self.port, ) diff --git a/cratedb_toolkit/testing/testcontainers/util.py b/cratedb_toolkit/testing/testcontainers/util.py index b947d075..84781c64 100644 --- a/cratedb_toolkit/testing/testcontainers/util.py +++ b/cratedb_toolkit/testing/testcontainers/util.py @@ -19,6 +19,7 @@ from docker.errors import DockerException from docker.models.containers import Container from testcontainers.core.container import DockerContainer +from testcontainers.core.generic import DbContainer from cratedb_toolkit.util.data import asbool @@ -176,7 +177,7 @@ class PytestTestcontainerAdapter: """ def __init__(self): - self.container: DockerContainer + self.container: t.Optional[DbContainer] = None self.run_setup() @abstractmethod @@ -204,7 +205,11 @@ def run_setup(self): self.start() def start(self): + if self.container is None: + raise ValueError("Container not initialized") self.container.start() def stop(self): + if self.container is None: + raise ValueError("Container not initialized") self.container.stop() diff --git a/cratedb_toolkit/util/client.py b/cratedb_toolkit/util/client.py index 2e00ae9a..53f2d8f7 100644 --- a/cratedb_toolkit/util/client.py +++ b/cratedb_toolkit/util/client.py @@ -14,17 +14,18 @@ class jwt_token_patch(contextlib.ContextDecorator): def __init__(self, jwt_token: Optional[str] = None): self.jwt_token = jwt_token + self._patch_stack = contextlib.ExitStack() def __enter__(self): if self.jwt_token: - self.patcher = patch.object(crate.client.http.Client, "_request", _mk_crate_client_request(self.jwt_token)) - self.patcher.start() + self._patch_stack.enter_context( + patch.object(crate.client.http.Client, "_request", _mk_crate_client_request(self.jwt_token)) + ) return self def __exit__(self, type, value, tb): # noqa: A002 - if self.jwt_token: - self.patcher.stop() - return self + self._patch_stack.close() + return False def _mk_crate_client_request(jwt_token: Optional[str] = None): diff --git a/cratedb_toolkit/util/database.py b/cratedb_toolkit/util/database.py index 9eefd0d2..ae10a989 100644 --- a/cratedb_toolkit/util/database.py +++ b/cratedb_toolkit/util/database.py @@ -110,7 +110,7 @@ def run_sql( self, sql: t.Union[str, Path, io.IOBase], parameters: t.Optional[t.Mapping[str, str]] = None, - records: t.Optional[bool] = False, + records: bool = False, ignore: t.Optional[str] = None, ): """ @@ -136,9 +136,7 @@ def run_sql( raise return None - def run_sql_real( - self, sql: str, parameters: t.Optional[t.Mapping[str, str]] = None, records: t.Optional[bool] = False - ): + def run_sql_real(self, sql: str, parameters: t.Optional[t.Mapping[str, str]] = None, records: bool = False): """ Invoke an SQL statement and return results. """ @@ -368,7 +366,7 @@ def import_csv_dask( chunksize=1000, if_exists="replace", npartitions: t.Optional[int] = None, - progress: t.Optional[bool] = False, + progress: bool = False, ): """ Import CSV data using Dask. diff --git a/tests/io/mongodb/conftest.py b/tests/io/mongodb/conftest.py index 13d211d9..45254b15 100644 --- a/tests/io/mongodb/conftest.py +++ b/tests/io/mongodb/conftest.py @@ -37,8 +37,8 @@ def __init__(self, container_class): ) self.container_class = container_class - self.container: t.Union[MongoDbContainerWithKeepalive, MongoDbReplicasetContainer] - self.client: MongoClient + self.container: t.Optional[t.Union[MongoDbContainerWithKeepalive, MongoDbReplicasetContainer]] = None + self.client: t.Optional[MongoClient] = None super().__init__() def setup(self): @@ -54,19 +54,28 @@ def setup(self): self.client = self.container.get_connection_client() def finalize(self): + if not self.container: + return self.container.stop() def reset(self): """ Drop all databases used for testing. """ + if self.client is None: + return + for database_name in RESET_DATABASES: self.client.drop_database(database_name) def get_connection_url(self): + if not self.container: + raise RuntimeError("Container has not been initialized") return self.container.get_connection_url() def get_connection_client(self): + if not self.container: + raise RuntimeError("Container has not been initialized") return self.container.get_connection_client() diff --git a/tests/io/mongodb/test_integration.py b/tests/io/mongodb/test_integration.py index 7451800c..11ab3626 100644 --- a/tests/io/mongodb/test_integration.py +++ b/tests/io/mongodb/test_integration.py @@ -49,7 +49,7 @@ def setUpClass(cls): logger.debug(f"MongoDB server info: {server_info}") except pymongo.errors.ServerSelectionTimeoutError as ex: if cls.SKIP_IF_NOT_RUNNING: - raise cls.skipTest(cls, reason="MongoDB server not running") from ex # ty: ignore[invalid-argument-type] + raise unittest.SkipTest("MongoDB server not running") from ex else: # noqa: RET506 raise