[DOP-13849] - put away Version annotations

MobileTeleSystems · Apr 16, 2024 · 2c84b33 · 2c84b33
1 parent 461a139
commit 2c84b33
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 21 deletions.
diff --git a/onetl/connection/db_connection/clickhouse/connection.py b/onetl/connection/db_connection/clickhouse/connection.py
@@ -110,18 +110,18 @@ class Clickhouse(JDBCConnection):
     @classmethod
     def get_packages(
         cls,
-        package_version: str | Version | None = None,
-        apache_http_client_version: str | Version | None = None,
+        package_version: str | None = None,
+        apache_http_client_version: str | None = None,
     ) -> list[str]:
         """
         Get package names to be downloaded by Spark. Allows specifying custom JDBC and Apache HTTP Client versions. |support_hooks|
 
         Parameters
         ----------
-        package_version : str | Version, optional
+        package_version : str, optional
              ClickHouse JDBC version client packages. Defaults to ``0.6.0``.
 
-        apache_http_client_version : str | Version, optional
+        apache_http_client_version : str, optional
              Apache HTTP Client version package. Defaults to ``5.3.1``.
 
         Examples
@@ -142,20 +142,16 @@ def get_packages(
         default_jdbc_version = "0.6.0"
         default_http_version = "5.3.1"
 
-        jdbc_version = default_jdbc_version if package_version is None else str(Version(package_version).min_digits(3))
-        http_version = (
-            default_http_version
-            if apache_http_client_version is None
-            else str(Version(apache_http_client_version).min_digits(3))
-        )
+        jdbc_version_obj = Version(package_version or default_jdbc_version).min_digits(3)
+        http_version_obj = Version(apache_http_client_version or default_http_version).min_digits(3)
 
         result = [
-            f"com.clickhouse:clickhouse-jdbc:{jdbc_version}",
-            f"com.clickhouse:clickhouse-http-client:{jdbc_version}",
+            f"com.clickhouse:clickhouse-jdbc:{jdbc_version_obj}",
+            f"com.clickhouse:clickhouse-http-client:{jdbc_version_obj}",
         ]
 
-        if Version(jdbc_version) >= Version("0.5.0"):
-            result.append(f"org.apache.httpcomponents.client5:httpclient5:{http_version}")
+        if jdbc_version_obj >= Version("0.5.0"):
+            result.append(f"org.apache.httpcomponents.client5:httpclient5:{http_version_obj}")
 
         return result
 

diff --git a/onetl/connection/db_connection/postgres/connection.py b/onetl/connection/db_connection/postgres/connection.py
@@ -106,13 +106,13 @@ class Postgres(JDBCConnection):
 
     @slot
     @classmethod
-    def get_packages(cls, package_version: str | Version | None = None) -> list[str]:
+    def get_packages(cls, package_version: str | None = None) -> list[str]:
         """
         Get package names to be downloaded by Spark.  Allows specifying a custom JDBC driver version. |support_hooks|
 
         Parameters
         ----------
-        package_version : str | Version, optional
+        package_version : str, optional
             Specifies the version of the PostgreSQL JDBC driver to use.  Defaults to ``42.7.3``.
 
         Examples
@@ -128,11 +128,10 @@ def get_packages(cls, package_version: str | Version | None = None) -> list[str]
             Postgres.get_packages(package_version="42.6.0")
 
         """
-
         default_version = "42.7.3"
-        version = default_version if package_version is None else str(Version(package_version).min_digits(3))
+        version_obj = Version(package_version or default_version).min_digits(3)
 
-        return [f"org.postgresql:postgresql:{version}"]
+        return [f"org.postgresql:postgresql:{version_obj}"]
 
     @classproperty
     def package(cls) -> str:

diff --git a/tests/tests_unit/tests_db_connection_unit/test_clickhouse_unit.py b/tests/tests_unit/tests_db_connection_unit/test_clickhouse_unit.py
@@ -94,7 +94,7 @@ def test_clickhouse_get_packages(package_version, apache_http_client_version, ex
         ("a.b.c", "5.3.1"),
     ],
 )
-def test_clickhouse_invalid_version(package_version, apache_http_client_version):
+def test_clickhouse_get_packages_invalid_version(package_version, apache_http_client_version):
     with pytest.raises(
         ValueError,
         match=rf"Version '{package_version}' does not have enough numeric components for requested format \(expected at least 3\).",

diff --git a/tests/tests_unit/tests_db_connection_unit/test_postgres_unit.py b/tests/tests_unit/tests_db_connection_unit/test_postgres_unit.py
@@ -39,7 +39,7 @@ def test_postgres_get_packages(package_version, expected_packages):
         "abc",
     ],
 )
-def test_postgres_invalid_version(package_version):
+def test_postgres_get_packages_invalid_version(package_version):
     with pytest.raises(
         ValueError,
         match=rf"Version '{package_version}' does not have enough numeric components for requested format \(expected at least 3\).",