From 66807a2a4b4d09b4df9c8a77d627d7e6fabf4b52 Mon Sep 17 00:00:00 2001
From: Nick Proskurin <42863572+MCPN@users.noreply.github.com>
Date: Tue, 7 Nov 2023 16:25:43 +0100
Subject: [PATCH] Clean up some TODOs (#75)

---
 .../dl_connector_chyt/core/data_source.py                | 6 ++++--
 .../core/clickhouse/data_source.py                       | 2 +-
 .../core/clickhouse_base/ch_commons.py                   | 9 +++++++--
 .../core/clickhouse_base/data_source.py                  | 4 ++--
 4 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/lib/dl_connector_chyt/dl_connector_chyt/core/data_source.py b/lib/dl_connector_chyt/dl_connector_chyt/core/data_source.py
index 3dc63504d..3eae9d55a 100644
--- a/lib/dl_connector_chyt/dl_connector_chyt/core/data_source.py
+++ b/lib/dl_connector_chyt/dl_connector_chyt/core/data_source.py
@@ -96,14 +96,16 @@ def spec(self) -> CHYTTableDataSourceSpec:
 
     @property
     def default_title(self) -> str:
-        return self.spec.table_name.split("/")[-1]  # type: ignore  # TODO: fix
+        assert self.spec.table_name is not None
+        return self.spec.table_name.split("/")[-1]
 
     @require_table_name
     def get_sql_source(self, alias: Optional[str] = None) -> Any:
         if alias:
             return sa.alias(self.get_sql_source(), name=alias)
         path = self.spec.table_name
-        path = self.normalize_path(path)  # type: ignore  # TODO: fix
+        assert path is not None
+        path = self.normalize_path(path)
         return sa.table(path)
 
     @require_table_name
diff --git a/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/data_source.py b/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/data_source.py
index 2c67ed7f5..e6dcab660 100644
--- a/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/data_source.py
+++ b/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/data_source.py
@@ -25,7 +25,7 @@ def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
         }
 
 
-class ClickHouseSubselectDataSource(ActualClickHouseBaseMixin, CommonClickHouseSubselectDataSource):  # type: ignore  # TODO: fix
+class ClickHouseSubselectDataSource(ActualClickHouseBaseMixin, CommonClickHouseSubselectDataSource):
     """Clickhouse subselect"""
 
     conn_type = CONNECTION_TYPE_CLICKHOUSE
diff --git a/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/ch_commons.py b/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/ch_commons.py
index 126134198..b49c4f178 100644
--- a/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/ch_commons.py
+++ b/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/ch_commons.py
@@ -17,6 +17,7 @@
 
 from dl_core import exc
 from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
+from dl_core.connectors.base.error_transformer import DBExcKWArgs
 from dl_core.db import (
     SchemaColumn,
     make_sa_type,
@@ -187,11 +188,12 @@ def create_column_sql(
 ) -> str:
     native_type = tt.type_user_to_native(user_t=col.user_type, native_t=col.native_type)
 
+    nullable: Optional[bool]
     if partition_fields and col.name in partition_fields:
         # Partition column cannot be nullable. Enforcing it in here.
         nullable = False
     else:
-        nullable = None  # type: ignore  # TODO: fix
+        nullable = None
 
     if nullable is not None:
         native_type = native_type.as_common().clone(nullable=nullable)
@@ -204,7 +206,10 @@ def create_column_sql(
     return sa_ddl_obj.string
 
 
-def ensure_db_message(exc_cls, kw):  # type: ignore  # TODO: fix
+def ensure_db_message(
+    exc_cls: Type[exc.DatabaseQueryError],
+    kw: DBExcKWArgs,
+) -> tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
     db_message = kw.get("db_message")
     details = kw.get("details")
     if db_message and details is not None and not details.get("db_message"):
diff --git a/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/data_source.py b/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/data_source.py
index 9a30a53fe..aa805db3f 100644
--- a/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/data_source.py
+++ b/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/data_source.py
@@ -48,7 +48,7 @@ class ClickHouseBaseMixin(BaseSQLDataSource):
     compiler_cls = ClickHouseQueryCompiler
 
     def get_connect_args(self) -> dict:
-        return dict(super().get_connect_args(), server_version=self.db_version or self.default_server_version)  # type: ignore  # TODO: fix  # noqa
+        return dict(super().get_connect_args(), server_version=self.db_version or self.default_server_version)
 
 
 class ActualClickHouseBaseMixin(ClickHouseBaseMixin):
@@ -59,5 +59,5 @@ class ActualClickHouseBaseMixin(ClickHouseBaseMixin):
     conn_type = CONNECTION_TYPE_CLICKHOUSE
 
 
-class ClickHouseDataSourceBase(ActualClickHouseBaseMixin, StandardSQLDataSource, metaclass=abc.ABCMeta):  # type: ignore  # TODO: fix
+class ClickHouseDataSourceBase(ActualClickHouseBaseMixin, StandardSQLDataSource, metaclass=abc.ABCMeta):
     """ClickHouse table. Might not work correctly for views."""