analytics-toolkit 1.3.6.18__tar.gz → 1.3.7.3__tar.gz

{analytics_toolkit-1.3.6.18 → analytics_toolkit-1.3.7.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: analytics-toolkit
-Version: 1.3.6.18
+Version: 1.3.7.3
 Summary: Shared Python utilities for SQL, Excel, and date helpers.
 Author: analytics_toolkit contributors
 License: MIT License
@@ -75,7 +75,7 @@ Python toolkit for AB-test analysis, SQL workflows, Excel reports, and date help
 
 [Documentation Overview](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/README.md)
 
-**Version:** `1.3.6.18`<br>
+**Version:** `1.3.7.3`<br>
 **Depends:** Python (`>=3.8,<3.15`)<br>
 **Imports:** [clickhouse-connect](https://pypi.org/project/clickhouse-connect/) (`>=0.5.14,<1`), [lz4](https://pypi.org/project/lz4/) (`>=4.3.2,<5`), [numpy](https://pypi.org/project/numpy/) (`>=1.24.2,<2`), [openpyxl](https://pypi.org/project/openpyxl/) (`>=3.1.1,<4`), [orjson](https://pypi.org/project/orjson/) (`>=3.8.7,<4`), [pandas](https://pypi.org/project/pandas/) (`>=1.4.4,<3`), [psycopg2-binary](https://pypi.org/project/psycopg2-binary/) (`>=2.9.5,<3`), [python-dateutil](https://pypi.org/project/python-dateutil/) (`>=2.8.2,<3`), [pytz](https://pypi.org/project/pytz/) (`>=2022.7`), [requests](https://pypi.org/project/requests/) (`>=2.28.2,<3`), [scipy](https://pypi.org/project/scipy/) (`>=1.10.1,<2`), [sqlglot](https://pypi.org/project/sqlglot/) (`>=20,<31`), [sqlparse](https://pypi.org/project/sqlparse/) (`>=0.4.3,<1`), [tqdm](https://pypi.org/project/tqdm/) (`>=4.65.0,<5`), [trino](https://pypi.org/project/trino/) (`>=0.320,<1`), [zstandard](https://pypi.org/project/zstandard/) (`>=0.20.0,<1`)<br>
 **Suggests:** [apache-airflow](https://pypi.org/project/apache-airflow/) (`>=2.4,<3`; optional extra `airflow`)<br>

{analytics_toolkit-1.3.6.18 → analytics_toolkit-1.3.7.3}/README.md

@@ -6,7 +6,7 @@ Python toolkit for AB-test analysis, SQL workflows, Excel reports, and date help
 
 [Documentation Overview](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/README.md)
 
-**Version:** `1.3.6.18`<br>
+**Version:** `1.3.7.3`<br>
 **Depends:** Python (`>=3.8,<3.15`)<br>
 **Imports:** [clickhouse-connect](https://pypi.org/project/clickhouse-connect/) (`>=0.5.14,<1`), [lz4](https://pypi.org/project/lz4/) (`>=4.3.2,<5`), [numpy](https://pypi.org/project/numpy/) (`>=1.24.2,<2`), [openpyxl](https://pypi.org/project/openpyxl/) (`>=3.1.1,<4`), [orjson](https://pypi.org/project/orjson/) (`>=3.8.7,<4`), [pandas](https://pypi.org/project/pandas/) (`>=1.4.4,<3`), [psycopg2-binary](https://pypi.org/project/psycopg2-binary/) (`>=2.9.5,<3`), [python-dateutil](https://pypi.org/project/python-dateutil/) (`>=2.8.2,<3`), [pytz](https://pypi.org/project/pytz/) (`>=2022.7`), [requests](https://pypi.org/project/requests/) (`>=2.28.2,<3`), [scipy](https://pypi.org/project/scipy/) (`>=1.10.1,<2`), [sqlglot](https://pypi.org/project/sqlglot/) (`>=20,<31`), [sqlparse](https://pypi.org/project/sqlparse/) (`>=0.4.3,<1`), [tqdm](https://pypi.org/project/tqdm/) (`>=4.65.0,<5`), [trino](https://pypi.org/project/trino/) (`>=0.320,<1`), [zstandard](https://pypi.org/project/zstandard/) (`>=0.20.0,<1`)<br>
 **Suggests:** [apache-airflow](https://pypi.org/project/apache-airflow/) (`>=2.4,<3`; optional extra `airflow`)<br>

{analytics_toolkit-1.3.6.18 → analytics_toolkit-1.3.7.3}/analytics_toolkit/sql/connection/config.py

@@ -56,8 +56,7 @@ class TrinoConfig:
     auth_mode: str
     http_scheme: str
     verify_value: str
-    use_keychain_certs: bool
-    keychain_cert_names: list[str]
+    ca_certs: list[str]
     insert_chunk_size: int | None
     request_timeout: int | None
     source: str | None
@@ -77,6 +76,10 @@ class GpConfig:
     keepalives_idle: int
     keepalives_interval: int
     keepalives_count: int
+    sslmode: str | None
+    ca_certs: list[str]
+    ssl_cert: str | None
+    ssl_key: str | None
 
 
 @dataclass(frozen=True)
@@ -90,8 +93,8 @@ class ChConfig:
     database: str | None
     secure: bool
     verify_value: str | None
-    ca_cert: str | None
-    ca_cert_variable: str | None
+    ca_certs: list[str]
+    ca_certs_variable: str | None
     connect_timeout: int | None
     send_receive_timeout: int | None
     settings: dict[str, Any] | None
@@ -174,6 +177,11 @@ def _build_connection_config(
     backend = _require_backend(connection_key, raw_config)
 
     if backend == "trino":
+        _reject_removed_fields(
+            raw_config,
+            connection_key,
+            ["use_keychain_certs", "keychain_cert_names", "ca_cert"],
+        )
         return TrinoConfig(
             connection_key=connection_key,
             backend=backend,
@@ -203,16 +211,10 @@ def _build_connection_config(
                 )
                 or "true"
             ),
-            use_keychain_certs=_optional_bool(
+            ca_certs=_optional_string_or_string_list(
                 raw_config,
                 connection_key,
-                "use_keychain_certs",
-                False,
-            ),
-            keychain_cert_names=_optional_string_list(
-                raw_config,
-                connection_key,
-                "keychain_cert_names",
+                "ca_certs",
             ),
             insert_chunk_size=_optional_positive_int(
                 raw_config,
@@ -227,6 +229,15 @@ def _build_connection_config(
             source=_optional_string(raw_config, connection_key, "source"),
         )
     if backend == "gp":
+        _reject_removed_fields(raw_config, connection_key, ["ca_cert"])
+        ca_certs = _optional_string_or_string_list(
+            raw_config,
+            connection_key,
+            "ca_certs",
+        )
+        sslmode = _optional_string(raw_config, connection_key, "sslmode")
+        if ca_certs and sslmode is None:
+            sslmode = "verify-full"
         return GpConfig(
             connection_key=connection_key,
             backend=backend,
@@ -265,24 +276,29 @@ def _build_connection_config(
                 "keepalives_count",
                 DEFAULT_GP_KEEPALIVES_COUNT,
             ),
+            sslmode=sslmode,
+            ca_certs=ca_certs,
+            ssl_cert=_optional_string(raw_config, connection_key, "ssl_cert"),
+            ssl_key=_optional_string(raw_config, connection_key, "ssl_key"),
         )
     if backend == "ch":
-        ca_cert = _optional_string(raw_config, connection_key, "ca_cert")
-        if ca_cert is None:
-            ca_cert = _optional_string(raw_config, connection_key, "ca_certs")
-        ca_cert_variable = _optional_string(
+        _reject_removed_fields(
             raw_config,
             connection_key,
-            "ca_cert_variable",
+            ["ca_cert", "ca_cert_variable"],
         )
-        if ca_cert_variable is None:
-            ca_cert_variable = _optional_string(
-                raw_config,
-                connection_key,
-                "ca_certs_variable",
-            )
-        if ca_cert is not None:
-            ca_cert_variable = None
+        ca_certs = _optional_string_or_string_list(
+            raw_config,
+            connection_key,
+            "ca_certs",
+        )
+        ca_certs_variable = _optional_string(
+            raw_config,
+            connection_key,
+            "ca_certs_variable",
+        )
+        if ca_certs:
+            ca_certs_variable = None
         return ChConfig(
             connection_key=connection_key,
             backend=backend,
@@ -297,8 +313,8 @@ def _build_connection_config(
                 connection_key,
                 "verify",
             ),
-            ca_cert=ca_cert,
-            ca_cert_variable=ca_cert_variable,
+            ca_certs=ca_certs,
+            ca_certs_variable=ca_certs_variable,
             connect_timeout=_optional_positive_int(
                 raw_config,
                 connection_key,
@@ -338,19 +354,35 @@ def generate_dummy_connections(airflow: bool = False) -> Path:
     connections_path = Path.cwd() / CONNECTIONS_FILE_NAME
     if connections_path.exists():
         raise ValueError(f"SQL connections file already exists: {connections_path}")
+    certs_dir = Path.cwd() / ".certs"
 
     content = (
         _build_dummy_airflow_connections()
         if airflow
         else _build_dummy_direct_connections()
     )
+    certs_dir.mkdir(exist_ok=True)
     connections_path.write_text(
         json.dumps(content, indent=2) + "\n",
         encoding="utf-8",
     )
+    _print_dummy_connections_cert_instructions(certs_dir)
     return connections_path
 
 
+def _print_dummy_connections_cert_instructions(certs_dir: Path) -> None:
+    print(f"Created {certs_dir} for local certificate files.")
+    print(
+        "Greenplum: put CA PEMs in .certs/ and set ca_certs; "
+        "optional ssl_cert, ssl_key, sslmode."
+    )
+    print("Trino: put HTTPS CA PEMs in .certs/ and set ca_certs.")
+    print(
+        "ClickHouse: put HTTPS CA PEMs in .certs/ and set ca_certs; "
+        "Airflow can use ca_certs_variable."
+    )
+
+
 def _build_dummy_direct_connections() -> dict[str, dict[str, Any]]:
     return {
         "gp": {
@@ -360,6 +392,7 @@ def _build_dummy_direct_connections() -> dict[str, dict[str, Any]]:
             "user": "user",
             "password": "password",
             "database": "db",
+            "ca_certs": "gp-ca.pem",
         },
         "trino": {
             "type": "trino",
@@ -369,6 +402,8 @@ def _build_dummy_direct_connections() -> dict[str, dict[str, Any]]:
             "password": "password",
             "catalog": "iceberg",
             "schema": "sandbox",
+            "http_scheme": "https",
+            "ca_certs": "trino-ca.pem",
         },
         "ch": {
             "type": "ch",
@@ -377,6 +412,8 @@ def _build_dummy_direct_connections() -> dict[str, dict[str, Any]]:
             "user": "user",
             "password": "password",
             "database": "default",
+            "secure": True,
+            "ca_certs": "clickhouse-ca.pem",
         },
     }
 
@@ -385,9 +422,9 @@ def _build_dummy_airflow_connections() -> dict[str, Any]:
     return {
         "source": "airflow",
         "connections": {
-            "gp": {"type": "gp"},
-            "trino": {"type": "trino"},
-            "ch": {"type": "ch"},
+            "gp": {"type": "gp", "ca_certs": "gp-ca.pem"},
+            "trino": {"type": "trino", "ca_certs": "trino-ca.pem"},
+            "ch": {"type": "ch", "ca_certs": "clickhouse-ca.pem"},
         },
     }
 
@@ -763,6 +800,10 @@ def _get_airflow_raw_connection_config_and_extras(
                 "keepalives_idle",
                 "keepalives_interval",
                 "keepalives_count",
+                "sslmode",
+                "ca_certs",
+                "ssl_cert",
+                "ssl_key",
             ],
         )
     elif resolved_backend == "trino":
@@ -775,8 +816,7 @@ def _get_airflow_raw_connection_config_and_extras(
                 "auth_mode",
                 "http_scheme",
                 "verify",
-                "use_keychain_certs",
-                "keychain_cert_names",
+                "ca_certs",
                 "insert_chunk_size",
                 "request_timeout",
                 "source",
@@ -794,9 +834,7 @@ def _get_airflow_raw_connection_config_and_extras(
             [
                 "secure",
                 "verify",
-                "ca_cert",
                 "ca_certs",
-                "ca_cert_variable",
                 "ca_certs_variable",
                 "connect_timeout",
                 "send_receive_timeout",
@@ -1208,7 +1246,21 @@ def _optional_mapping(
     return dict(value)
 
 
-def _optional_string_list(
+def _reject_removed_fields(
+    config: dict[str, Any],
+    connection_key: str,
+    field_names: Sequence[str],
+) -> None:
+    for field_name in field_names:
+        if field_name in config:
+            raise SqlConfigError(
+                f"SQL connection '{connection_key}' field '{field_name}' "
+                "is not supported. "
+                "Use 'ca_certs' for CA certificate files."
+            )
+
+
+def _optional_string_or_string_list(
     config: dict[str, Any],
     connection_key: str,
     field_name: str,
@@ -1217,7 +1269,7 @@ def _optional_string_list(
     if value is None:
         return []
     if isinstance(value, str):
-        names = [name.strip() for name in value.split("|")]
+        names = [value.strip()]
     elif isinstance(value, list):
         names = []
         for item in value:

{analytics_toolkit-1.3.6.18 → analytics_toolkit-1.3.7.3}/analytics_toolkit/sql/connection/get_sql_connection.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import os
-import subprocess
 from dataclasses import replace
 from functools import wraps
 from pathlib import Path
@@ -91,9 +89,8 @@ def _get_trino_connection(config: TrinoConfig) -> Any:
             "The 'trino' package is required for Trino connections."
         ) from exc
 
-    verify_value = config.verify_value
-    if config.use_keychain_certs:
-        verify_value = str(_build_trino_keychain_bundle(config))
+    ca_certs = _resolve_ca_certs(config.connection_key, config.ca_certs)
+    verify_value = ca_certs or config.verify_value
 
     if config.auth_mode == "oauth2":
         auth = trino.auth.OAuth2Authentication()
@@ -150,6 +147,25 @@ def _get_gp_connection(config: GpConfig) -> Any:
                 "keepalives_count": config.keepalives_count,
             }
         )
+    if config.sslmode is not None:
+        connect_kwargs["sslmode"] = config.sslmode
+    ca_certs = _resolve_ca_certs(config.connection_key, config.ca_certs)
+    if ca_certs is not None:
+        connect_kwargs["sslrootcert"] = ca_certs
+    if config.ssl_cert is not None:
+        ssl_cert = _resolve_single_cert_path(
+            config.connection_key,
+            config.ssl_cert,
+            field_name="ssl_cert",
+        )
+        connect_kwargs["sslcert"] = str(ssl_cert)
+    if config.ssl_key is not None:
+        ssl_key = _resolve_single_cert_path(
+            config.connection_key,
+            config.ssl_key,
+            field_name="ssl_key",
+        )
+        connect_kwargs["sslkey"] = str(ssl_key)
 
     return psycopg2.connect(**connect_kwargs)
 
@@ -176,9 +192,9 @@ def _get_ch_connection(config: ChConfig) -> Any:
         client_kwargs["database"] = config.database
     if config.verify_value is not None:
         client_kwargs["verify"] = _parse_verify_value(config.verify_value)
-    ca_cert = _resolve_ch_ca_cert(config)
-    if ca_cert is not None:
-        client_kwargs["ca_cert"] = ca_cert
+    ca_certs = _resolve_ch_ca_certs(config)
+    if ca_certs is not None:
+        client_kwargs["ca_cert"] = ca_certs
     if config.connect_timeout is not None:
         client_kwargs["connect_timeout"] = config.connect_timeout
     if config.send_receive_timeout is not None:
@@ -197,35 +213,35 @@ def _get_ch_connection(config: ChConfig) -> Any:
     return clickhouse_connect.get_client(**client_kwargs)
 
 
-def _resolve_ch_ca_cert(config: ChConfig) -> str | None:
-    if config.ca_cert:
-        return config.ca_cert
-    if config.ca_cert_variable is None:
+def _resolve_ch_ca_certs(config: ChConfig) -> str | None:
+    if config.ca_certs:
+        return _resolve_ca_certs(config.connection_key, config.ca_certs)
+    if config.ca_certs_variable is None:
         return None
 
     try:
         from airflow.models.variable import Variable
     except ImportError as exc:
         raise SqlConfigError(
-            f"SQL connection '{config.connection_key}' uses ca_cert_variable "
+            f"SQL connection '{config.connection_key}' uses ca_certs_variable "
             "but Airflow Variable support is unavailable."
         ) from exc
 
     try:
-        ca_cert = Variable.get(config.ca_cert_variable)
+        ca_certs = Variable.get(config.ca_certs_variable)
     except Exception as exc:
         raise SqlConfigError(
-            f"Could not resolve Airflow Variable '{config.ca_cert_variable}' "
+            f"Could not resolve Airflow Variable '{config.ca_certs_variable}' "
             f"for SQL connection '{config.connection_key}'. "
             f"{type(exc).__name__}: {exc}"
         ) from exc
 
-    if not isinstance(ca_cert, str) or not ca_cert.strip():
+    if not isinstance(ca_certs, str) or not ca_certs.strip():
         raise SqlConfigError(
-            f"Airflow Variable '{config.ca_cert_variable}' for SQL connection "
+            f"Airflow Variable '{config.ca_certs_variable}' for SQL connection "
             f"'{config.connection_key}' must be a non-empty string."
         )
-    return ca_cert.strip()
+    return _resolve_ca_certs(config.connection_key, [ca_certs.strip()])
 
 
 def _parse_verify_value(value: str) -> bool | str:
@@ -238,69 +254,56 @@ def _parse_verify_value(value: str) -> bool | str:
     return normalized
 
 
-def _build_trino_keychain_bundle(config: TrinoConfig) -> Path:
-    if not config.keychain_cert_names:
-        raise SqlConfigError(
-            f"SQL connection '{config.connection_key}' enables keychain certs "
-            "but does not define keychain_cert_names."
-        )
+def _resolve_ca_certs(connection_key: str, ca_certs: list[str]) -> str | None:
+    if not ca_certs:
+        return None
 
-    certs_dir = _state_dir() / "certs"
-    certs_dir.mkdir(parents=True, exist_ok=True)
-    bundle_path = certs_dir / f"trino-{_safe_file_key(config.connection_key)}-keychain-ca.pem"
-    keychains = [
-        str(Path.home() / "Library/Keychains/login.keychain-db"),
-        "/Library/Keychains/System.keychain",
+    resolved_paths = [
+        _resolve_single_cert_path(connection_key, cert_path, field_name="ca_certs")
+        for cert_path in ca_certs
     ]
-
-    certificates: list[str] = []
-    for cert_name in config.keychain_cert_names:
-        certificate = _export_keychain_certificate(cert_name, keychains)
-        if not certificate:
-            raise SqlConfigError(
-                f"Could not export '{cert_name}' from macOS Keychain."
-            )
-        certificates.append(certificate.strip())
-    bundle_contents = "\n".join(certificates) + "\n"
-    if not bundle_path.exists() or bundle_path.read_text(encoding="utf-8") != bundle_contents:
+    if len(resolved_paths) == 1:
+        return str(resolved_paths[0])
+
+    connections_dir = get_connections_file_path().parent
+    bundle_dir = connections_dir / ".certs" / ".generated"
+    bundle_dir.mkdir(parents=True, exist_ok=True)
+    bundle_path = bundle_dir / f"{_safe_file_key(connection_key)}-ca-bundle.pem"
+    bundle_contents = "\n".join(
+        path.read_text(encoding="utf-8").strip()
+        for path in resolved_paths
+    ) + "\n"
+    if (
+        not bundle_path.exists()
+        or bundle_path.read_text(encoding="utf-8") != bundle_contents
+    ):
         bundle_path.write_text(bundle_contents, encoding="utf-8")
-    return bundle_path
-
-
-def _export_keychain_certificate(cert_name: str, keychains: list[str]) -> str:
-    for keychain in keychains:
-        if not Path(keychain).exists():
-            continue
-
-        result = subprocess.run(
-            [
-                "security",
-                "find-certificate",
-                "-a",
-                "-c",
-                cert_name,
-                "-p",
-                keychain,
-            ],
-            capture_output=True,
-            text=True,
-            check=False,
+    return str(bundle_path)
+
+
+def _resolve_single_cert_path(
+    connection_key: str,
+    cert_path: str,
+    *,
+    field_name: str,
+) -> Path:
+    raw_path = cert_path.strip()
+    path = Path(raw_path)
+    if path.is_absolute():
+        resolved = path
+    else:
+        connections_dir = get_connections_file_path().parent
+        if "/" in raw_path or "\\" in raw_path:
+            resolved = connections_dir / path
+        else:
+            resolved = connections_dir / ".certs" / path
+    resolved = resolved.resolve()
+    if not resolved.is_file():
+        raise SqlConfigError(
+            f"SQL connection '{connection_key}' field '{field_name}' references "
+            f"missing certificate file: {resolved}"
         )
-        if result.returncode == 0 and "BEGIN CERTIFICATE" in result.stdout:
-            return result.stdout
-
-    return ""
-
-
-def _state_dir() -> Path:
-    state_override = os.getenv("MAGNIT_UTILS_HOME")
-    if state_override:
-        return Path(state_override).expanduser()
-
-    try:
-        return get_connections_file_path().parent
-    except SqlConfigError:
-        return Path.cwd().resolve()
+    return resolved
 
 
 def _safe_file_key(connection_key: str) -> str:

{analytics_toolkit-1.3.6.18 → analytics_toolkit-1.3.7.3}/analytics_toolkit.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: analytics-toolkit
-Version: 1.3.6.18
+Version: 1.3.7.3
 Summary: Shared Python utilities for SQL, Excel, and date helpers.
 Author: analytics_toolkit contributors
 License: MIT License
@@ -75,7 +75,7 @@ Python toolkit for AB-test analysis, SQL workflows, Excel reports, and date help
 
 [Documentation Overview](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/README.md)
 
-**Version:** `1.3.6.18`<br>
+**Version:** `1.3.7.3`<br>
 **Depends:** Python (`>=3.8,<3.15`)<br>
 **Imports:** [clickhouse-connect](https://pypi.org/project/clickhouse-connect/) (`>=0.5.14,<1`), [lz4](https://pypi.org/project/lz4/) (`>=4.3.2,<5`), [numpy](https://pypi.org/project/numpy/) (`>=1.24.2,<2`), [openpyxl](https://pypi.org/project/openpyxl/) (`>=3.1.1,<4`), [orjson](https://pypi.org/project/orjson/) (`>=3.8.7,<4`), [pandas](https://pypi.org/project/pandas/) (`>=1.4.4,<3`), [psycopg2-binary](https://pypi.org/project/psycopg2-binary/) (`>=2.9.5,<3`), [python-dateutil](https://pypi.org/project/python-dateutil/) (`>=2.8.2,<3`), [pytz](https://pypi.org/project/pytz/) (`>=2022.7`), [requests](https://pypi.org/project/requests/) (`>=2.28.2,<3`), [scipy](https://pypi.org/project/scipy/) (`>=1.10.1,<2`), [sqlglot](https://pypi.org/project/sqlglot/) (`>=20,<31`), [sqlparse](https://pypi.org/project/sqlparse/) (`>=0.4.3,<1`), [tqdm](https://pypi.org/project/tqdm/) (`>=4.65.0,<5`), [trino](https://pypi.org/project/trino/) (`>=0.320,<1`), [zstandard](https://pypi.org/project/zstandard/) (`>=0.20.0,<1`)<br>
 **Suggests:** [apache-airflow](https://pypi.org/project/apache-airflow/) (`>=2.4,<3`; optional extra `airflow`)<br>

{analytics_toolkit-1.3.6.18 → analytics_toolkit-1.3.7.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "analytics-toolkit"
-version = "1.3.6.18"
+version = "1.3.7.3"
 description = "Shared Python utilities for SQL, Excel, and date helpers."
 readme = "README.md"
 requires-python = ">=3.8,<3.15"