def test_query_data_source() -> None:
"""
Tests using the Query as a data source
"""
query = Query(
Table("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression(
"col1", Column(alias="col1", table_name=None, column_name="col1")
),
SelectedExpression(
"some_func",
FunctionCall(
"some_func",
"f",
(Column(alias="col1", table_name=None, column_name="col1"),),
),
),
SelectedExpression(
None, Column(alias="col2", table_name=None, column_name="col2")
),
],
)
assert query.get_columns() == ColumnSet(
[("col1", Any()), ("some_func", Any()), ("_invalid_alias_2", Any())]
)
def test_nested_query() -> None:
"""
Simply builds a nested query.
"""
nested = LogicalQuery(
Entity(EntityKey.EVENTS, ColumnSet([("event_id", String())])),
selected_columns=[
SelectedExpression("string_evt_id",
Column("string_evt_id", None, "event_id"))
],
)
composite = CompositeQuery(
from_clause=nested,
selected_columns=[
SelectedExpression("output", Column("output", None,
"string_evt_id"))
],
)
# The iterator methods on the composite query do not descend into
# the nested query
assert composite.get_all_ast_referenced_columns() == {
Column("output", None, "string_evt_id")
}
# The schema of the nested query is the selected clause of that query.
assert composite.get_from_clause().get_columns() == ColumnSet([
("string_evt_id", Any())
])
def test_complex_joins() -> None:
e = Entity(key=EntityKey.EVENTS, schema=ERRORS_SCHEMA)
node_err = IndividualNode(alias="err", data_source=e)
g = Entity(key=EntityKey.GROUPEDMESSAGES, schema=GROUPS_SCHEMA)
node_group = IndividualNode(alias="groups", data_source=g)
a = Entity(key=EntityKey.GROUPASSIGNEE, schema=GROUPS_ASSIGNEE)
query = Query(
from_clause=a,
selected_columns=[
SelectedExpression("id", Column("id", None, "id")),
SelectedExpression("assigned_user",
Column("assigned_user", None, "user")),
],
)
node_query = IndividualNode(alias="assignee", data_source=query)
join = JoinClause(
left_node=JoinClause(
left_node=node_err,
right_node=node_group,
keys=[
JoinCondition(
left=JoinConditionExpression("err", "group_id"),
right=JoinConditionExpression("groups", "id"),
)
],
join_type=JoinType.INNER,
),
right_node=node_query,
keys=[
JoinCondition(
left=JoinConditionExpression("err", "group_id"),
right=JoinConditionExpression("assignee", "id"),
)
],
join_type=JoinType.INNER,
)
assert join.get_column_sets() == {
"err": ERRORS_SCHEMA,
"assignee": ColumnSet([("id", Any()), ("assigned_user", Any())]),
"groups": GROUPS_SCHEMA,
}
def get_columns(self) -> ColumnSet:
"""
From the DataSource class. It returns the schema exposed by this
query when used as a Data Source for another query.
"""
ret = []
for index, selected_col in enumerate(self.__selected_columns):
name = selected_col.name
# TODO: Make the type of the columns precise onn the type
# when possible. It may become useful for query validation
# but it would be best effort because we cannot infer the
# type of complex expressions.
ret.append(
(name, Any()) if name is not None
# This should never happen for nested queries.
# Though the type of the name oof a selected column is
# still optional soo we need to fix that first.
else (f"_invalid_alias_{index}", Any()))
return ColumnSet(ret)