From 2c84b33092b41d225850b40c485644a707e83bac Mon Sep 17 00:00:00 2001 From: maxim-lixakov Date: Tue, 16 Apr 2024 14:11:23 +0300 Subject: [PATCH] [DOP-13849] - put away Version annotations --- .../db_connection/clickhouse/connection.py | 24 ++++++++----------- .../db_connection/postgres/connection.py | 9 ++++--- .../test_clickhouse_unit.py | 2 +- .../test_postgres_unit.py | 2 +- 4 files changed, 16 insertions(+), 21 deletions(-) diff --git a/onetl/connection/db_connection/clickhouse/connection.py b/onetl/connection/db_connection/clickhouse/connection.py index 658df3238..749b7677b 100644 --- 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 index 59d67916b..5c82ad038 100644 --- 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 index 62d30e6b4..30154daca 100644 --- 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 index 6f63d4d11..87045ee87 100644 --- 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\).",