rules.py

"""Metadata classes for rule steps."""

import warnings
from abc import ABCMeta, abstractmethod
from typing import (
    Any,
    ClassVar,
    Dict,
    Iterator,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    TypeVar,
    Union,
)

from pydantic import BaseModel, Extra, Field, root_validator, validate_arguments, validator
from typing_extensions import Literal

from dve.core_engine.backends.base.reference_data import ReferenceConfigUnion
from dve.core_engine.backends.metadata.reporting import LegacyReportingConfig, ReportingConfig
from dve.core_engine.templating import template_object
from dve.core_engine.type_hints import (
    Alias,
    EntityName,
    Expression,
    JSONable,
    MultipleExpressions,
    TemplateVariables,
)

__all__ = [
    "AbstractStep",
    "Aggregation",
    "AntiJoin",
    "BaseStep",
    "ColumnAddition",
    "ColumnRemoval",
    "ConfirmJoinHasMatch",
    "CopyEntity",
    "DeferredFilter",
    "EntityRemoval",
    "HeaderJoin",
    "ImmediateFilter",
    "InnerJoin",
    "LeftJoin",
    "OneToOneJoin",
    "OneToOneJoin",
    "OrphanIdentification",
    "ParentMetadata",
    "RenameEntity",
    "Rule",
    "RuleMetadata",
    "SelectColumns",
    "SemiJoin",
    "TableUnion",
]

ASSelf = TypeVar("ASSelf", bound="AbstractStep")
RSelf = TypeVar("RSelf", bound="Rule")

Stage = Literal["Pre-sync", "Sync", "Post-sync"]
"""The stage that the rule falls under."""


# pylint: disable=too-few-public-methods
class ParentMetadata(BaseModel):
    """Data about the parent rule for a step."""

    rule: Union["Rule", str]
    """The rule the step belongs to, or a string representing the rule."""
    index: int
    """The index of the step within the parent rule."""
    stage: Stage
    """The name of the stage in the rule the step belongs to."""

    def __repr__(self) -> str:
        components: List[str] = []
        if isinstance(self.rule, Rule):
            components.append(f"rule=Rule(name={self.rule.name!r}, ...)")
        else:
            components.append(f"rule=Rule(name={self.rule!r}, ...)")

        for key, value in self.dict(exclude={"rule"}).items():
            components.append(f"{key}={value!r}")

        return f"ParentMetadata({', '.join(components)})"

    class Config:  # pylint: disable=too-few-public-methods
        """`pydantic configuration options`"""

        frozen = True
        extra = Extra.forbid


# pylint: disable=too-few-public-methods
class AbstractStep(BaseModel, metaclass=ABCMeta):
    """An abstract transformation step."""

    id: Optional[str] = None
    """An ID for the rule step. This will not be templated."""
    description: Optional[str] = None
    """An optional description for the step."""
    parent: Optional[ParentMetadata] = None
    """Data about the parent rule and the step's place within it."""

    UNTEMPLATED_KEYS: ClassVar[Set[str]] = {"id", "description", "parent"}
    """A set of aliases which are exempted from templating."""

    class Config:  # pylint: disable=too-few-public-methods
        """`pydantic configuration options`"""

        frozen = True
        extra = Extra.forbid

    def __repr_args__(self) -> Sequence[Tuple[Optional[str], Any]]:
        # Exclude nulls from 'repr' for conciseness.
        return [(key, value) for key, value in super().__repr_args__() if value is not None]

    @abstractmethod
    def get_required_entities(self) -> Set[EntityName]:
        """Get a set of the required entity names for the transformation."""
        raise NotImplementedError()  # pragma: no cover

    @abstractmethod
    def get_created_entities(self) -> Set[EntityName]:
        """Get a set of the entity names created by the transformation."""
        raise NotImplementedError()  # pragma: no cover

    def get_removed_entities(self) -> Set[EntityName]:
        """Get a set of the entity names removed by the transformation."""
        return set()

    def with_parent(self: ASSelf, parent: ParentMetadata) -> ASSelf:
        """Return a new step with different parent metadata."""
        return type(self)(**self.dict(exclude={"parent"}), parent=parent)

    def template(
        self: ASSelf,
        local_variables: TemplateVariables,
        *,
        global_variables: TemplateVariables = None,
    ) -> ASSelf:
        """Template the rule, given the global and local variables."""
        type_ = type(self)
        if global_variables:
            variables = global_variables.copy()
            variables.update(local_variables)
        else:
            variables = local_variables

        untemplated_data = {key: getattr(self, key) for key in self.UNTEMPLATED_KEYS}
        data_to_template = self.dict(exclude=self.UNTEMPLATED_KEYS)
        templated_data = template_object(data_to_template, variables, "jinja")

        return type_(**templated_data, **untemplated_data)

    def __str__(self):  # pydantic's default __str__ strips the model name.
        return super().__repr__()

    @root_validator(pre=True)
    @classmethod
    def _warn_for_deprecated_aliases(cls, values: Dict[str, JSONable]) -> Dict[str, JSONable]:
        for deprecated_name, replacement in (
            ("entity", "entity_name"),
            ("target", "target_entity_name"),
        ):
            if deprecated_name in values:
                warnings.warn(
                    f"Using deprecated name {deprecated_name!r}, should be {replacement!r}"
                )
                if replacement in values:
                    raise ValueError(
                        f"{replacement!r} and deprecated name {deprecated_name!r} both "
                        + "provided"
                    )
                values[replacement] = values.pop(deprecated_name)
        return values


class BaseStep(AbstractStep, metaclass=ABCMeta):
    """A base transformation step which performs a transformation on an entity
    and provides for a new entity being created instead of modification.

    """

    entity_name: EntityName
    """The entity to apply the step to."""
    new_entity_name: Optional[EntityName] = None
    """Optionally, a new entity to create after the operation."""

    def get_required_entities(self) -> Set[EntityName]:
        """Get a set of the required entity names for the transformation."""
        return {self.entity_name}

    def get_created_entities(self) -> Set[EntityName]:
        """Get a set of the entity names created by the transformation."""
        return {self.new_entity_name or self.entity_name}


class ImmediateFilter(BaseStep):
    """A filter which removes records from an entity (or creates a new entity without
    those rows) according to a provided expression. Only records for which the
    provided expression is truthy will be retained.

    """

    expression: Expression
    """
    The expression for the filter. Records for which this expression is falsey will be
    removed.

    """


class DeferredFilter(AbstractStep):
    """A filter which should be applied in a synchronised step with other filters for
    the same entity, providing feedback to users.

    Messages will be  emitted for each record that fails a filter according to the
    configured reporting, and any record-level errors will be removed from source
    entities at the end of the stage. This enables multiple issues with the same record
    to be flagged to a data provider.

    """

    entity_name: EntityName
    """The entity to apply the step to."""
    expression: Expression
    """
    The expression for the filter. Records for which this expression is falsey will be
    removed from the source entity if the reporting level is a record-level error.

    """
    reporting: Union[ReportingConfig, LegacyReportingConfig]
    """The reporting information for the filter."""

    def template(
        self: "DeferredFilter",
        local_variables: Dict[Alias, Any],
        *,
        global_variables: Dict[Alias, Any] = None,
    ) -> "DeferredFilter":
        """Template the rule, given the global and local variables."""
        type_ = type(self)
        if global_variables:
            variables = global_variables.copy()
            variables.update(local_variables)
        else:
            variables = local_variables

        untemplated_data = {key: getattr(self, key) for key in self.UNTEMPLATED_KEYS}
        data_to_template = self.dict(exclude={*self.UNTEMPLATED_KEYS, "reporting"})
        templated_data = template_object(data_to_template, variables, "jinja")
        templated_data["reporting"] = self.reporting.template(
            local_variables, global_variables=global_variables
        )

        return type_(**templated_data, **untemplated_data)

    def get_required_entities(self) -> Set[EntityName]:
        """Get a set of the required entity names for the transformation."""
        return {self.entity_name}

    def get_created_entities(self) -> Set[EntityName]:
        """Get a set of the required entity names for the transformation."""
        return {self.entity_name}


class Notification(AbstractStep):
    """An implementation of a notification according to a particular reporting configuration.
    This will emit notifications for all records that meet the expression, and will
    not modify the source entity.

    NOTE: This is not intended to be used directly, but is used in the implementation
    of the sync filter stage.

    """

    entity_name: EntityName
    """The entity to apply the step to."""
    expression: Expression
    """
    The expression for the notification. Records for which this expression is truthy will
    emit notifications according to the reporting config.

    """
    excluded_columns: List[Alias] = Field(default_factory=list)
    """Columns to be excluded from the record in the report."""
    reporting: ReportingConfig
    """The reporting information for the filter."""

    def get_required_entities(self) -> Set[EntityName]:
        return {self.entity_name}

    def get_created_entities(self) -> Set[EntityName]:
        return set()


class AbstractJoin(AbstractStep, metaclass=ABCMeta):
    """An abstract table join configuration. This joins `target` into `entity`."""

    entity_name: EntityName
    """The name of the source entity to join from."""
    target_name: EntityName
    """The target dataset to join to."""
    new_entity_name: Optional[EntityName] = None
    """Optionally, a new entity to create after the operation."""

    def get_required_entities(self) -> Set[EntityName]:
        return {self.entity_name, self.target_name}

    def get_created_entities(self) -> Set[EntityName]:
        return {self.new_entity_name or self.entity_name}


class AbstractConditionalJoin(AbstractJoin, metaclass=ABCMeta):
    """An abstract table join configuration with a join condition.
    This joins `target` into `entity`.

    """

    join_condition: Expression
    """
    A SQL expression representing the join. Columns will be namespaced
    within the source and target dataset names.

    """

    JOIN_TYPE: ClassVar[str]
    """The type of join that the class uses in Spark."""


class AbstractNewColumnConditionalJoin(AbstractConditionalJoin, metaclass=ABCMeta):
    """A join type which adds new columns to the finished entity."""

    new_columns: MultipleExpressions = Field(default_factory=dict)
    """
    A mapping of SQL expressions to the names of the new columns to be added
    to `entity`.

    Expressions can access columns from the source and the target, and
    columns will be namespaced within the source and target dataset names.

    """


class ColumnAddition(BaseStep):
    """A transformation step which adds a new column based on an expression."""

    column_name: Alias
    """The name of the column to be created."""
    expression: Expression
    """The SQL expression for the transformation."""


class ColumnRemoval(BaseStep):
    """A transformation step which removes a column by name."""

    column_name: Alias
    """The name of the column to be dropped."""


class SelectColumns(BaseStep):
    """A transformation step which selects columns from an entity."""

    columns: MultipleExpressions
    """Multiple expressions to select from the specified entity."""
    distinct: bool = False
    """Whether to keep only distinct columns from the select statement."""


class Aggregation(BaseStep):
    """A transformation which performs aggregation."""

    group_by: MultipleExpressions
    """Multiple expressions to group by."""
    pivot_column: Optional[Alias] = None
    """An optional pivot column for the table."""
    pivot_values: Optional[List[Any]] = None
    """A list of values to translate to columns when pivoting."""
    agg_columns: Optional[MultipleExpressions] = None
    """Multiple aggregate expressions to take from the group by (for spark backend)"""
    agg_function: Optional[Alias] = None
    """The aggregate function to apply to the agg_columns (for duckdb backend)"""

    @validator("pivot_values")
    @classmethod
    def _ensure_column_if_values(
        cls,
        value: Optional[Any],
        values: Dict[str, Any],
    ):
        """Ensure that `pivot_column` is not null if pivot values are provided."""
        if value and not values["pivot_column"]:
            raise ValueError("`pivot_values` specified, but no `pivot_column`")
        return value

    @validator("agg_function")
    @classmethod
    def _ensure_column_if_function(
        cls,
        agg_function: Optional[Any],
        values: Dict[str, Any],
    ):
        """Ensure that `pivot_column` is not null if pivot values are provided."""
        if agg_function and not values["agg_columns"]:
            raise ValueError("`agg_function` specified, but no `agg_columns`")
        return agg_function


class EntityRemoval(AbstractStep):
    """A transformation which drops entities."""

    entity_name: Union[EntityName, List[EntityName]]
    """The entity to drop."""

    def get_required_entities(self) -> Set[EntityName]:
        """Get a set of the required entity names for the transformation."""
        return set()

    def get_created_entities(self) -> Set[EntityName]:
        """Get a set of the entity names created by the transformation."""
        return set()

    def get_removed_entities(self) -> Set[EntityName]:
        """Get a set of the entity names created by the transformation."""
        if isinstance(self.entity_name, list):
            return set(self.entity_name)
        return {self.entity_name}


class CopyEntity(AbstractStep):
    """A transformation which copies an entity."""

    entity_name: EntityName
    """The entity to copy."""
    new_entity_name: EntityName
    """The new name for the copied entity."""

    def get_required_entities(self) -> Set[EntityName]:
        """Get a set of the required entity names for the transformation."""
        return {self.entity_name}

    def get_created_entities(self) -> Set[EntityName]:
        """Get a set of the entity names created by the transformation."""
        return {self.new_entity_name}

    def get_removed_entities(self) -> Set[EntityName]:
        """Gets the entity which has been removed"""
        return set()


class RenameEntity(CopyEntity):
    """A transformation which renames an entity."""

    def get_removed_entities(self) -> Set[EntityName]:
        """Get a set of the entity names removed by the transformation."""
        return {self.entity_name}


class LeftJoin(AbstractNewColumnConditionalJoin):
    """A table join configuration. This joins `target` into `entity`.

    The table join will result in all of the columns from `entity` being retained,
    new columns can be added based on data from `target` by adding SQL expressions
    in `new_columns`.

    """


class InnerJoin(AbstractNewColumnConditionalJoin):
    """A table join configuration. This joins `target` into `entity`, retaining only
    rows which match in both tables.

    The table join will result in all of the columns from `entity` being retained,
    new columns can be added based on data from `target` by adding SQL expressions
    in `new_columns`.

    """


class OneToOneJoin(LeftJoin):
    """A table join configuration. This joins `target` into `entity`. This
    will not alter the number of rows in `target`.

    The table join will result in all of the columns from `entity` being retained,
    new columns can be added based on data from `target` by adding SQL expressions
    in `new_columns`.

    """

    perform_integrity_check: bool = True
    """
    Whether to perform the integrity check on the number of records. If this is
    disabled, it is not guaranteed that joins will be one to one.

    """


class SemiJoin(AbstractConditionalJoin):
    """A table join configuration. This keeps the rows in `entity` where the join
    condition with `target` is true.

    The table join will result in all of the columns from `entity` being retained.
    No rows can be added from `target`.

    """


class AntiJoin(AbstractConditionalJoin):
    """A table join configuration. This keeps the rows in `entity` where the join
    condition with `target` is _not_ true.

    The table join will result in all of the columns from `entity` being retained.
    No rows can be added from `target`.

    """


class ConfirmJoinHasMatch(AbstractConditionalJoin):
    """Add a boolean column with name `column_name` to `entity`. The new column
    will indicate whether the `entity` has a corresponding row in `target`.

    """

    column_name: Alias
    """
    The name of the new column to add to `entity`. This will indicate whether
    there is a corresponding match for each row.

    """
    perform_integrity_check: bool = True
    """
    Whether to perform the integrity check on the number of records. If this is
    disabled, it is not guaranteed that rows in `entity` will not be duplicated.

    """


class HeaderJoin(AbstractJoin):
    """Add a 'header' entity (`target`) to each row in `entity`. The header entity
    must contain only a single row.

    """

    header_column_name: Alias = "_Header"
    """The name of the column to add to `entity`, which will contain the header."""


class TableUnion(AbstractJoin):
    """Union two tables together, taking the columns from each by name.

    Where columns have the same name, they must be the same type or coerceable.
    Where column casing differs, the casing from the `source` entity will be kept.

    Column order will be preserved, with columns from `source` taken first and extra
    columns in `target` added in order afterwards.

    """


class OrphanIdentification(AbstractConditionalJoin):
    """Identify records in `entity` which don't have at least one corresponding
    match in `target`. A new boolean column will be added to `entity` ('IsOrphaned')
    indicating whether the condition matched.

    If there is already an 'IsOrphaned' column in the entity, this will be set to the
    logical OR of its current value and the value it would have been set to otherwise.

    """


Step = Union[AbstractStep, Literal["sync"]]
"""A step within a rule. This is either a rule config or the literal string 'sync'."""


class Rule(BaseModel):
    """A rule, made up of multiple steps."""

    name: str
    """The name of the rule."""
    pre_sync_steps: List[AbstractStep]
    """The pre-sync steps in the rule."""
    sync_filter_steps: List[DeferredFilter]
    """The sync filter steps in the rule."""
    post_sync_steps: List[AbstractStep]
    """The post-sync steps in the rule."""

    def __str__(self):  # pydantic's default __str__ strips the model name.
        return super().__repr__()

    @classmethod
    @validate_arguments
    def from_step_list(cls, name: str, steps: List[Step]):
        """Load the rule from a single step list."""
        pre_sync_steps: List[AbstractStep] = []
        sync_filter_steps: List[DeferredFilter] = []
        post_sync_steps: List[AbstractStep] = []
        self = cls(
            name=name,
            pre_sync_steps=pre_sync_steps,
            sync_filter_steps=sync_filter_steps,
            post_sync_steps=post_sync_steps,
        )

        sync_stage_start: Optional[int] = None
        sync_stage_stop: Optional[int] = None
        for step_index, step in enumerate(steps):
            in_sync_stage = (sync_stage_start is not None) and (sync_stage_stop is None)
            sync_stage_completed = sync_stage_stop is not None

            step_is_sync_filter = isinstance(step, DeferredFilter)
            step_is_sync_keyword = isinstance(step, str) and step == "sync"
            step_is_sync = step_is_sync_filter or step_is_sync_keyword

            if step_is_sync:
                if sync_stage_completed:
                    raise ValueError(
                        "Synchronised steps must be contiguous. The step at index "
                        + f"{step_index} comes after the end of the first continuous "
                        + f"sync stage ([{sync_stage_start}:{sync_stage_stop}])"
                    )

                if sync_stage_start is None:
                    sync_stage_start = step_index

                if step_is_sync_filter:
                    sync_filter_steps.append(step)  # type: ignore

                continue

            if in_sync_stage:
                sync_stage_completed = True
                sync_stage_stop = step_index

            if not sync_stage_completed:
                pre_sync_steps.append(step)  # type: ignore
            else:
                post_sync_steps.append(step)  # type: ignore

        return self

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        global_index = 0
        for stage, step_list in (
            ("Pre-sync", self.pre_sync_steps),
            ("Sync", self.sync_filter_steps),
            ("Post-sync", self.post_sync_steps),
        ):
            for stage_step_index, step in enumerate(step_list):
                step_list[stage_step_index] = step.with_parent(
                    ParentMetadata(rule=self, index=global_index, stage=stage)
                )
                global_index += 1

    def template(
        self: RSelf,
        local_variables: TemplateVariables,
        *,
        global_variables: TemplateVariables = None,
    ) -> RSelf:
        """Template the rule, returning the new templated rule. This is only really useful
        for 'upfront' templating, as all stages of the rule will be templated at once.

        """
        rule_lists: Dict[str, Union[List[AbstractStep], List[DeferredFilter]]] = {
            "pre_sync_steps": self.pre_sync_steps,
            "sync_filter_steps": self.sync_filter_steps,
            "post_sync_steps": self.post_sync_steps,
        }
        for list_name, step_list in rule_lists.items():
            templated_rule_list = []
            for step in step_list:
                step = step.template(local_variables, global_variables=global_variables)
                templated_rule_list.append(step)
            rule_lists[list_name] = templated_rule_list

        type_ = type(self)
        return type_(name=self.name, **rule_lists)


class RuleMetadata(BaseModel):
    """Metadata about the rules."""

    rules: List[Rule]
    """A list of rules to be applied to to the entities."""
    local_variables: Optional[List[TemplateVariables]] = None
    """
    An optional list of local, rule-level template variables.

    If provided, this must be the same length as `rules`.

    """
    global_variables: TemplateVariables = Field(default_factory=dict)
    """An optional mapping of global template variables."""
    templating_strategy: Literal["upfront", "runtime"] = "upfront"
    """
    The templating strategy for the rules.

    If 'upfront', template all rules at the beginning of the rule stage
    execution.

    If 'runtime', template rules immediately before evaluating them.

    Runtime templating doesn't currently add any value, but will provide
    the ability to set variables from data items in the future. With runtime
    rules, it is not possible to check that all entities are present before
    performing work.

    """
    reference_data_config: Dict[EntityName, ReferenceConfigUnion]
    """
    Per-entity configuration options for the reference data.

    Options are likely to vary by backend (e.g. some backends will have
    a database and table name, some might have a remote URI).

    """

    @root_validator()
    @classmethod
    def _ensure_locals_same_length_as_rules(cls, values: Dict[str, List[Any]]):
        """Ensure that if 'local_variables' is provided, it's the same length as 'rules'."""
        local_vars = values["local_variables"]
        if local_vars is not None:
            n_local_vars = len(local_vars)
            n_rules = len(values["rules"])
            if n_local_vars != n_rules:
                raise ValueError(
                    f"Mismatch between number of local variables {n_local_vars} and number of "
                    + f"rules {n_rules}"
                )
        return values

    def __iter__(self) -> Iterator[Tuple[Rule, TemplateVariables]]:  # type: ignore
        """Iterate over the rules and local variables."""
        if self.local_variables is None:
            yield from ((rule, {}) for rule in self.rules)  # type: ignore
        else:
            yield from zip(self.rules, self.local_variables)


ParentMetadata.update_forward_refs()