loaders.py

"""
Module containing functions and classes for loading configuration files.
"""
#  Copyright 2023 Cognite AS
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import argparse
import dataclasses
import json
import logging
import os
import re
import sys
from collections.abc import Callable, Iterable
from enum import Enum
from hashlib import sha256
from pathlib import Path
from typing import Any, Generic, TextIO, TypeVar, cast

import dacite
import yaml
from azure.core.credentials import TokenCredential
from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError
from azure.identity import ClientSecretCredential, DefaultAzureCredential
from azure.keyvault.secrets import SecretClient
from yaml.scanner import ScannerError

from cognite.client import CogniteClient
from cognite.extractorutils.configtools._util import _to_snake_case
from cognite.extractorutils.configtools.elements import (
    BaseConfig,
    CastableInt,
    ConfigType,
    IgnorePattern,
    PortNumber,
    TimeIntervalConfig,
    _BaseConfig,
)
from cognite.extractorutils.exceptions import InvalidConfigError

_logger = logging.getLogger(__name__)


CustomConfigClass = TypeVar("CustomConfigClass", bound=BaseConfig)


class KeyVaultAuthenticationMethod(Enum):
    """
    Enum representing the authentication methods for Azure KeyVault.
    """

    DEFAULT = "default"
    CLIENTSECRET = "client-secret"


class KeyVaultLoader:
    """
    Class responsible for configuring keyvault for clients using Azure.

    Args:
        config: A dictionary containing the configuration for the keyvault.
    """

    def __init__(self, config: dict | None) -> None:
        self.config = config

        self.client: SecretClient | None = None

    def _init_client(self) -> None:
        from dotenv import load_dotenv

        if not self.config:
            raise InvalidConfigError(
                "Attempted to load values from Azure key vault with no key vault configured. "
                "Include an `azure-keyvault` section in your config to use the !keyvault tag."
            )

        keyvault_name = self.config.get("keyvault-name")
        if not keyvault_name:
            raise InvalidConfigError("Please add the keyvault-name")

        if "authentication-method" not in self.config:
            raise InvalidConfigError(
                "Please enter the authentication method to access Azure KeyVault"
                "Possible values are: default or client-secret"
            )

        vault_url = f"https://{keyvault_name}.vault.azure.net"

        credentials: TokenCredential
        if self.config["authentication-method"] == KeyVaultAuthenticationMethod.DEFAULT.value:
            _logger.info("Using Azure DefaultCredentials to access KeyVault")
            credentials = DefaultAzureCredential()

        elif self.config["authentication-method"] == KeyVaultAuthenticationMethod.CLIENTSECRET.value:
            auth_parameters = ("client-id", "tenant-id", "secret")

            _logger.info("Using Azure ClientSecret credentials to access KeyVault")

            env_file_found = load_dotenv("./.env", override=True)

            if not env_file_found:
                _logger.info(f"Local environment file not found at {Path.cwd() / '.env'}")

            if all(param in self.config for param in auth_parameters):
                tenant_id = os.path.expandvars(self.config["tenant-id"])
                client_id = os.path.expandvars(self.config["client-id"])
                secret = os.path.expandvars(self.config["secret"])

                credentials = ClientSecretCredential(
                    tenant_id=tenant_id,
                    client_id=client_id,
                    client_secret=secret,
                )
            else:
                raise InvalidConfigError(
                    "Missing client secret parameters. client-id, tenant-id and client-secret are mandatory"
                )
        else:
            raise InvalidConfigError(
                "Invalid KeyVault authentication method. Possible values : default or client-secret"
            )

        self.client = SecretClient(vault_url=vault_url, credential=credentials)

    def __call__(self, _: yaml.SafeLoader, node: yaml.Node) -> str:
        """
        Method to be called when the !keyvault tag is encountered in the YAML file.
        """
        self._init_client()
        try:
            return self.client.get_secret(node.value).value  # type: ignore  # _init_client guarantees not None
        except (ResourceNotFoundError, ServiceRequestError, HttpResponseError) as e:
            raise InvalidConfigError(str(e)) from e


class _EnvLoader(yaml.SafeLoader):
    pass


class SafeLoaderIgnoreUnknown(yaml.SafeLoader):
    """
    Variant of PyYAML's SafeLoader that ignores unknown tags.
    """

    def ignore_unknown(self, node: yaml.Node) -> None:
        """
        Constructor for unknown tags that does nothing.
        """
        return None


def _env_constructor(_: yaml.SafeLoader, node: yaml.Node) -> bool:
    bool_values = {
        "true": True,
        "false": False,
    }
    expanded_value = os.path.expandvars(node.value)
    return bool_values.get(expanded_value.lower(), expanded_value)


def _load_yaml_dict_raw(
    source: TextIO | str,
    expand_envvars: bool = True,
    keyvault_loader: KeyVaultLoader | None = None,
) -> dict[str, Any]:
    loader = _EnvLoader if expand_envvars else yaml.SafeLoader

    class SafeLoaderIgnoreUnknown(yaml.SafeLoader):
        def ignore_unknown(self, node: yaml.Node) -> None:
            return None

        # Ignoring types since the key can be None.

    SafeLoaderIgnoreUnknown.add_constructor(None, SafeLoaderIgnoreUnknown.ignore_unknown)  # type: ignore
    initial_load = yaml.load(source, Loader=SafeLoaderIgnoreUnknown)  # noqa: S506

    if not isinstance(initial_load, dict):
        raise InvalidConfigError("The root node of the YAML document must be an object")

    if not isinstance(source, str):
        source.seek(0)

    if keyvault_loader:
        _EnvLoader.add_constructor("!keyvault", keyvault_loader)
    else:
        keyvault_config = initial_load.get("azure-keyvault", initial_load.get("key-vault"))
        _EnvLoader.add_constructor("!keyvault", KeyVaultLoader(keyvault_config))

    _EnvLoader.add_implicit_resolver("!env", re.compile(r"\$\{([^}^{]+)\}"), None)
    _EnvLoader.add_constructor("!env", _env_constructor)

    try:
        config_dict = yaml.load(source, Loader=loader)  # noqa: S506
    except ScannerError as e:
        location = e.problem_mark or e.context_mark
        formatted_location = (
            f" at line {location.line + 1}, column {location.column + 1}" if location is not None else ""
        )
        cause = e.problem or e.context
        raise InvalidConfigError(f"Invalid YAML{formatted_location}: {cause or ''}") from e

    return config_dict


def _load_yaml_dict(
    source: TextIO | str,
    case_style: str = "hyphen",
    expand_envvars: bool = True,
    dict_manipulator: Callable[[dict[str, Any]], dict[str, Any]] = lambda x: x,
    keyvault_loader: KeyVaultLoader | None = None,
) -> dict[str, Any]:
    config_dict = _load_yaml_dict_raw(source, expand_envvars, keyvault_loader)

    config_dict = dict_manipulator(config_dict)
    config_dict = _to_snake_case(config_dict, case_style)

    if "azure_keyvault" in config_dict:
        config_dict.pop("azure_keyvault")
    if "key_vault" in config_dict:
        config_dict.pop("key_vault")

    return config_dict


def _load_yaml(
    source: TextIO | str,
    config_type: type[CustomConfigClass],
    case_style: str = "hyphen",
    expand_envvars: bool = True,
    dict_manipulator: Callable[[dict[str, Any]], dict[str, Any]] = lambda x: x,
    keyvault_loader: KeyVaultLoader | None = None,
) -> CustomConfigClass:
    config_dict = _load_yaml_dict(
        source,
        case_style=case_style,
        expand_envvars=expand_envvars,
        dict_manipulator=dict_manipulator,
        keyvault_loader=keyvault_loader,
    )

    try:
        config = dacite.from_dict(
            data=config_dict,
            data_class=config_type,
            config=dacite.Config(strict=True, cast=[Enum, TimeIntervalConfig, Path, CastableInt, PortNumber]),
        )
    except dacite.UnexpectedDataError as e:
        unknowns = [f'"{k.replace("_", "-") if case_style == "hyphen" else k}"' for k in e.keys]
        raise InvalidConfigError(
            f"Unknown config parameter{'s' if len(unknowns) > 1 else ''} {', '.join(unknowns)}"
        ) from e

    except (dacite.WrongTypeError, dacite.MissingValueError, dacite.UnionMatchError) as e:
        path = (e.field_path.replace("_", "-") if case_style == "hyphen" else e.field_path) if e.field_path else None

        def name(type_: type) -> str:
            return type_.__name__ if hasattr(type_, "__name__") else str(type_)

        def all_types(type_: type) -> Iterable[type]:
            return type_.__args__ if hasattr(type_, "__args__") else [type_]

        if isinstance(e, dacite.WrongTypeError | dacite.UnionMatchError) and e.value is not None:
            got_type = name(type(e.value))
            need_type = ", ".join(name(t) for t in all_types(e.field_type))

            raise InvalidConfigError(
                f'Wrong type for field "{path}" - got "{e.value}" of type {got_type} instead of {need_type}'
            ) from e
        raise InvalidConfigError(f'Missing mandatory field "{path}"') from e

    except dacite.ForwardReferenceError as e:
        raise ValueError(f"Invalid config class: {e!s}") from e

    config._file_hash = sha256(json.dumps(config_dict).encode("utf-8")).hexdigest()

    return config


def load_yaml(
    source: TextIO | str,
    config_type: type[CustomConfigClass],
    case_style: str = "hyphen",
    expand_envvars: bool = True,
    keyvault_loader: KeyVaultLoader | None = None,
) -> CustomConfigClass:
    """
    Read a YAML file, and create a config object based on its contents.

    Args:
        source: Input stream (as returned by open(...)) or string containing YAML.
        config_type: Class of config type (i.e. your custom subclass of BaseConfig).
        case_style: Casing convention of config file. Valid options are 'snake', 'hyphen' or 'camel'. Should be
            'hyphen'.
        expand_envvars: Substitute values with the pattern ${VAR} with the content of the environment variable VAR
        keyvault_loader: Pre-built loader for keyvault tags. Will be loaded from config if not set.

    Returns:
        An initialized config object.

    Raises:
        InvalidConfigError: If any config field is given as an invalid type, is missing or is unknown
    """
    return _load_yaml(
        source=source,
        config_type=config_type,
        case_style=case_style,
        expand_envvars=expand_envvars,
        keyvault_loader=keyvault_loader,
    )


def load_yaml_dict(
    source: TextIO | str,
    case_style: str = "hyphen",
    expand_envvars: bool = True,
    keyvault_loader: KeyVaultLoader | None = None,
) -> dict[str, Any]:
    """
    Read a YAML file and return a dictionary from its contents.

    Args:
        source: Input stream (as returned by open(...)) or string containing YAML.
        case_style: Casing convention of config file. Valid options are 'snake', 'hyphen' or 'camel'. Should be
            'hyphen'.
        expand_envvars: Substitute values with the pattern ${VAR} with the content of the environment variable VAR
        keyvault_loader: Pre-built loader for keyvault tags. Will be loaded from config if not set.

    Returns:
        A raw dict with the contents of the config file.

    Raises:
        InvalidConfigError: If any config field is given as an invalid type, is missing or is unknown
    """
    return _load_yaml_dict(
        source=source, case_style=case_style, expand_envvars=expand_envvars, keyvault_loader=keyvault_loader
    )


def compile_patterns(ignore_patterns: list[str | IgnorePattern]) -> list[re.Pattern[str]]:
    """
    List of patterns to compile.

    Args:
        ignore_patterns: A list of strings or IgnorePattern to be compiled.

    Returns:
        A list of compiled RegExp patterns.
    """
    compiled = []
    for p in ignore_patterns:
        if isinstance(p, IgnorePattern):
            compiled.append(re.compile(p.compile()))
        else:
            compiled.append(re.compile(p))
    return compiled


class ConfigResolver(Generic[CustomConfigClass]):
    """
    Class for resolving configuration files, either from a local file or a remote CDF extraction pipeline.

    Automatically reloads the configuration file if it has changed
    """

    def __init__(self, config_path: str, config_type: type[CustomConfigClass]) -> None:
        self.config_path = config_path
        self.config_type = config_type

        self._config: CustomConfigClass | None = None
        self._next_config: CustomConfigClass | None = None

        self._cognite_client: CogniteClient | None = None

    def _reload_file(self) -> None:
        with open(self.config_path) as stream:
            self._config_text = stream.read()

    @property
    def cognite_client(self) -> CogniteClient | None:
        """
        Returns a CogniteClient instance based on the configuration.
        """
        if self._cognite_client is None and self._config is not None:
            self._cognite_client = self._config.cognite.get_cognite_client("config_resolver")
        return self._cognite_client

    @cognite_client.setter
    def cognite_client(self, client: CogniteClient) -> None:
        if not isinstance(client, CogniteClient):
            raise AttributeError("cognite_client must be set to a CogniteClient instance")
        self._cognite_client = client

    @property
    def is_remote(self) -> bool:
        """
        Returns True if the configuration is a remote CDF extraction pipeline config, False if it is a local file.
        """
        raw_config_type = load_yaml_dict(self._config_text).get("type")
        if raw_config_type is None:
            _logger.warning("No config type specified, default to local")
            raw_config_type = "local"
        config_type = ConfigType(raw_config_type)
        return config_type == ConfigType.REMOTE

    @property
    def has_changed(self) -> bool:
        """
        Returns True if the configuration file has changed since the last accepted configuration.
        """
        try:
            self._resolve_config()
        except Exception:
            _logger.exception("Failed to reload configuration file")
            return False
        return self._config._file_hash != self._next_config._file_hash if self._config else True  # type: ignore

    @property
    def config(self) -> CustomConfigClass:
        """
        Returns the current configuration object. If it has not been resolved yet, it will resolve it first.
        """
        if self._config is None:
            self._resolve_config()
            self.accept_new_config()
        return self._config  # type: ignore

    def accept_new_config(self) -> None:
        """
        Accepts the new configuration, making it the current configuration.
        """
        self._config = self._next_config

    @classmethod
    def from_cli(
        cls, name: str, description: str, version: str, config_type: type[CustomConfigClass]
    ) -> "ConfigResolver":
        """
        Creates a ConfigResolver instance from command line arguments.

        Args:
            name: The name of the extractor.
            description: A description of the extractor.
            version: The version of the extractor.
            config_type: The type of the configuration class to be used.

        Returns:
            A ConfigResolver instance initialized with the configuration file path.
        """
        argument_parser = argparse.ArgumentParser(sys.argv[0], description=description)
        argument_parser.add_argument(
            "config", nargs=1, type=str, help="The YAML file containing configuration for the extractor."
        )
        argument_parser.add_argument("-v", "--version", action="version", version=f"{name} v{version}")
        args = argument_parser.parse_args()

        return cls(args.config[0], config_type)

    def _inject_cognite(self, local_part: _BaseConfig, remote_part: dict[str, Any]) -> dict[str, Any]:
        # We can not dump 'local_part.cognite' directly because e.g. 'data_set' may be set remote only...
        remote_part.setdefault("cognite", {})
        remote_part["cognite"]["idp_authentication"] = dataclasses.asdict(local_part.cognite.idp_authentication)
        remote_part["cognite"]["extraction-pipeline"] = dataclasses.asdict(
            local_part.cognite.extraction_pipeline  # type: ignore [arg-type]
        )

        if local_part.cognite.host is not None:
            remote_part["cognite"]["host"] = local_part.cognite.host
        remote_part["cognite"]["project"] = local_part.cognite.project

        return remote_part

    def _use_cached_cognite_client(self, tmp_config: _BaseConfig) -> bool:
        # Ideally we'd check tmp_config == self._config, but due to 'is_remote & _inject_...', this is not
        # reliable to avoid new unneeded instantiations of CogniteClient:
        return (
            self.cognite_client is not None
            and self._config is not None
            and tmp_config.cognite.host == self._config.cognite.host
            and tmp_config.cognite.project == self._config.cognite.project
            and tmp_config.cognite.idp_authentication == self._config.cognite.idp_authentication
        )

    def _get_keyvault_loader(self) -> KeyVaultLoader:
        temp_config = _load_yaml_dict_raw(self._config_text)
        return KeyVaultLoader(temp_config.get("azure-keyvault", temp_config.get("key-vault")))

    def _resolve_config(self) -> None:
        self._reload_file()

        if self.is_remote:
            _logger.debug("Loading remote config file")
            tmp_config: _BaseConfig = load_yaml(self._config_text, _BaseConfig)  # type: ignore
            if self._use_cached_cognite_client(tmp_config):
                # Use existing client to avoid invoking a token refresh, if possible. Reason: this is run every 5 min
                # by default ('ConfigReloader' thread) which for certain OAuth providers like Auth0, incurs a cost:
                client = cast(CogniteClient, self.cognite_client)
            else:
                # Credentials towards CDF may have changed, instantiate (and store) a new client:
                client = self.cognite_client = tmp_config.cognite.get_cognite_client("config_resolver")

            response = client.extraction_pipelines.config.retrieve(
                tmp_config.cognite.get_extraction_pipeline(client).external_id  # type: ignore  # ignoring extpipe None
            )

            if response.config is None:
                _logger.error("No config included in response from extraction pipelines")
                return

            self._next_config = _load_yaml(
                source=response.config,
                config_type=self.config_type,
                dict_manipulator=lambda d: self._inject_cognite(tmp_config, d),
                keyvault_loader=self._get_keyvault_loader(),
            )

        else:
            _logger.debug("Loading local config file")
            self._next_config = load_yaml(self._config_text, self.config_type)