def test_codegen_group_simple(self):
tree = self._compile_to_tree('''
select (group Issue by .status) {
name := .key.status.name,
num := count(.elements),
} order by .name
''')
child = ast_visitor.find_children(
tree,
lambda x: isinstance(x, pgast.SelectStmt) and x.group_clause,
terminate_early=True)
group_sql = pg_compiler.run_codegen(child, pretty=True)
# We want no array_agg in the group - it should just be able
# to do a count
self.assertNotIn(
"array_agg",
group_sql,
"group has unnecessary array_agg",
)
# And we want no uuid generation, which is a huge perf killer
self.assertNotIn(
"uuid_generate",
group_sql,
"group has unnecessary uuid_generate",
)
def test_common_ast_find_children(self):
node = tast.UnaryOp(
op='NamedTuple',
operand=[
('foo', tast.Constant(value=2)),
('bar', [tast.UnaryOp(op='-',
operand=tast.Constant(value=3))]),
],
)
flt = lambda n: isinstance(n, tast.Constant)
children = visitor.find_children(node, flt)
assert {x.value for x in children} == {2, 3}
def test_codegen_order_by_not_subquery_01(self):
sql = self._compile_to_tree('''
select User order by .name
''')
child = ast_visitor.find_children(
sql,
lambda x: isinstance(x, pgast.SelectStmt) and x.sort_clause,
terminate_early=True)
# Make sure that a simple order by on a property is not compiled
# as a subquery in the ORDER BY, which pg fails to use an index for.
self.assertIsInstance(
child.sort_clause[0].node,
pgast.ColumnRef,
"simple sort clause is not a column ref",
)
def test_codegen_order_by_not_subquery_02(self):
# Same as above but a bit more involved
sql = self._compile_to_tree('''
select User { z := .name ++ "!" } order by .z
''')
child = ast_visitor.find_children(
sql,
lambda x: isinstance(x, pgast.SelectStmt) and x.sort_clause,
terminate_early=True)
# Make sure that a simple order by on a property is not compiled
# as a subquery in the ORDER BY, which pg fails to use an index for.
self.assertIsInstance(
child.sort_clause[0].node,
pgast.Expr,
"simple sort clause is not a op expr",
)
def _fixup_materialized_sets(ir: irast.Base, *,
ctx: context.ContextLevel) -> None:
# Make sure that all materialized sets have their views compiled
flt = lambda n: isinstance(n, irast.Stmt)
children: List[irast.Stmt] = ast_visitor.find_children(ir, flt)
for nobe in ctx.source_map.values():
if nobe.irexpr:
children += ast_visitor.find_children(nobe.irexpr, flt)
for stmt in set(children):
if not stmt.materialized_sets:
continue
for key in list(stmt.materialized_sets):
mat_set = stmt.materialized_sets[key]
assert not mat_set.finalized
if len(mat_set.uses) <= 1:
del stmt.materialized_sets[key]
continue
# Find the right set to compile by looking for the one
# with a matching rptr.
if not mat_set.materialized:
for use in mat_set.use_sets:
if use.rptr and use.rptr.source == stmt.result:
mat_set.materialized = use
break
else:
raise AssertionError(
f"couldn't find the source for {mat_set.uses} on "
f"{stmt.result}!")
ir_set = mat_set.materialized
assert ir_set.path_scope_id is not None
new_scope = ctx.env.scope_tree_nodes[ir_set.path_scope_id]
assert new_scope.parent
parent = new_scope.parent
good_reason = False
for x in mat_set.reason:
if isinstance(x, irast.MaterializeVolatile):
good_reason = True
elif isinstance(x, irast.MaterializeVisible):
# If any of the bindings that the set uses are *visible*
# at the binding point, we need to materialize, to make
# sure that things get correlated properly. If it's not
# visible, then it's just being used internally and we
# don't need any special work.
if any(parent.is_visible(b) for b in x.paths):
good_reason = True
if not good_reason:
del stmt.materialized_sets[key]
continue
# Compile the view shapes in the set
with ctx.new() as subctx:
subctx.implicit_tid_in_shapes = False
subctx.implicit_tname_in_shapes = False
subctx.path_scope = new_scope
viewgen.compile_view_shapes(ir_set, ctx=subctx)
assert (not any(use.src_path() for use in mat_set.uses)
or mat_set.materialized.rptr
), f"materialized ptr {mat_set.uses} missing rptr"
mat_set.finalized = True
def _find_visible_binding_refs(ir: irast.Base, *,
ctx: context.ContextLevel) -> List[irast.Set]:
flt = lambda n: isinstance(n, irast.Set) and n.is_visible_binding_ref
children: List[irast.Set] = ast_visitor.find_children(ir, flt)
return children
def _fixup_materialized_sets(ir: irast.Base, *,
ctx: context.ContextLevel) -> List[irast.Set]:
# Make sure that all materialized sets have their views compiled
flt = lambda n: isinstance(n, irast.Stmt)
children: List[irast.Stmt] = ast_visitor.find_children(ir, flt)
for nobe in ctx.source_map.values():
if nobe.irexpr:
children += ast_visitor.find_children(nobe.irexpr, flt)
to_clear = []
for stmt in ordered.OrderedSet(children):
if not stmt.materialized_sets:
continue
for key in list(stmt.materialized_sets):
mat_set = stmt.materialized_sets[key]
assert not mat_set.finalized
if len(mat_set.uses) <= 1:
del stmt.materialized_sets[key]
continue
ir_set = mat_set.materialized
assert ir_set.path_scope_id is not None
new_scope = ctx.env.scope_tree_nodes[ir_set.path_scope_id]
assert new_scope.parent
parent = new_scope.parent
good_reason = False
for x in mat_set.reason:
if isinstance(x, irast.MaterializeVolatile):
good_reason = True
elif isinstance(x, irast.MaterializeVisible):
# If any of the bindings that the set uses are
# *visible* at the definition point and *not
# visible* from at least one use point, we need to
# materialize, to make sure that the use site sees
# the same value for the binding as the definition
# point. If it's not visible, then it's just being
# used internally and we don't need any special
# work.
use_scopes = [
ctx.env.scope_tree_nodes.get(x.path_scope_id)
if x.path_scope_id is not None else None
for x in mat_set.use_sets
]
for b, _ in x.sets:
if parent.is_visible(b, allow_group=True) and not all(
use_scope and use_scope.parent and use_scope.
parent.is_visible(b, allow_group=True)
for use_scope in use_scopes):
good_reason = True
break
if not good_reason:
del stmt.materialized_sets[key]
continue
# Compile the view shapes in the set
with ctx.new() as subctx:
subctx.implicit_tid_in_shapes = False
subctx.implicit_tname_in_shapes = False
subctx.path_scope = new_scope
viewgen.late_compile_view_shapes(ir_set, ctx=subctx)
for use_set in mat_set.use_sets:
if use_set != mat_set.materialized:
use_set.is_materialized_ref = True
# XXX: Deleting it on linkprops breaks a bunch of
# linkprop related DML...
if not use_set.path_id.is_linkprop_path():
to_clear.append(use_set)
assert (not any(use.src_path() for use in mat_set.uses)
or mat_set.materialized.rptr
), f"materialized ptr {mat_set.uses} missing rptr"
mat_set.finalized = True
return to_clear