9@dataclass(frozen=True)
10class AliasedExpression(Expression):
11    """
12    Used to alias the name of an expression in the results of a query. It is not used
13    anywhere in Snuba except to change the names in the results set. Right now this is
14    limited to Columns only because Functions have a separate alias. Eventually the
15    two will be combined.
16
17    :param Expression: The expression to alias.
18    :type Expression: Column
19    :raises InvalidExpressionError: If the expression or alias is invalid.
20    """
21
22    # TODO: We should eventually allow Functions here as well, once we think through
23    # how this should work with functions that already have aliases.
24    exp: Column
25    alias: Optional[str] = None
26
27    def validate(self) -> None:
28        if not isinstance(self.exp, Column):
29            raise InvalidExpressionError(
30                "aliased expressions can only contain a Column"
31            )
32
33        if self.alias is not None:
34            if not isinstance(self.alias, str) or self.alias == "":
35                raise InvalidExpressionError(
36                    f"alias '{self.alias}' of expression must be None or a non-empty string"
37                )
38            if not ALIAS_RE.match(self.alias):
39                raise InvalidExpressionError(
40                    f"alias '{self.alias}' of expression contains invalid characters"
41                )

135@dataclass(frozen=True)
136class And(BooleanCondition):
137    op: BooleanOp = field(init=False, default=BooleanOp.AND)
138    conditions: ConditionGroup = field(default_factory=list)

23class ArithmeticOperator(Enum):
24    PLUS = "plus"
25    MINUS = "minus"
26    MULTIPLY = "multiply"
27    DIVIDE = "divide"

108@dataclass(frozen=True)
109class BooleanCondition(Expression):
110    op: BooleanOp
111    conditions: ConditionGroup
112
113    def validate(self) -> None:
114        if not isinstance(self.op, BooleanOp):
115            raise InvalidConditionError(
116                "invalid boolean: operator of a boolean must be a BooleanOp"
117            )
118
119        if not isinstance(self.conditions, (list, tuple)):
120            raise InvalidConditionError(
121                "invalid boolean: conditions must be a list of other conditions"
122            )
123        elif len(self.conditions) < 2:
124            raise InvalidConditionError(
125                "invalid boolean: must supply at least two conditions"
126            )
127
128        for con in self.conditions:
129            if not isinstance(con, (Condition, BooleanCondition)):
130                raise InvalidConditionError(
131                    f"invalid boolean: {con} is not a valid condition"
132                )

103class BooleanOp(Enum):
104    AND = "AND"
105    OR = "OR"

18@dataclass(frozen=True)
19class Column(Expression):
20    """
21    A representation of a single column in the database. Columns are
22    expected to be alpha-numeric, with '.', '_`, and `:` allowed as well.
23    If the column is subscriptable then you can specify the column in the
24    form `subscriptable[key]`. The `subscriptable` attribute will contain the outer
25    column and `key` will contain the inner key.
26
27    :param name: The column name.
28    :type name: str
29    :param entity: The entity for that column
30    :type name: Optional[Entity]
31
32    :raises InvalidColumnError: If the column name is not a string or has an
33        invalid format.
34
35    """
36
37    name: str
38    entity: Optional[Entity] = None
39    subscriptable: Optional[str] = field(init=False, default=None)
40    key: Optional[str] = field(init=False, default=None)
41
42    def validate(self) -> None:
43        if not isinstance(self.name, str):
44            raise InvalidColumnError(f"column '{self.name}' must be a string")
45        if not column_name_re.match(self.name):
46            raise InvalidColumnError(
47                f"column '{self.name}' is empty or contains invalid characters"
48            )
49
50        if self.entity is not None:
51            if not isinstance(self.entity, Entity):
52                raise InvalidColumnError(f"column '{self.name}' expects an Entity")
53            if not self.entity.alias:
54                raise InvalidColumnError(
55                    f"column '{self.name}' expects an Entity with an alias"
56                )
57
58            self.validate_data_model(self.entity)
59
60        # If this is a subscriptable set these values to help with debugging etc.
61        # Because this is frozen we can't set the value directly.
62        if "[" in self.name:
63            subscriptable, key = self.name.split("[", 1)
64            key = key.strip("]")
65            super().__setattr__("subscriptable", subscriptable)
66            super().__setattr__("key", key)
67
68    def validate_data_model(self, match: Union[Entity, Storage]) -> None:
69        if match.data_model is None:
70            return
71
72        to_check = self.subscriptable if self.subscriptable else self.name
73        if not match.data_model.contains(to_check):
74            raise InvalidColumnError(
75                f"'{match.name}' does not support the column '{self.name}'"
76            )

 73@dataclass(frozen=True)
 74class Condition(Expression):
 75    lhs: Union[Column, CurriedFunction, Function]
 76    op: Op
 77    rhs: Optional[Union[Column, CurriedFunction, Function, ScalarType]] = None
 78
 79    def validate(self) -> None:
 80        if not isinstance(self.lhs, (Column, CurriedFunction, Function)):
 81            raise InvalidConditionError(
 82                f"invalid condition: LHS of a condition must be a Column, CurriedFunction or Function, not {type(self.lhs)}"
 83            )
 84        if not isinstance(self.op, Op):
 85            raise InvalidConditionError(
 86                "invalid condition: operator of a condition must be an Op"
 87            )
 88
 89        if is_unary(self.op):
 90            if self.rhs is not None:
 91                raise InvalidConditionError(
 92                    "invalid condition: unary operators don't have rhs conditions"
 93                )
 94
 95        if not isinstance(
 96            self.rhs, (Column, CurriedFunction, Function)
 97        ) and not is_scalar(self.rhs):
 98            raise InvalidConditionError(
 99                f"invalid condition: RHS of a condition must be a Column, CurriedFunction, Function or Scalar not {type(self.rhs)}"
100            )

 32@dataclass(frozen=True)
 33class CurriedFunction(Expression):
 34    function: str
 35    initializers: Optional[Sequence[Union[ScalarLiteralType, Column]]] = None
 36    parameters: Optional[
 37        Sequence[
 38            Union[
 39                ScalarType,
 40                Column,
 41                CurriedFunction,
 42                Function,
 43                Identifier,
 44                Lambda,
 45            ]
 46        ]
 47    ] = None
 48    alias: Optional[str] = None
 49
 50    def validate(self) -> None:
 51        if not isinstance(self.function, str):
 52            raise InvalidFunctionError(f"function '{self.function}' must be a string")
 53        if self.function == "":
 54            # TODO: Have a whitelist of valid functions to check, maybe even with more
 55            # specific parameter type checking
 56            raise InvalidFunctionError("function cannot be empty")
 57        if not function_name_re.match(self.function):
 58            raise InvalidFunctionError(
 59                f"function '{self.function}' contains invalid characters"
 60            )
 61
 62        if self.initializers is not None:
 63            if not isinstance(self.initializers, Sequence):
 64                raise InvalidFunctionError(
 65                    f"initializers of function {self.function} must be a Sequence"
 66                )
 67            elif not all(
 68                isinstance(param, Column) or is_literal(param)
 69                for param in self.initializers
 70            ):
 71                raise InvalidFunctionError(
 72                    f"initializers to function {self.function} must be a scalar or column"
 73                )
 74
 75        if self.alias is not None:
 76            if not isinstance(self.alias, str) or self.alias == "":
 77                raise InvalidFunctionError(
 78                    f"alias '{self.alias}' of function {self.function} must be None or a non-empty string"
 79                )
 80            if not ALIAS_RE.match(self.alias):
 81                raise InvalidFunctionError(
 82                    f"alias '{self.alias}' of function {self.function} contains invalid characters"
 83                )
 84
 85        if self.parameters is not None:
 86            if not isinstance(self.parameters, Sequence):
 87                raise InvalidFunctionError(
 88                    f"parameters of function {self.function} must be a Sequence"
 89                )
 90            for param in self.parameters:
 91                if not isinstance(
 92                    param,
 93                    (Column, CurriedFunction, Function, Identifier, Lambda),
 94                ) and not is_scalar(param):
 95                    assert not isinstance(param, bytes)  # mypy
 96                    raise InvalidFunctionError(
 97                        f"parameter '{param}' of function {self.function} is an invalid type"
 98                    )
 99
100    def __eq__(self, other: object) -> bool:
101        # Don't use the alias to compare equality
102        if not isinstance(other, CurriedFunction):
103            return False
104
105        return (
106            self.function == other.function
107            and self.initializers == other.initializers
108            and self.parameters == other.parameters
109        )

11class Direction(Enum):
12    ASC = "ASC"
13    DESC = "DESC"

16@dataclass(frozen=True, repr=False)
17class Entity(Expression):
18    name: str
19    alias: Optional[str] = None
20    sample: Optional[float] = None
21    data_model: Optional[DataModel] = field(hash=False, default=None)
22
23    def validate(self) -> None:
24        # TODO: There should be a whitelist of entity names at some point
25        if not isinstance(self.name, str) or not entity_name_re.match(self.name):
26            raise InvalidEntityError(f"'{self.name}' is not a valid entity name")
27
28        if self.sample is not None:
29            if not isinstance(self.sample, float):
30                raise InvalidEntityError("sample must be a float")
31            elif self.sample <= 0.0:
32                raise InvalidEntityError("samples must be greater than 0.0")
33
34        if self.alias is not None:
35            if not isinstance(self.alias, str) or not self.alias:
36                raise InvalidEntityError(f"'{self.alias}' is not a valid alias")
37
38        if self.data_model is not None:
39            if not isinstance(self.data_model, DataModel):
40                raise InvalidEntityError("data_model must be an instance of DataModel")
41
42    def __repr__(self) -> str:
43        alias = f", alias='{self.alias}'" if self.alias is not None else ""
44        sample = f", sample={self.sample}" if self.sample is not None else ""
45        return f"Entity('{self.name}'{alias}{sample})"

124@dataclass(frozen=True)
125class Extrapolate(Expression):
126    extrapolate: bool
127
128    def validate(self) -> None:
129        if not isinstance(self.extrapolate, bool):
130            raise InvalidExpressionError("extrapolate must be a boolean")

22@dataclass
23class Flags:
24    totals: bool | None = None
25    consistent: bool | None = None
26    turbo: bool | None = None
27    debug: bool | None = None
28    dry_run: bool | None = None
29    legacy: bool | None = None
30
31    def validate(self) -> None:
32        flags_fields = fields(self)
33        for ff in flags_fields:
34            flag = getattr(self, ff.name)
35            if flag is not None and not isinstance(flag, bool):
36                raise InvalidFlagError(f"{ff.name} must be a boolean")
37
38    def to_dict(self) -> dict[str, bool]:
39        self.validate()
40        values = asdict(self)
41        return {f: v for f, v in values.items() if v is not None}

 40@dataclass(frozen=True)
 41class Formula:
 42    function_name: str
 43    parameters: Optional[Sequence[FormulaParameterGroup]] = None
 44    aggregate_params: list[Any] | None = None
 45    filters: Optional[ConditionGroup] = None
 46    groupby: Optional[list[Column | AliasedExpression]] = None
 47
 48    def __validate_consistency(self) -> None:
 49        """
 50        Ensure that the groupby columns are consistent across all Timeseries
 51        and Formulas within this Formula."""
 52        if self.parameters is None:
 53            raise InvalidFormulaError("Formula must have parameters")
 54
 55        groupbys = set()
 56        has_timeseries = False
 57
 58        stack: list[FormulaParameterGroup] = [self]
 59        while stack:
 60            param = stack.pop()
 61            if isinstance(param, Formula):
 62                if param.groupby is not None:
 63                    groupbys.add(tuple(param.groupby))
 64
 65                if param.parameters:
 66                    stack.extend(param.parameters)
 67            elif isinstance(param, Timeseries):
 68                has_timeseries = True
 69                if param.groupby is not None:
 70                    groupbys.add(tuple(param.groupby))
 71
 72        if not has_timeseries:
 73            raise InvalidFormulaError(
 74                "Formulas must operate on at least one Timeseries"
 75            )
 76        if len(set(groupbys)) > 1:
 77            raise InvalidFormulaError(
 78                "Formula parameters must group by the same columns"
 79            )
 80
 81    def validate(self) -> None:
 82        if not isinstance(self.function_name, str):
 83            raise InvalidFormulaError(f"formula '{self.function_name}' must be a str")
 84        if self.parameters is None:
 85            raise InvalidFormulaError("Formula must have parameters")
 86        elif not isinstance(self.parameters, Sequence):
 87            raise InvalidFormulaError(
 88                f"parameters of formula {self.function_name} must be a Sequence"
 89            )
 90
 91        for param in self.parameters:
 92            if not isinstance(param, tuple(FormulaParameter)):
 93                raise InvalidFormulaError(
 94                    f"parameter '{param}' of formula {self.function_name} is an invalid type"
 95                )
 96        self.__validate_consistency()
 97
 98    def _replace(self, field: str, value: Any) -> Formula:
 99        new = replace(self, **{field: value})
100        return new
101
102    def set_parameters(self, parameters: Sequence[FormulaParameterGroup]) -> Formula:
103        if parameters is not None and not list_type(
104            parameters, (Formula, Timeseries, float, int)
105        ):
106            raise InvalidFormulaError(
107                "parameters must be a list of either Timeseries, floats, or ints"
108            )
109        return self._replace("parameters", parameters)
110
111    def set_filters(self, filters: ConditionGroup | None) -> Formula:
112        if filters is not None and not list_type(
113            filters, (BooleanCondition, Condition)
114        ):
115            raise InvalidFormulaError("filters must be a list of Conditions")
116        return self._replace("filters", filters)
117
118    def set_groupby(self, groupby: list[Column | AliasedExpression] | None) -> Formula:
119        if groupby is not None and not list_type(groupby, (Column, AliasedExpression)):
120            raise InvalidFormulaError("groupby must be a list of Columns")
121        return self._replace("groupby", groupby)

112@dataclass(frozen=True)
113class Function(CurriedFunction):
114    initializers: Optional[Sequence[Union[ScalarLiteralType, Column]]] = field(
115        init=False, default=None
116    )

107@dataclass(frozen=True)
108class Granularity(Expression):
109    granularity: int
110
111    def validate(self) -> None:
112        _validate_int_literal("granularity", self.granularity, 1, None)

126@dataclass(frozen=True)
127class Identifier(Expression):
128    name: str
129
130    def validate(self) -> None:
131        if not isinstance(self.name, str) or not identifier_re.match(self.name):
132            raise InvalidLambdaError(f"`{self.name}` is not a valid identifier")

51@dataclass(frozen=True)
52class Join(Expression):
53    """
54    A collection of relationships that is used in the MATCH section of
55    the SnQL query. Must contain at least one Relationship, and will make
56    sure that Entity aliases are not used by different Entities.
57
58    :param relationships: The relationships in the join.
59    :type name: Sequence[Relationship]
60
61    :raises InvalidExpressionError: If two different Entities are using
62        the same alias, this will be raised.
63
64    """
65
66    relationships: Sequence[Relationship]
67
68    def get_alias_mappings(self) -> set[tuple[str, str]]:
69        aliases = set()
70        for rel in self.relationships:
71            if rel.lhs.alias is not None:
72                aliases.add((rel.lhs.alias, rel.lhs.name))
73            if rel.rhs.alias is not None:
74                aliases.add((rel.rhs.alias, rel.rhs.name))
75
76        return aliases
77
78    def get_entities(self) -> set[Entity]:
79        entities = set()
80        for rel in self.relationships:
81            entities.add(rel.lhs)
82            entities.add(rel.rhs)
83        return entities
84
85    def validate(self) -> None:
86        if not isinstance(self.relationships, (list, tuple)) or not self.relationships:
87            raise InvalidExpressionError("Join must have at least one Relationship")
88        elif not all(isinstance(x, Relationship) for x in self.relationships):
89            raise InvalidExpressionError("Join expects a list of Relationship objects")
90
91        seen: MutableMapping[str, str] = {}
92        for alias, entity in self.get_alias_mappings():
93            if alias in seen and seen[alias] != entity:
94                entities = sorted([entity, seen[alias]])
95                raise InvalidExpressionError(
96                    f"alias '{alias}' is duplicated for entities {', '.join(entities)}"
97                )
98            seen[alias] = entity