def __init__(self, uuid_columns: Set[str]) -> None:
self.__unique_uuid_columns = uuid_columns
self.__uuid_column_match = Or([String(u_col) for u_col in uuid_columns])
self.uuid_in_condition = FunctionCallMatch(
Or((String(ConditionFunctions.IN), String(ConditionFunctions.NOT_IN))),
(
self.formatted_uuid_pattern(),
Param("params", FunctionCallMatch(String("tuple"), None)),
),
)
self.uuid_condition = FunctionCallMatch(
Or(
[
String(op)
for op in FUNCTION_TO_OPERATOR
if op not in (ConditionFunctions.IN, ConditionFunctions.NOT_IN)
]
),
(
Or(
(
Param("literal_0", LiteralMatch(AnyOptionalString())),
self.formatted_uuid_pattern("_0"),
)
),
Or(
(
Param("literal_1", LiteralMatch(AnyOptionalString())),
self.formatted_uuid_pattern("_1"),
)
),
),
)
self.formatted: Optional[str] = None
def build_match(
col: str, ops: Sequence[str], param_type: Any, alias: Optional[str] = None
) -> Or[Expression]:
# The IN condition has to be checked separately since each parameter
# has to be checked individually.
alias_match = AnyOptionalString() if alias is None else String(alias)
column_match = Param("column", ColumnPattern(alias_match, String(col)))
return Or(
[
FunctionCallPattern(
Or([String(op) for op in ops]),
(column_match, Param("rhs", LiteralPattern(AnyPattern(param_type)))),
),
FunctionCallPattern(
String(ConditionFunctions.IN),
(
column_match,
Param(
"rhs",
FunctionCallPattern(
Or([String("array"), String("tuple")]),
all_parameters=LiteralPattern(AnyPattern(param_type)),
),
),
),
),
]
)
def __init__(self, column_name: str, hash_map_name: str, killswitch: str) -> None:
self.__column_name = column_name
self.__hash_map_name = hash_map_name
self.__killswitch = killswitch
# TODO: Add the support for IN conditions.
self.__optimizable_pattern = FunctionCall(
function_name=String("equals"),
parameters=(
Or(
[
mapping_pattern,
FunctionCall(
function_name=String("ifNull"),
parameters=(mapping_pattern, Literal(String(""))),
),
]
),
Param("right_hand_side", Literal(Any(str))),
),
)
self.__tag_exists_patterns = [
FunctionCall(
function_name=String("notEquals"),
parameters=(
Or(
[
mapping_pattern,
FunctionCall(
function_name=String("ifNull"),
parameters=(mapping_pattern, Literal(String(""))),
),
]
),
Param("right_hand_side", Literal(Any(str))),
),
),
FunctionCall(
function_name=String("has"),
parameters=(
ColumnMatcher(
Param(TABLE_MAPPING_PARAM, AnyOptionalString()),
Param(VALUE_COL_MAPPING_PARAM, String(f"{column_name}.key")),
),
Literal(Param(KEY_MAPPING_PARAM, Any(str))),
),
),
]
def build_match(
col: str,
ops: Sequence[str],
param_type: Any,
alias: Optional[str] = None,
key: Optional[str] = None,
) -> Or[Expression]:
# The IN condition has to be checked separately since each parameter
# has to be checked individually.
alias_match = AnyOptionalString() if alias is None else String(alias)
pattern: Union[ColumnPattern, SubscriptableReferencePattern]
if key is not None:
pattern = SubscriptableReferencePattern(table_name=alias_match,
column_name=String(col),
key=String(key))
else:
pattern = ColumnPattern(table_name=alias_match,
column_name=String(col))
column_match = Param("column", pattern)
return Or([
FunctionCallPattern(
Or([String(op) for op in ops]),
(column_match, Param("rhs", LiteralPattern(
AnyPattern(param_type)))),
),
FunctionCallPattern(
String(ConditionFunctions.IN),
(
column_match,
Param(
"rhs",
FunctionCallPattern(
Or([String("array"), String("tuple")]),
all_parameters=LiteralPattern(AnyPattern(param_type)),
),
),
),
),
])
LiteralExpr(None, None),
),
)
TABLE_MAPPING_PARAM = "table_name"
VALUE_COL_MAPPING_PARAM = "value_column"
KEY_COL_MAPPING_PARAM = "key_column"
KEY_MAPPING_PARAM = "key"
mapping_pattern = FunctionCall(
None,
String("arrayElement"),
(
Column(
None,
Param(TABLE_MAPPING_PARAM, AnyOptionalString()),
Param(VALUE_COL_MAPPING_PARAM, Any(str)),
),
FunctionCall(
None,
String("indexOf"),
(
Column(None, None, Param(KEY_COL_MAPPING_PARAM, Any(str))),
Literal(None, Param(KEY_MAPPING_PARAM, Any(str))),
),
),
),
)
# TODO: build more of these mappers.
),
(
"Single node match",
Column(OptionalString("table"), String("test_col")),
ColumnExpr("alias_we_don't_care_of", "table", "test_col"),
MatchResult(),
),
(
"Single node no match",
Column(None, String("test_col")),
ColumnExpr(None, None, "not_a_test_col"),
None,
),
(
"Matches a None table name",
Column(Param("table_name", AnyOptionalString()), None),
ColumnExpr(None, None, "not_a_test_col"),
MatchResult({"table_name": None}),
),
(
"Matches None as table name",
Column(Param("table_name", OptionalString(None)), None),
ColumnExpr(None, None, "not_a_test_col"),
MatchResult({"table_name": None}),
),
(
"Not matching a non None table",
Column(Param("table_name", OptionalString(None)), None),
ColumnExpr(None, "not None", "not_a_test_col"),
None,
),
def validate_required_conditions(
self, query: Query, alias: Optional[str] = None
) -> bool:
if not self._required_filter_columns and not self._required_time_column:
return True
condition = query.get_condition_from_ast()
top_level = get_first_level_and_conditions(condition) if condition else []
if not top_level:
return False
alias_match = AnyOptionalString() if alias is None else StringMatch(alias)
def build_match(
col: str, ops: Sequence[str], param_type: Any
) -> Or[Expression]:
# The IN condition has to be checked separately since each parameter
# has to be checked individually.
column_match = ColumnMatch(alias_match, StringMatch(col))
return Or(
[
FunctionCallMatch(
Or([StringMatch(op) for op in ops]),
(column_match, LiteralMatch(AnyMatch(param_type))),
),
FunctionCallMatch(
StringMatch(ConditionFunctions.IN),
(
column_match,
FunctionCallMatch(
Or([StringMatch("array"), StringMatch("tuple")]),
all_parameters=LiteralMatch(AnyMatch(param_type)),
),
),
),
]
)
if self._required_filter_columns:
for col in self._required_filter_columns:
match = build_match(col, [ConditionFunctions.EQ], int)
found = any(match.match(cond) for cond in top_level)
if not found:
return False
if self._required_time_column:
match = build_match(
self._required_time_column, [ConditionFunctions.EQ], datetime,
)
found = any(match.match(cond) for cond in top_level)
if found:
return True
lower, upper = get_time_range_expressions(
top_level, self._required_time_column, alias
)
if not lower or not upper:
return False
# At this point we have valid conditions. However we need to align them and
# make sure they don't exceed the max_days. Replace the conditions.
self._replace_time_condition(query, *lower, *upper)
return True
),
(
"Single node match",
Column(OptionalString("table"), String("test_col")),
ColumnExpr("alias_we_don't_care_of", "table", "test_col"),
MatchResult(),
),
(
"Single node no match",
Column(None, String("test_col")),
ColumnExpr(None, None, "not_a_test_col"),
None,
),
(
"Matches a None table name",
Column(Param("table_name", AnyOptionalString()), None),
ColumnExpr(None, None, "not_a_test_col"),
MatchResult({"table_name": None}),
),
(
"Matches None as table name",
Column(Param("table_name", OptionalString(None)), None),
ColumnExpr(None, None, "not_a_test_col"),
MatchResult({"table_name": None}),
),
(
"Not matching a non None table",
Column(Param("table_name", OptionalString(None)), None),
ColumnExpr(None, "not None", "not_a_test_col"),
None,
),
result = DATETIME_MATCH.match(exp)
if result is not None:
date_string = result.expression("date_string")
assert isinstance(date_string, Literal) # mypy
assert isinstance(date_string.value, str) # mypy
return Literal(exp.alias, parse_datetime(date_string.value))
return exp
query.transform_expressions(parse)
ARRAY_JOIN_MATCH = FunctionCallMatch(
Param("function_name", Or([StringMatch("arrayExists"), StringMatch("arrayAll")])),
(
Param("column", ColumnMatch(AnyOptionalString(), AnyMatch(str))),
Param("op", Or([LiteralMatch(StringMatch(op)) for op in OPERATOR_TO_FUNCTION])),
Param("value", AnyExpression()),
),
)
def _array_join_transformation(
query: Union[CompositeQuery[QueryEntity], LogicalQuery]
) -> None:
def parse(exp: Expression) -> Expression:
result = ARRAY_JOIN_MATCH.match(exp)
if result:
function_name = result.string("function_name")
column = result.expression("column")
assert isinstance(column, Column)
