al-data-utils 0.3.1__py3-none-any.whl

al_data_utils/__init__.py

@@ -0,0 +1,10 @@
+__all__ = [
+    'airflow_utils',
+    'config_utils',
+    'geo_utils',
+    'logging_utils',
+    'azure_as_utils',
+    'sql_utils',
+    'tg_utils',
+    'ones_api_utils',
+]

al_data_utils/airflow_utils/__init__.py

@@ -0,0 +1,13 @@
+from .airflow_utils import (
+    get_variable_value,
+    get_variable_value_dict,
+    get_variable_value_int,
+    get_variable_value_str,
+)
+
+__all__ = [
+    "get_variable_value",
+    "get_variable_value_dict",
+    "get_variable_value_int",
+    "get_variable_value_str",
+]

al_data_utils/airflow_utils/airflow_utils.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import functools
+import json
+from pathlib import Path
+from sys import platform
+from typing import Any
+
+import yaml
+
+
+def _get_dags_folder() -> str:
+    try:
+        from airflow import settings  # type: ignore
+        return settings.DAGS_FOLDER
+    except Exception:
+        return "./dags"
+
+
+def _airflow_variable_get(name: str) -> str:
+    from airflow.sdk import Variable  # type: ignore
+    return Variable.get(name)
+
+
+@functools.cache
+def get_variable_value(
+    airflow_variable_name: str,
+    pack_to_json: bool = False,
+    source: str = "airflow",
+    source_section: str = "variables",
+) -> Any:
+    """Get variable value different ways for windows and linux."""
+
+    if platform == "win32":
+        base = Path(".local/variables")
+        data = json.loads((base / "variables.json").read_text(encoding="utf-8"))
+        environment_code = data["ENVIRONMENT"]
+
+        env_data = json.loads(
+            (base / f"{environment_code.lower()}_variables.json").read_text(encoding="utf-8")
+        )
+        variable_value = env_data[airflow_variable_name]
+
+    else:
+        if source == "config":
+            dags_folder = _get_dags_folder()
+            p = Path(dags_folder) / ".configs" / "_variables.yaml"
+            variables = yaml.safe_load(p.read_text(encoding="utf-8")) or {}
+            variable_value = variables[source_section][airflow_variable_name]
+        else:
+            variable_value = _airflow_variable_get(airflow_variable_name)
+
+    if pack_to_json and not isinstance(variable_value, dict):
+        try:
+            variable_value = json.loads(variable_value)
+        except Exception:
+            pass
+
+    return variable_value
+
+
+def get_variable_value_str(
+    airflow_variable_name: str,
+    source: str = "airflow",
+    source_section: str = "variables",
+) -> str:
+    return str(get_variable_value(airflow_variable_name, False, source, source_section))
+
+
+def get_variable_value_int(
+    airflow_variable_name: str,
+    source: str = "airflow",
+    source_section: str = "variables",
+) -> int:
+    return int(get_variable_value(airflow_variable_name, False, source, source_section))
+
+
+def get_variable_value_dict(
+    airflow_variable_name: str,
+    source: str = "airflow",
+    source_section: str = "variables",
+) -> dict:
+    return dict(get_variable_value(airflow_variable_name, True, source, source_section))

al_data_utils/azure_as_utils/__init__.py

@@ -0,0 +1,31 @@
+from .auth import ServicePrincipalTokenProvider, StaticTokenProvider, TokenProvider
+from .azure_as_utils import SemanticModels
+from .client import AasRestClient
+from .errors import (
+    AasAuthError,
+    AasError,
+    AasHttpError,
+    AasRefreshError,
+    AasRefreshTimeout,
+    AasTransportError,
+)
+from .models import RefreshObject, RefreshRequest, RefreshStatus
+from .poller import RefreshPoller
+
+__all__ = [
+    'AasAuthError',
+    'AasError',
+    'AasHttpError',
+    'AasRefreshError',
+    'AasRefreshTimeout',
+    'AasRestClient',
+    'AasTransportError',
+    'RefreshObject',
+    'RefreshPoller',
+    'RefreshRequest',
+    'RefreshStatus',
+    'SemanticModels',
+    'ServicePrincipalTokenProvider',
+    'StaticTokenProvider',
+    'TokenProvider',
+]

al_data_utils/azure_as_utils/auth.py

@@ -0,0 +1,82 @@
+"""Access tokens for the AAS REST API via an Azure AD service principal."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, Protocol
+
+import msal
+
+from .errors import AasAuthError
+
+#: Resource identifier of Azure Analysis Services; the wildcard is literal.
+AAS_SCOPE = 'https://*.asazure.windows.net/.default'
+DEFAULT_AUTHORITY_HOST = 'https://login.microsoftonline.com'
+
+
+class TokenProvider(Protocol):
+    """Anything that can hand out a bearer token for the AAS resource."""
+
+    def get_token(self) -> str: ...
+
+
+class StaticTokenProvider:
+    """Returns a fixed token. Useful for tests and pre-issued tokens."""
+
+    def __init__(self, token: str) -> None:
+        if not token:
+            raise ValueError('token must be a non-empty string')
+        self._token = token
+
+    def get_token(self) -> str:
+        return self._token
+
+
+class ServicePrincipalTokenProvider:
+    """Client-credentials flow through MSAL.
+
+    MSAL keeps the issued token in its in-memory cache and only goes back to
+    Azure AD when it is about to expire, so calling :meth:`get_token` before
+    every request is cheap.
+    """
+
+    def __init__(
+        self,
+        tenant_id: str,
+        client_id: str,
+        client_secret: str,
+        *,
+        authority_host: str = DEFAULT_AUTHORITY_HOST,
+        scope: str = AAS_SCOPE,
+    ) -> None:
+        for name, value in (
+            ('tenant_id', tenant_id),
+            ('client_id', client_id),
+            ('client_secret', client_secret),
+        ):
+            if not isinstance(value, str) or not value.strip():
+                raise ValueError(f'{name} must be a non-empty string')
+        self._scopes = [scope]
+        self._authority = f'{authority_host.rstrip("/")}/{tenant_id}'
+        self._app = msal.ConfidentialClientApplication(
+            client_id, authority=self._authority, client_credential=client_secret
+        )
+
+    def get_token(self) -> str:
+        result: Mapping[str, Any] | None = self._app.acquire_token_for_client(
+            scopes=self._scopes
+        )
+        token = result.get('access_token') if isinstance(result, Mapping) else None
+        if not token:
+            raise AasAuthError(_describe_failure(result))
+        return token
+
+
+def _describe_failure(result: Mapping[str, Any] | None) -> str:
+    if not isinstance(result, Mapping):
+        return f'token request returned no result ({result!r})'
+    error = result.get('error') or 'unknown_error'
+    description = result.get('error_description') or 'no description'
+    correlation = result.get('correlation_id')
+    suffix = f' (correlation_id={correlation})' if correlation else ''
+    return f'token request failed: {error}: {description}{suffix}'

al_data_utils/azure_as_utils/azure_as_utils.py

@@ -0,0 +1,223 @@
+"""Refresh Azure Analysis Services tabular models over the REST API.
+
+Typical use::
+
+    from al_data_utils.azure_as_utils import SemanticModels
+
+    sm = SemanticModels(
+        tenant_id=..., client_id=..., client_secret=...,
+        rollout='westeurope', server_name='myserver',
+    )
+    sm.refresh_models([
+        ('MyModel', {'Type': 'Full', 'CommitMode': 'transactional',
+                     'MaxParallelism': 2, 'RetryCount': 2, 'Objects': []}),
+    ])
+
+The class only glues together the smaller pieces in this package: a token
+provider (``auth``), the REST client (``client``) and the status poller
+(``poller``). Nothing here touches Airflow, databases or environment
+variables; the caller supplies every credential and setting.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable, Iterable, Mapping
+from typing import Any
+
+import requests
+
+import al_data_utils.logging_utils  # noqa: F401  (configures root logging)
+
+from .auth import ServicePrincipalTokenProvider, TokenProvider
+from .client import AasRestClient
+from .errors import AasHttpError, AasRefreshError, AasRefreshTimeout, AasTransportError
+from .models import RefreshRequest, RefreshStatus
+from .poller import (
+    DEFAULT_POLL_INTERVAL_SECONDS,
+    DEFAULT_TIMEOUT_SECONDS,
+    RefreshPoller,
+)
+
+log = logging.getLogger(__name__)
+
+ModelSpec = tuple[str, RefreshRequest | Mapping[str, Any]]
+
+#: Errors that mean "this refresh unit failed" rather than "the setup is broken".
+_UNIT_FAILURES = (AasRefreshError, AasRefreshTimeout, AasHttpError, AasTransportError)
+
+
+class SemanticModels:
+    """Refreshes one or more tabular models on a single AAS server.
+
+    Args:
+        tenant_id, client_id, client_secret: Azure AD service principal.
+        rollout: AAS region, e.g. ``'westeurope'``.
+        server_name: AAS server name (without the ``asazure://`` prefix).
+        process_by_partitions: when ``False`` every ``partition`` key is dropped
+            from the request objects and duplicate tables are collapsed, so the
+            refresh runs at table level.
+        debug_mode: when ``True`` each object of a request is refreshed in its
+            own call and gets its own log line. Failures do not stop the run;
+            one :class:`AasRefreshError` summarising them is raised at the end.
+        timeout_seconds: how long a single refresh may run before
+            :class:`AasRefreshTimeout` is raised.
+        poll_interval_seconds: initial delay between status checks; the poller
+            backs off from there.
+        request_timeout_seconds: socket timeout for each HTTP call.
+        token_provider, session: injection points, mainly for tests.
+    """
+
+    def __init__(
+        self,
+        tenant_id: str,
+        client_id: str,
+        client_secret: str,
+        rollout: str,
+        server_name: str,
+        process_by_partitions: bool = True,
+        debug_mode: bool = False,
+        *,
+        timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
+        poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS,
+        request_timeout_seconds: float = 60.0,
+        token_provider: TokenProvider | None = None,
+        session: requests.Session | None = None,
+        sleep: Callable[[float], None] = time.sleep,
+    ) -> None:
+        self.process_by_partitions = bool(process_by_partitions)
+        self.debug_mode = bool(debug_mode)
+        if token_provider is None:
+            token_provider = ServicePrincipalTokenProvider(
+                tenant_id, client_id, client_secret
+            )
+        self.client = AasRestClient(
+            rollout,
+            server_name,
+            token_provider,
+            session=session,
+            request_timeout_seconds=request_timeout_seconds,
+        )
+        self.poller = RefreshPoller(
+            self.client,
+            poll_interval_seconds=poll_interval_seconds,
+            timeout_seconds=timeout_seconds,
+            sleep=sleep,
+        )
+
+    # -- Public API ----------------------------------------------------------
+
+    def refresh_models(self, models: Iterable[ModelSpec]) -> None:
+        """Refresh every ``(model_name, refresh_request)`` pair in order.
+
+        Outside debug mode the first failure raises immediately. In debug mode
+        all units of all models are attempted and a single
+        :class:`AasRefreshError` is raised afterwards if any of them failed.
+        """
+        specs = list(models)
+        failures: list[str] = []
+        for model_name, raw_request in specs:
+            failures.extend(self._refresh_model_collecting(model_name, raw_request))
+        if failures:
+            raise AasRefreshError(
+                f'{len(failures)} refresh unit(s) failed: ' + '; '.join(failures),
+                messages=failures,
+            )
+
+    def refresh_model(
+        self, model_name: str, request: RefreshRequest | Mapping[str, Any]
+    ) -> list[RefreshStatus]:
+        """Refresh a single model and return the terminal status of each unit."""
+        results: list[RefreshStatus] = []
+        failures: list[str] = []
+        for unit in self._units(self.prepare_request(request)):
+            status = self._run_unit(model_name, unit, failures)
+            if status is not None:
+                results.append(status)
+        if failures:
+            raise AasRefreshError(
+                f'{len(failures)} refresh unit(s) of model {model_name!r} failed: '
+                + '; '.join(failures),
+                model_name=model_name,
+                messages=failures,
+            )
+        return results
+
+    def prepare_request(
+        self, request: RefreshRequest | Mapping[str, Any]
+    ) -> RefreshRequest:
+        """Apply the instance settings (partition handling) to a raw request."""
+        prepared = RefreshRequest.coerce(request)
+        if not self.process_by_partitions:
+            prepared = prepared.without_partitions()
+        return prepared
+
+    # -- Internals -----------------------------------------------------------
+
+    def _refresh_model_collecting(
+        self, model_name: str, raw_request: RefreshRequest | Mapping[str, Any]
+    ) -> list[str]:
+        failures: list[str] = []
+        for unit in self._units(self.prepare_request(raw_request)):
+            self._run_unit(model_name, unit, failures)
+        return failures
+
+    def _units(self, request: RefreshRequest) -> list[RefreshRequest]:
+        if not self.debug_mode:
+            return [request]
+        if request.is_whole_model:
+            log.warning(
+                'debug_mode is on but the request has no Objects; '
+                'refreshing the whole model in one call'
+            )
+        return request.per_object()
+
+    def _run_unit(
+        self, model_name: str, unit: RefreshRequest, failures: list[str]
+    ) -> RefreshStatus | None:
+        """Run one refresh; append to ``failures`` in debug mode, raise otherwise."""
+        label = unit.describe()
+        try:
+            status = self._start_and_wait(model_name, unit)
+        except _UNIT_FAILURES as exc:
+            if not self.debug_mode:
+                raise
+            log.error('model %s | %s | FAILED: %s', model_name, label, exc)
+            failures.append(f'{model_name} | {label}: {exc}')
+            return None
+        if self.debug_mode:
+            log.info('model %s | %s | %s', model_name, label, status.status)
+        return status
+
+    def _start_and_wait(self, model_name: str, unit: RefreshRequest) -> RefreshStatus:
+        log.info(
+            'model %s: starting %s refresh (%s) on %s/%s',
+            model_name,
+            unit.refresh_type,
+            unit.describe(),
+            self.client.host,
+            self.client.server_name,
+        )
+        refresh_id = self.client.start_refresh(model_name, unit)
+        log.info('model %s: refresh %s accepted', model_name, refresh_id)
+        status = self.poller.wait(model_name, refresh_id)
+        if not status.succeeded:
+            details = list(status.messages) or status.failed_objects()
+            detail_text = f': {"; ".join(details)}' if details else ''
+            raise AasRefreshError(
+                f'refresh {refresh_id} of model {model_name!r} ended with status '
+                f'{status.status!r}{detail_text}',
+                model_name=model_name,
+                refresh_id=refresh_id,
+                status=status.status,
+                messages=details,
+            )
+        log.info(
+            'model %s: refresh %s succeeded (%s -> %s)',
+            model_name,
+            refresh_id,
+            status.start_time,
+            status.end_time,
+        )
+        return status

al_data_utils/azure_as_utils/client.py

@@ -0,0 +1,158 @@
+"""Thin HTTP layer over the AAS refresh endpoints.
+
+Endpoints (see the Azure Analysis Services REST reference):
+
+* ``POST   /servers/{server}/models/{model}/refreshes``          start
+* ``GET    /servers/{server}/models/{model}/refreshes/{id}``     status
+* ``GET    /servers/{server}/models/{model}/refreshes``          history
+* ``DELETE /servers/{server}/models/{model}/refreshes/{id}``     cancel
+
+A successful start answers ``202 Accepted`` with a ``Location`` header that
+points at the status resource; the last path segment is the refresh id.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import quote, urlsplit
+
+import requests
+
+from .auth import TokenProvider
+from .errors import AasError, AasHttpError, AasTransportError
+from .models import RefreshRequest, RefreshStatus
+
+AAS_DOMAIN = 'asazure.windows.net'
+_BODY_EXCERPT_CHARS = 2000
+
+
+def rollout_host(rollout: str) -> str:
+    """Normalise ``'westeurope'`` or a full host / URL into the API host name."""
+    if not isinstance(rollout, str) or not rollout.strip():
+        raise ValueError('rollout must be a non-empty string')
+    value = rollout.strip().lower()
+    if '://' in value:
+        value = urlsplit(value).netloc or value
+    value = value.strip('/')
+    if '.' in value:
+        return value
+    return f'{value}.{AAS_DOMAIN}'
+
+
+class AasRestClient:
+    """Sends authenticated requests to one AAS server."""
+
+    def __init__(
+        self,
+        rollout: str,
+        server_name: str,
+        token_provider: TokenProvider,
+        *,
+        session: requests.Session | None = None,
+        request_timeout_seconds: float = 60.0,
+    ) -> None:
+        if not isinstance(server_name, str) or not server_name.strip():
+            raise ValueError('server_name must be a non-empty string')
+        if request_timeout_seconds <= 0:
+            raise ValueError('request_timeout_seconds must be positive')
+        self.host = rollout_host(rollout)
+        self.server_name = server_name.strip()
+        self._token_provider = token_provider
+        self._session = session or requests.Session()
+        self._timeout = request_timeout_seconds
+
+    # -- URL helpers ---------------------------------------------------------
+
+    @property
+    def server_url(self) -> str:
+        return f'https://{self.host}/servers/{quote(self.server_name, safe="")}'
+
+    def refreshes_url(self, model_name: str) -> str:
+        if not isinstance(model_name, str) or not model_name.strip():
+            raise ValueError('model_name must be a non-empty string')
+        return f'{self.server_url}/models/{quote(model_name, safe="")}/refreshes'
+
+    def refresh_url(self, model_name: str, refresh_id: str) -> str:
+        return f'{self.refreshes_url(model_name)}/{quote(refresh_id, safe="")}'
+
+    # -- Operations ----------------------------------------------------------
+
+    def start_refresh(self, model_name: str, request: RefreshRequest) -> str:
+        """Kick off a refresh and return its id."""
+        url = self.refreshes_url(model_name)
+        response = self._send('POST', url, json=request.to_payload())
+        refresh_id = _refresh_id_from_response(response)
+        if not refresh_id:
+            raise AasError(
+                f'POST {url} accepted the refresh but returned no refresh id '
+                f'(status {response.status_code}, headers {dict(response.headers)})'
+            )
+        return refresh_id
+
+    def get_refresh(self, model_name: str, refresh_id: str) -> RefreshStatus:
+        url = self.refresh_url(model_name, refresh_id)
+        response = self._send('GET', url)
+        try:
+            payload = response.json()
+        except ValueError as exc:
+            raise AasError(
+                f'GET {url} returned a non-JSON body: {_excerpt(response)}'
+            ) from exc
+        if not isinstance(payload, dict):
+            raise AasError(f'GET {url} returned unexpected JSON: {payload!r}')
+        return RefreshStatus.from_response(refresh_id, payload)
+
+    def list_refreshes(self, model_name: str) -> list[dict[str, Any]]:
+        response = self._send('GET', self.refreshes_url(model_name))
+        payload = response.json()
+        return (
+            [item for item in payload if isinstance(item, dict)]
+            if isinstance(payload, list)
+            else []
+        )
+
+    def cancel_refresh(self, model_name: str, refresh_id: str) -> None:
+        self._send('DELETE', self.refresh_url(model_name, refresh_id))
+
+    # -- Plumbing ------------------------------------------------------------
+
+    def _send(self, method: str, url: str, **kwargs: Any) -> requests.Response:
+        headers = {
+            'Authorization': f'Bearer {self._token_provider.get_token()}',
+            'Accept': 'application/json',
+        }
+        try:
+            response = self._session.request(
+                method, url, headers=headers, timeout=self._timeout, **kwargs
+            )
+        except requests.RequestException as exc:
+            raise AasTransportError(method, url, exc) from exc
+        if response.status_code >= 400:
+            raise AasHttpError(method, url, response.status_code, _excerpt(response))
+        return response
+
+
+def _excerpt(response: requests.Response) -> str:
+    text = response.text or ''
+    if len(text) > _BODY_EXCERPT_CHARS:
+        return text[:_BODY_EXCERPT_CHARS] + '...'
+    return text
+
+
+def _refresh_id_from_response(response: requests.Response) -> str | None:
+    location = response.headers.get('Location') or response.headers.get('location')
+    if location:
+        path = urlsplit(location).path or location
+        tail = path.rstrip('/').rsplit('/', 1)[-1]
+        if tail:
+            return tail
+    try:
+        body = response.json()
+    except ValueError:
+        return None
+    if isinstance(body, dict):
+        for key in ('operationId', 'refreshId', 'id'):
+            value = body.get(key)
+            if isinstance(value, str) and value:
+                return value
+    return None

al_data_utils/azure_as_utils/errors.py

@@ -0,0 +1,78 @@
+"""Exception hierarchy for the Azure Analysis Services helpers.
+
+Every error raised on purpose by this package derives from :class:`AasError`,
+so callers can catch one type when they do not care about the details.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+
+class AasError(Exception):
+    """Base class for all errors raised by ``azure_as_utils``."""
+
+
+class AasAuthError(AasError):
+    """Acquiring an access token for the AAS server failed."""
+
+
+class AasTransportError(AasError):
+    """The HTTP request never produced a response (DNS, TLS, timeout, ...)."""
+
+    def __init__(self, method: str, url: str, reason: BaseException) -> None:
+        self.method = method
+        self.url = url
+        self.reason = reason
+        super().__init__(f'{method} {url} failed: {reason!r}')
+
+
+class AasHttpError(AasError):
+    """The AAS REST endpoint answered with a non-success HTTP status."""
+
+    def __init__(self, method: str, url: str, status_code: int, body: str) -> None:
+        self.method = method
+        self.url = url
+        self.status_code = status_code
+        self.body = body
+        super().__init__(f'{method} {url} -> HTTP {status_code}: {body}')
+
+    @property
+    def is_transient(self) -> bool:
+        """True for throttling and server-side failures worth retrying."""
+        return self.status_code == 429 or self.status_code >= 500
+
+
+class AasRefreshError(AasError):
+    """A refresh ended in a state other than ``succeeded``."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        model_name: str | None = None,
+        refresh_id: str | None = None,
+        status: str | None = None,
+        messages: Sequence[str] = (),
+    ) -> None:
+        self.model_name = model_name
+        self.refresh_id = refresh_id
+        self.status = status
+        self.messages = tuple(messages)
+        super().__init__(message)
+
+
+class AasRefreshTimeout(AasError):
+    """A refresh was still running when the polling deadline passed."""
+
+    def __init__(
+        self, model_name: str, refresh_id: str, timeout_seconds: float, status: str
+    ) -> None:
+        self.model_name = model_name
+        self.refresh_id = refresh_id
+        self.timeout_seconds = timeout_seconds
+        self.status = status
+        super().__init__(
+            f'refresh {refresh_id} of model {model_name!r} still {status!r} after '
+            f'{timeout_seconds:.0f}s'
+        )