data_validation/query_builder/query_builder.py

# Copyright 2020 Google LLC
#
# 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 logging
import ibis
from data_validation import consts
from ibis.expr.types import StringScalar
from third_party.ibis.ibis_addon import api, operations


class AggregateField(object):
    def __init__(self, ibis_expr, field_name=None, alias=None, cast=None):
        """A representation of a table or column aggregate in Ibis

        Args:
            ibis_expr (ColumnExpr): A column aggregation to use from Ibis
            field_name (String: A field to act on in the table.
                Table level expr do not have a field name
            alias (String): A field to use as the aggregate alias name
        """
        self.expr = ibis_expr
        self.field_name = field_name
        self.alias = alias
        self.cast = cast

    @staticmethod
    def count(field_name=None, alias=None, cast=None):
        if field_name:
            return AggregateField(
                ibis.expr.types.ColumnExpr.count,
                field_name=field_name,
                alias=alias,
                cast=cast,
            )
        else:
            return AggregateField(
                ibis.expr.types.TableExpr.count,
                field_name=field_name,
                alias=alias,
                cast=cast,
            )

    @staticmethod
    def min(field_name=None, alias=None, cast=None):
        return AggregateField(
            ibis.expr.types.ColumnExpr.min,
            field_name=field_name,
            alias=alias,
            cast=cast,
        )

    @staticmethod
    def avg(field_name=None, alias=None, cast=None):
        return AggregateField(
            ibis.expr.types.NumericColumn.mean,
            field_name=field_name,
            alias=alias,
            cast=cast,
        )

    @staticmethod
    def max(field_name=None, alias=None, cast=None):
        return AggregateField(
            ibis.expr.types.ColumnExpr.max,
            field_name=field_name,
            alias=alias,
            cast=cast,
        )

    @staticmethod
    def sum(field_name=None, alias=None, cast=None):
        return AggregateField(
            ibis.expr.api.IntegerColumn.sum,
            field_name=field_name,
            alias=alias,
            cast=cast,
        )

    @staticmethod
    def bit_xor(field_name=None, alias=None, cast=None):
        return AggregateField(
            ibis.expr.api.IntegerColumn.bit_xor,
            field_name=field_name,
            alias=alias,
            cast=cast,
        )

    @staticmethod
    def std(field_name=None, alias=None, cast=None):
        return AggregateField(
            ibis.expr.types.NumericColumn.std,
            field_name=field_name,
            alias=alias,
            cast=cast,
        )

    def compile(self, ibis_table):
        if self.field_name:
            agg_field = self.expr(ibis_table[self.field_name])
        else:
            agg_field = self.expr(ibis_table)

        if self.cast:
            agg_field = agg_field.force_cast(self.cast)

        if self.alias:
            agg_field = agg_field.name(self.alias)

        return agg_field


class FilterField(object):
    def __init__(
        self, ibis_expr, left=None, right=None, left_field=None, right_field=None
    ):
        """A representation of a query filter to be used while building a query.
            You can alternatively use either (left or left_field) and
            (right or right_field).

        Args:
            ibis_expr (ColumnExpr): A column expression to be used for comparisons (None represents a custom filter).
            left (Object): A value to compare on the left side of the expression
            left_field (String): A column name to be used to filter against
            right (Object): A value to compare on the right side of the expression
            right_field (String): A column name to be used to filter against

        """
        self.expr = ibis_expr
        self.left = left
        self.right = right
        self.left_field = left_field
        self.right_field = right_field

    @staticmethod
    def greater_than(field_name, value):
        # Build Left and Right Objects
        return FilterField(
            ibis.expr.types.ColumnExpr.__gt__, left_field=field_name, right=value
        )

    @staticmethod
    def less_than(field_name, value):
        # Build Left and Right Objects
        return FilterField(
            ibis.expr.types.ColumnExpr.__lt__, left_field=field_name, right=value
        )

    @staticmethod
    def equal_to(field_name, value):
        # Build Left and Right Objects
        return FilterField(
            ibis.expr.types.ColumnExpr.__eq__, left_field=field_name, right=value
        )

    @staticmethod
    def isin(field_name, values):
        # Build Left and Right Objects
        return FilterField(
            ibis.expr.types.ColumnExpr.isin, left_field=field_name, right=values
        )

    @staticmethod
    def custom(expr):
        """Returns a FilterField instance built for any custom SQL using a supported operator.

        Args:
            expr (Str): A custom SQL expression used to filter a query.
        """
        return FilterField(None, left=expr)

    @staticmethod
    def or_(field_list: list):
        return FilterField(ibis.or_, left=field_list)

    def compile(self, ibis_table):
        if self.expr is None:
            return operations.compile_raw_sql(ibis_table, self.left)

        if self.left_field:
            self.left = ibis_table[self.left_field]

        if self.right_field:
            self.right = ibis_table[self.right_field]

        if self.expr == ibis.or_:
            return self.expr(*[_.compile(ibis_table) for _ in self.left])
        else:
            return self.expr(self.left, self.right)


class ComparisonField(object):
    def __init__(
        self, field_name: str, alias: str = None, cast: str = None, trim: bool = None
    ):
        """A representation of a comparison field used to build a query.

        Args:
            field_name (String): A field to act on in the table
            alias (String): An alias to use for the group
            cast (String): A cast on the column if required
        """
        self.field_name = field_name
        self.alias = alias
        self.cast = cast
        self.trim = trim

    def compile(self, ibis_table):
        # Fields are supplied on compile or on build
        comparison_field = ibis_table[self.field_name]
        alias = self.alias or self.field_name
        if self.cast:
            comparison_field = comparison_field.force_cast(self.cast)
        elif self.trim and comparison_field.type().is_string():
            comparison_field = comparison_field.rstrip()
        comparison_field = comparison_field.name(alias)

        return comparison_field


class GroupedField(object):
    def __init__(self, field_name, alias=None, cast=None):
        """A representation of a group by field used to build a query.

        Args:
            field_name (String): A field to act on in the table
            alias (String): An alias to use for the group
            cast (String): A cast on the column if required
        """
        self.field_name = field_name
        self.alias = alias
        self.cast = cast

    def compile(self, ibis_table):
        # Fields are supplied on compile or on build
        group_field = ibis_table[self.field_name]

        # TODO: generate cast for known types not specified
        if self.cast:
            group_field = group_field.cast(self.cast)
        elif isinstance(group_field.type(), ibis.expr.datatypes.Timestamp):
            group_field = group_field.cast("date")
        else:
            # TODO: need to build Truncation Int support
            # TODO: should be using a logger
            logging.warning("Unknown cast types can cause memory errors")

        # The Casts require we also supply a name.
        alias = self.alias or self.field_name
        group_field = group_field.name(alias)

        return group_field


class ColumnReference(object):
    def __init__(self, column_name):
        """A representation of an calculated field to build a query.

        Args:
            column_name (String): The column name used in a complex expr
        """
        self.column_name = column_name

    def compile(self, ibis_table):
        """Return an ibis object referencing the column.

        Args:
            ibis_table (IbisTable): The table obj reference
        """
        return ibis_table[self.column_name]


class CalculatedField(object):
    def __init__(self, ibis_expr, config, fields, cast=None, **kwargs):
        """A representation of an calculated field to build a query.

        Args:
            config dict: Configurations object explaining calc field details
            fields list: List of fields to transform into a single column
        """
        self.expr = ibis_expr
        self.config = config
        self.fields = fields
        self.cast = cast
        self.kwargs = kwargs

    def __repr__(self):
        return (
            f"CalculatedField(fields={self.fields}, expr={self.expr}, cast={self.cast})"
        )

    @staticmethod
    def concat(config, fields):
        if config.get("default_concat_separator") is None:
            config["default_concat_separator"] = ibis.literal("")
        fields = [config["default_concat_separator"], fields]
        cast = "string"
        return CalculatedField(
            ibis.expr.types.StringValue.join,
            config,
            fields,
            cast=cast,
        )

    @staticmethod
    def hash(config, fields):
        if config.get("default_hash_function") is None:
            how = "sha256"
            return CalculatedField(
                ibis.expr.types.StringValue.hashbytes,
                config,
                fields,
                how=how,
            )
        else:
            how = "farm_fingerprint"
            return CalculatedField(
                ibis.expr.types.Value.hash,
                config,
                fields,
                how=how,
            )

    @staticmethod
    def to_char(config, fields):
        fmt = ibis.literal(config.get("default_to_char_fmt", "FM90.099"))
        return CalculatedField(
            ibis.expr.api.NumericValue.to_char, config, fields, fmt=fmt
        )

    @staticmethod
    def ifnull(config, fields):
        default_null_string = ibis.literal(
            config.get("default_null_string", "DEFAULT_REPLACEMENT_STRING")
        )
        fields = [fields[0], default_null_string]
        return CalculatedField(
            ibis.expr.types.Value.fillna,
            config,
            fields,
        )

    @staticmethod
    def length(config, fields):
        return CalculatedField(
            ibis.expr.types.StringValue.length,
            config,
            fields,
        )

    @staticmethod
    def padded_char_length(config, fields):
        return CalculatedField(
            ibis.expr.types.StringValue.padded_char_length,
            config,
            fields,
        )

    @staticmethod
    def byte_length(config, fields):
        return CalculatedField(
            ibis.expr.types.BinaryValue.byte_length,
            config,
            fields,
        )

    @staticmethod
    def rstrip(config, fields):
        return CalculatedField(
            ibis.expr.types.StringValue.rstrip,
            config,
            fields,
        )

    @staticmethod
    def upper(config, fields):
        return CalculatedField(
            ibis.expr.types.StringValue.upper,
            config,
            fields,
        )

    @staticmethod
    def epoch_seconds(config, fields):
        return CalculatedField(
            ibis.expr.types.TimestampValue.epoch_seconds,
            config,
            fields,
        )

    @staticmethod
    def cast(config, fields):
        target_type = config.get(consts.CONFIG_DEFAULT_CAST, "string")
        return CalculatedField(
            api.cast,
            config,
            fields,
            target_type=target_type,
        )

    @staticmethod
    def custom(config, fields):
        """Returns a CalculatedField instance built for any custom ibis expression
        e.g. 'ibis.expr.api.StringValue.replace'. For a list of supported functions,
        see https://github.com/ibis-project/ibis/blob/1.4.0/ibis/expr/api.py
        Args:
            expr (Str): A custom ibis expression to be used as a calc field
        """
        ibis_expr = config.get(consts.CONFIG_CUSTOM_IBIS_EXPR)
        expr_params = config.get(consts.CONFIG_CUSTOM_PARAMS, [])
        params = {k: v for d in expr_params for k, v in d.items()}
        return CalculatedField(eval(ibis_expr), config, fields, **params)

    def _compile_fields(self, ibis_table, fields):
        compiled_fields = []

        for field in fields:
            if type(field) in [StringScalar]:
                compiled_fields.append(field)
            elif isinstance(field, list):
                compiled_fields.append(self._compile_fields(ibis_table, field))
            else:
                if self.cast:
                    compiled_fields.append(ibis_table[field].cast(self.cast))
                else:
                    compiled_fields.append(ibis_table[field])
        return compiled_fields

    def compile(self, ibis_table):
        compiled_fields = self._compile_fields(ibis_table, self.fields)
        calc_field = self.expr(*compiled_fields, **self.kwargs)
        if self.config["field_alias"]:
            calc_field = calc_field.name(self.config["field_alias"])

        return calc_field


class QueryBuilder(object):
    def __init__(
        self,
        aggregate_fields,
        calculated_fields,
        filters,
        grouped_fields,
        comparison_fields,
        limit=None,
    ):
        """Build a QueryBuilder object which can be used to build queries easily

        Args:
            aggregate_fields (list[AggregateField]): AggregateField instances with Ibis expressions
            calculated_fields (list[CalculatedField]): A list of CalculatedField instances
            filters (list[FilterField]): A list of FilterField instances
            grouped_fields (list[GroupedField]): A list of GroupedField instances
            limit (int): A limit value for the number of records to pull (used for testing)
        """
        self.aggregate_fields = aggregate_fields
        self.calculated_fields = calculated_fields
        self.filters = filters
        self.grouped_fields = grouped_fields
        self.comparison_fields = comparison_fields
        self.limit = limit

    @staticmethod
    def build_count_validator(limit=None):
        """Return a basic template builder for most validations"""
        aggregate_fields = []
        filters = []
        grouped_fields = []
        comparison_fields = []
        calculated_fields = []

        return QueryBuilder(
            aggregate_fields,
            filters=filters,
            grouped_fields=grouped_fields,
            comparison_fields=comparison_fields,
            calculated_fields=calculated_fields,
        )

    def compile_aggregate_fields(self, table):
        aggs = [field.compile(table) for field in self.aggregate_fields]

        return aggs

    def compile_filter_fields(self, table):
        return [field.compile(table) for field in self.filters]

    def compile_group_fields(self, table):
        return [field.compile(table) for field in self.grouped_fields]

    def compile_comparison_fields(self, table):
        return [field.compile(table) for field in self.comparison_fields]

    def compile_calculated_fields(self, table, n=0):
        return [
            field.compile(table)
            for field in self.calculated_fields
            if field.config[consts.CONFIG_DEPTH] == n
        ]
        # if n is not None:
        #     return [
        #         field.compile(table)
        #         for field in self.calculated_fields
        #         if field.config[consts.CONFIG_DEPTH] == n
        #     ]
        # else:
        #     return [field.compile(table) for field in self.calculated_fields]

    def compile(self, validation_type, table):
        """Return an Ibis query object

        Args:
            table (IbisTable): The Ibis Table expression.
        """

        # Build Query Expressions
        compiled_filters = self.compile_filter_fields(table)
        filtered_table = table.filter(compiled_filters) if compiled_filters else table

        if self.calculated_fields:
            depth_limit = max(
                field.config.get(consts.CONFIG_DEPTH, 0)
                for field in self.calculated_fields
            )
            for n in range(0, (depth_limit + 1)):
                filtered_table = filtered_table.mutate(
                    self.compile_calculated_fields(filtered_table, n)
                )

        if (
            validation_type == consts.ROW_VALIDATION
            or validation_type == consts.CUSTOM_QUERY
        ):
            if self.comparison_fields:
                filtered_table = filtered_table.projection(
                    self.compile_comparison_fields(filtered_table)
                )
        else:
            if self.comparison_fields:
                filtered_table = filtered_table.mutate(
                    self.compile_comparison_fields(filtered_table)
                )
        compiled_groups = self.compile_group_fields(filtered_table)
        grouped_table = (
            filtered_table.group_by(compiled_groups)
            if compiled_groups
            else filtered_table
        )
        if self.aggregate_fields:
            query = grouped_table.aggregate(
                self.compile_aggregate_fields(filtered_table)
            )
        else:
            query = grouped_table

        if self.limit:
            query = query.limit(self.limit)

        return query

    def add_aggregate_field(self, aggregate_field):
        """Add an AggregateField instance to the query which
            will be used when compiling your query (ie. SUM(a))

        Args:
            aggregate_field (AggregateField): An AggregateField instance
        """
        self.aggregate_fields.append(aggregate_field)

    def add_comparison_field(self, comparison_field):
        """Add an ComparisonField instance to the query which
            will be used when compiling your query (ie. SUM(a))

        Args:
            comparison_field (ComparisonField): An ComparisonField instance
        """
        self.comparison_fields.append(comparison_field)

    def add_grouped_field(self, grouped_field):
        """Add a GroupedField instance to the query which
            represents adding a column to group by in the
            query being built.
        Args:
            grouped_field (GroupedField): A GroupedField instance
        """
        self.grouped_fields.append(grouped_field)

    def add_filter_field(self, filter_obj):
        """Add a FilterField instance to your query which
            will add the desired filter to your compiled
            query (ie. WHERE query_filter=True)
        Args:
            filter_obj (FilterField): A FilterField instance
        """
        self.filters.append(filter_obj)

    def add_calculated_field(self, calculated_field):
        """Add a CalculatedField instance to your query which
            will add the desired scalar function to your compiled
            query (ie. CONCAT(field_a, field_b))
        Args:
            calculated_field (CalculatedField): A CalculatedField instance
        """
        self.calculated_fields.append(calculated_field)