QA: Migrate type checker from mypy to ty -- CodeRabbit improvements

Author: amotl

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.

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:

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,

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):

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,
         )
 

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()

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):

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.

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()
 
 

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