def fini_toplevel(stmt: pgast.Query,
ctx: context.CompilerContextLevel) -> None:
scan_check_ctes(stmt, ctx.env.check_ctes, ctx=ctx)
# Type rewrites go first.
if stmt.ctes is None:
stmt.ctes = []
stmt.ctes[:0] = list(ctx.type_ctes.values())
stmt.argnames = argmap = ctx.argmap
if not ctx.env.use_named_params:
# Adding unused parameters into a CTE
targets = []
for param in ctx.env.query_params:
pgparam = argmap[param.name]
if pgparam.used:
continue
targets.append(
pgast.ResTarget(val=pgast.TypeCast(
arg=pgast.ParamRef(number=pgparam.index),
type_name=pgast.TypeName(
name=pg_types.pg_type_from_ir_typeref(
param.ir_type)))))
if targets:
stmt.append_cte(
pgast.CommonTableExpr(
name="__unused_vars",
query=pgast.SelectStmt(target_list=targets)))
def scan_check_ctes(
stmt: pgast.Query,
check_ctes: List[pgast.CommonTableExpr],
*,
ctx: context.CompilerContextLevel,
) -> None:
if not check_ctes:
return
# Scan all of the check CTEs to enforce constraints that are
# checked as explicit queries and not Postgres constraints or
# triggers.
# To make sure that Postgres can't optimize the checks away, we
# reference them in the where clause of an UPDATE to a dummy
# table.
# Add a big random number, so that different queries should try to
# access different "rows" of the table, in case that matters.
base_int = random.randint(0, (1 << 60) - 1)
val: pgast.BaseExpr = pgast.NumericConstant(val=str(base_int))
for check_cte in check_ctes:
# We want the CTE to be MATERIALIZED, because otherwise
# Postgres might not fully evaluate all its columns when
# scanning it.
check_cte.materialized = True
check = pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.FuncCall(name=('count', ),
args=[pgast.Star()]), )
],
from_clause=[
relctx.rvar_for_rel(check_cte, ctx=ctx),
],
)
val = pgast.Expr(kind=pgast.ExprKind.OP,
name='+',
lexpr=val,
rexpr=check)
update_query = pgast.UpdateStmt(
targets=[
pgast.UpdateTarget(name='flag',
val=pgast.BooleanConstant(val='true'))
],
relation=pgast.RelRangeVar(
relation=pgast.Relation(schemaname='edgedb', name='_dml_dummy')),
where_clause=pgast.Expr(
kind=pgast.ExprKind.OP,
name='=',
lexpr=pgast.ColumnRef(name=['id']),
rexpr=val,
))
stmt.append_cte(
pgast.CommonTableExpr(query=update_query,
name=ctx.env.aliases.get(hint='check_scan')))