def validate_parameters(self, exp: FunctionCall) -> None:
validator = SignatureValidator([])
try:
validator.validate(exp.parameters, self.__columnset)
except InvalidFunctionCall as err:
raise InvalidExpressionException(
exp,
f"Illegal function call to {exp.function_name}: {str(err)}"
) from err
def validate_parameters(self, exp: FunctionCall,
entity: QueryEntity) -> None:
validator = SignatureValidator([])
try:
validator.validate(exp.function_name, exp.parameters, entity)
except InvalidFunctionCall as err:
raise InvalidExpressionException.from_args(
exp,
f"Illegal function call to {exp.function_name}: {str(err)}",
should_report=False,
) from err
def __init__(
self,
name: str,
signature: Sequence[Tuple[str, ParamType]],
body: Expression,
) -> None:
self.__function_name = name
self.__param_names: Sequence[str] = []
param_types: Sequence[ParamType] = []
if len(signature) > 0:
self.__param_names, param_types = zip(*signature)
self.__body = body
self.__validator = SignatureValidator(param_types)
def test_like_validator(
expressions: Sequence[Expression],
expected_types: Sequence[ParamType],
extra_param: bool,
should_raise: bool,
) -> None:
schema = ColumnSet([
("event_id", String()),
("level", Nullable(String())),
("str_col", String()),
("timestamp", DateTime()),
("received", NullableOld(DateTime())),
])
validator = SignatureValidator(expected_types, extra_param)
if should_raise:
with pytest.raises(InvalidFunctionCall):
validator.validate(expressions, schema)
else:
validator.validate(expressions, schema)
def test_like_validator(
expressions: Sequence[Expression],
expected_types: Sequence[ParamType],
extra_param: bool,
should_raise: bool,
) -> None:
entity = QueryEntity(
EntityKey.EVENTS,
ColumnSet([
("event_id", String()),
("level", String(Modifiers(nullable=True))),
("str_col", String()),
("timestamp", DateTime()),
("received", DateTime(Modifiers(nullable=True))),
]),
)
func_name = "like"
validator = SignatureValidator(expected_types, extra_param)
if should_raise:
with pytest.raises(InvalidFunctionCall):
validator.validate(func_name, expressions, entity)
else:
validator.validate(func_name, expressions, entity)
class CustomFunction(QueryProcessor):
"""
Defines a custom snuba function.
The custom function has a name, a signature in the form of a list of
parameter names and types, an expanded expression which is the body
of the custom function and the dataset abstract schema for validation.
The custom function is invoked in a query as a standard snuba function.
Example:
CustomFunction(name="power_two", param_names=[value], body="value * value", schema)
would transform
`["power_two", [["f", ["something"]]], "alias"]`
into
`f(something) * f(something) AS alias`
Would raise InvalidCustomFunctionCall if a custom function is invoked with
the wrong number of parameters or if the parameters are of the wrong type
according to the SignatureValidator.
The validation has the same limitations of SignatureValidation in that it only deals
with columns: no literals or complex functions validation.
We need to pass the dataset abstract schema to the processor constructor
because the schema is not populated in the query object when dataset processors
are executed.
TODO: Assign the abstract dataset schema to the query object right after parsing.
"""
def __init__(
self,
dataset_schema: ColumnSet,
name: str,
signature: Sequence[Tuple[str, ParamType]],
body: Expression,
) -> None:
self.__dataset_schema = dataset_schema
self.__function_name = name
self.__param_names: Sequence[str] = []
param_types: Sequence[ParamType] = []
if len(signature) > 0:
self.__param_names, param_types = zip(*signature)
self.__body = body
self.__validator = SignatureValidator(param_types)
def process_query(self, query: Query,
request_settings: RequestSettings) -> None:
def apply_function(expression: Expression) -> Expression:
if (isinstance(expression, FunctionCall)
and expression.function_name == self.__function_name):
try:
self.__validator.validate(expression.parameters,
self.__dataset_schema)
except InvalidFunctionCall as exception:
raise InvalidCustomFunctionCall(
expression,
f"Illegal call to function {expression.function_name}: {str(exception)}",
) from exception
resolved_params = {
name: expression
for (name, expression
) in zip(self.__param_names, expression.parameters)
}
ret = replace_in_expression(self.__body, resolved_params)
return replace(ret, alias=expression.alias)
else:
return expression
query.transform_expressions(apply_function)
from snuba.clickhouse.columns import Array, String
from snuba.datasets.entity import Entity
from snuba.query.exceptions import InvalidExpressionException
from snuba.query.expressions import Expression, FunctionCall
from snuba.query.parser.validation import ExpressionValidator
from snuba.query.validation import FunctionCallValidator, InvalidFunctionCall
from snuba.query.validation.signature import Any, Column, SignatureValidator
logger = logging.getLogger(__name__)
default_validators: Mapping[str, FunctionCallValidator] = {
# like and notLike need to take care of Arrays as well since
# Arrays are exploded into strings if they are part of the arrayjoin
# clause.
# TODO: provide a more restrictive support for arrayjoin.
"like": SignatureValidator([Column({Array, String}),
Any()]),
"notLike": SignatureValidator([Column({Array, String}),
Any()]),
}
class FunctionCallsValidator(ExpressionValidator):
"""
Applies all function validators on the provided expression.
The individual function validators are divided in two mappings:
a default one applied to all queries and one a mapping per dataset.
"""
def validate(self, exp: Expression, entity: Entity) -> None:
if not isinstance(exp, FunctionCall):
return