def wrap_script_stmt(
stmt: pgast.SelectStmt,
ir_set: irast.Set,
*,
env: context.Environment,
) -> pgast.SelectStmt:
subrvar = pgast.RangeSubselect(
subquery=stmt, alias=pgast.Alias(aliasname=env.aliases.get('aggw')))
stmt_res = stmt.target_list[0]
if stmt_res.name is None:
stmt_res = stmt.target_list[0] = pgast.ResTarget(
name=env.aliases.get('v'),
val=stmt_res.val,
)
count_val = pgast.FuncCall(name=('count', ),
args=[pgast.ColumnRef(name=[stmt_res.name])]),
result = pgast.SelectStmt(target_list=[pgast.ResTarget(val=count_val, )],
from_clause=[subrvar])
result.ctes = stmt.ctes
result.argnames = stmt.argnames
stmt.ctes = []
return result
def aggregate_json_output(stmt: pgast.Query, ir_set: irast.Set, *,
env: context.Environment) -> pgast.Query:
subrvar = pgast.RangeSubselect(
subquery=stmt, alias=pgast.Alias(aliasname=env.aliases.get('aggw')))
stmt_res = stmt.target_list[0]
if stmt_res.name is None:
stmt_res = stmt.target_list[0] = pgast.ResTarget(
name=env.aliases.get('v'),
val=stmt_res.val,
)
new_val = pgast.FuncCall(name=_get_json_func('agg', env=env),
args=[pgast.ColumnRef(name=[stmt_res.name])])
new_val = pgast.CoalesceExpr(
args=[new_val, pgast.StringConstant(val='[]')])
result = pgast.SelectStmt(target_list=[pgast.ResTarget(val=new_val)],
from_clause=[subrvar])
result.ctes = stmt.ctes
result.argnames = stmt.argnames
stmt.ctes = []
return result
def top_output_as_config_op(
ir_set: irast.Set,
stmt: pgast.SelectStmt, *,
env: context.Environment) -> pgast.Query:
assert isinstance(ir_set.expr, irast.ConfigCommand)
if ir_set.expr.scope is qltypes.ConfigScope.SYSTEM:
alias = env.aliases.get('cfg')
subrvar = pgast.RangeSubselect(
subquery=stmt,
alias=pgast.Alias(
aliasname=alias,
)
)
stmt_res = stmt.target_list[0]
if stmt_res.name is None:
stmt_res = stmt.target_list[0] = pgast.ResTarget(
name=env.aliases.get('v'),
val=stmt_res.val,
)
result_row = pgast.RowExpr(
args=[
pgast.StringConstant(val='ADD'),
pgast.StringConstant(val=str(ir_set.expr.scope)),
pgast.StringConstant(val=ir_set.expr.name),
pgast.ColumnRef(name=[stmt_res.name]),
]
)
array = pgast.FuncCall(
name=('jsonb_build_array',),
args=result_row.args,
null_safe=True,
ser_safe=True,
)
result = pgast.SelectStmt(
target_list=[
pgast.ResTarget(
val=array,
),
],
from_clause=[
subrvar,
],
)
result.ctes = stmt.ctes
result.argnames = stmt.argnames
stmt.ctes = []
return result
else:
raise errors.InternalServerError(
f'CONFIGURE {ir_set.expr.scope} INSERT is not supported')
def range_for_ptrref(ptrref: irast.BasePointerRef,
*,
include_overlays: bool = True,
only_self: bool = False,
env: context.Environment) -> pgast.BaseRangeVar:
""""Return a Range subclass corresponding to a given ptr step.
The return value may potentially be a UNION of all tables
corresponding to a set of specialized links computed from the given
`ptrref` taking source inheritance into account.
"""
tgt_col = pgtypes.get_ptrref_storage_info(ptrref,
resolve_type=False,
link_bias=True).column_name
cols = ['source', tgt_col]
set_ops = []
if only_self:
ptrrefs = {ptrref}
else:
ptrrefs = {ptrref} | ptrref.descendants
for src_ptrref in ptrrefs:
table = table_from_ptrref(src_ptrref, env=env)
qry = pgast.SelectStmt()
qry.from_clause.append(table)
qry.rptr_rvar = table
# Make sure all property references are pulled up properly
for colname in cols:
selexpr = pgast.ColumnRef(name=[table.alias.aliasname, colname])
qry.target_list.append(pgast.ResTarget(val=selexpr, name=colname))
set_ops.append(('union', qry))
overlays = env.rel_overlays.get(src_ptrref.shortname)
if overlays and include_overlays:
for op, cte in overlays:
rvar = pgast.RangeVar(
relation=cte,
alias=pgast.Alias(aliasname=env.aliases.get(cte.name)))
qry = pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.ColumnRef(name=[col]))
for col in cols
],
from_clause=[rvar],
)
set_ops.append((op, qry))
rvar = range_from_queryset(set_ops, ptrref.shortname, env=env)
return rvar
def top_output_as_config_op(
ir_set: irast.Set,
stmt: pgast.Query, *,
env: context.Environment) -> pgast.Query:
if ir_set.expr.system:
alias = env.aliases.get('cfg')
subrvar = pgast.RangeSubselect(
subquery=stmt,
alias=pgast.Alias(
aliasname=alias,
)
)
stmt_res = stmt.target_list[0]
if stmt_res.name is None:
stmt_res = stmt.target_list[0] = pgast.ResTarget(
name=env.aliases.get('v'),
val=stmt_res.val,
)
result_row = pgast.RowExpr(
args=[
pgast.StringConstant(val='ADD'),
pgast.StringConstant(
val='SYSTEM' if ir_set.expr.system else 'SESSION'),
pgast.StringConstant(val=ir_set.expr.name),
pgast.ColumnRef(name=[stmt_res.name]),
]
)
result = pgast.FuncCall(
name=('jsonb_build_array',),
args=result_row.args,
null_safe=True,
ser_safe=True,
)
return pgast.SelectStmt(
target_list=[
pgast.ResTarget(
val=result,
),
],
from_clause=[
subrvar,
],
)
else:
raise errors.InternalServerError(
'CONFIGURE SESSION INSERT is not supported')
def wrap_script_stmt(
stmt: pgast.SelectStmt,
*,
suppress_all_output: bool = False,
env: context.Environment,
) -> pgast.SelectStmt:
subrvar = pgast.RangeSubselect(
subquery=stmt, alias=pgast.Alias(aliasname=env.aliases.get('aggw')))
stmt_res = stmt.target_list[0]
if stmt_res.name is None:
stmt_res = stmt.target_list[0] = pgast.ResTarget(
name=env.aliases.get('v'),
val=stmt_res.val,
)
assert stmt_res.name is not None
count_val = pgast.FuncCall(name=('count', ),
args=[pgast.ColumnRef(name=[stmt_res.name])])
result = pgast.SelectStmt(target_list=[
pgast.ResTarget(
val=count_val,
name=stmt_res.name,
),
],
from_clause=[
subrvar,
])
if suppress_all_output:
subrvar = pgast.RangeSubselect(
subquery=result, alias=pgast.Alias(aliasname=env.aliases.get('q')))
result = pgast.SelectStmt(
target_list=[],
from_clause=[
subrvar,
],
where_clause=pgast.NullTest(arg=pgast.ColumnRef(
name=[subrvar.alias.aliasname, stmt_res.name], ), ),
)
result.ctes = stmt.ctes
result.argnames = stmt.argnames
stmt.ctes = []
return result
def get_packed_path_var(
rvar: pgast.PathRangeVar, path_id: irast.PathId, aspect: str, *,
env: context.Environment) -> Tuple[pgast.OutputVar, bool]:
res = maybe_get_rvar_path_packed_output(
rvar, path_id, aspect, env=env)
if res:
return res
query = rvar.query
assert isinstance(query, pgast.Query)
rel_rvar = get_path_rvar(
query, path_id, flavor='packed', aspect=aspect, env=env)
# XXX: some duplication of path_output
ref, multi = get_packed_path_var(rel_rvar, path_id, aspect, env=env)
alias = get_path_output_alias(path_id, aspect, env=env)
restarget = pgast.ResTarget(
name=alias, val=ref, ser_safe=getattr(ref, 'ser_safe', False))
query.target_list.append(restarget)
nullable = is_nullable(ref, env=env)
optional = None
if isinstance(ref, pgast.ColumnRef):
optional = ref.optional
result = pgast.ColumnRef(
name=[alias], nullable=nullable, optional=optional)
_put_path_output_var(query, path_id, aspect, result, env=env)
return result, multi
def fini_stmt(stmt: pgast.Query, ctx: context.CompilerContextLevel,
parent_ctx: context.CompilerContextLevel) -> None:
if stmt is ctx.toplevel_stmt:
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:
if param.name in argmap:
continue
if param.name.isdecimal():
idx = int(param.name) + 1
else:
idx = len(argmap) + 1
argmap[param.name] = pgast.Param(
index=idx,
required=param.required,
)
targets.append(
pgast.ResTarget(val=pgast.TypeCast(
arg=pgast.ParamRef(number=idx),
type_name=pgast.TypeName(
name=pg_types.pg_type_from_ir_typeref(
param.ir_type)))))
if targets:
ctx.toplevel_stmt.ctes.append(
pgast.CommonTableExpr(
name="__unused_vars",
query=pgast.SelectStmt(target_list=targets)))
def get_path_serialized_output(rel: pgast.Query, path_id: irast.PathId, *,
env: context.Environment) -> pgast.OutputVar:
# Serialized output is a special case, we don't
# want this behaviour to be recursive, so it
# must be kept outside of get_path_output() generic.
aspect = 'serialized'
result = rel.path_outputs.get((path_id, aspect))
if result is not None:
return result
ref = get_path_serialized_or_value_var(rel, path_id, env=env)
refexpr = output.serialize_expr(ref, path_id=path_id, env=env)
alias = get_path_output_alias(path_id, aspect, env=env)
restarget = pgast.ResTarget(name=alias, val=refexpr, ser_safe=True)
rel.target_list.append(restarget)
result = pgast.ColumnRef(name=[alias],
nullable=refexpr.nullable,
ser_safe=True)
_put_path_output_var(rel, path_id, aspect, result, env=env)
return result
def _get_rel_object_id_output(rel: pgast.BaseRelation,
path_id: irast.PathId,
*,
aspect: str,
ptr_info: typing.Optional[
pg_types.PointerStorageInfo] = None,
env: context.Environment) -> pgast.OutputVar:
var = rel.path_outputs.get((path_id, aspect))
if var is not None:
return var
if isinstance(rel, pgast.NullRelation):
name = env.aliases.get('id')
val = pgast.TypeCast(arg=pgast.NullConstant(),
type_name=pgast.TypeName(name=('uuid', ), ))
rel.target_list.append(pgast.ResTarget(name=name, val=val))
result = pgast.ColumnRef(name=[name], nullable=True)
else:
result = pgast.ColumnRef(name=['id'], nullable=False)
_put_path_output_var(rel, path_id, aspect, result, env=env)
return result
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 _cast_operands(
lexpr: Optional[pgast.BaseExpr],
rexpr: Optional[pgast.BaseExpr],
sql_types: Tuple[str, ...],
) -> Tuple[Optional[pgast.BaseExpr], Optional[pgast.BaseExpr]]:
if lexpr is not None:
lexpr = pgast.TypeCast(arg=lexpr,
type_name=pgast.TypeName(name=(sql_types[0], )))
if rexpr is not None:
rexpr_qry = None
if (isinstance(rexpr, pgast.SubLink)
and isinstance(rexpr.expr, pgast.SelectStmt)):
rexpr_qry = rexpr.expr
elif isinstance(rexpr, pgast.SelectStmt):
rexpr_qry = rexpr
if rexpr_qry is not None:
# Handle cases like foo <op> ANY (SELECT) and
# foo <OP> (SELECT).
rexpr_qry.target_list[0] = pgast.ResTarget(
name=rexpr_qry.target_list[0].name,
val=pgast.TypeCast(
arg=rexpr_qry.target_list[0].val,
type_name=pgast.TypeName(name=(sql_types[1], ))))
else:
rexpr = pgast.TypeCast(
arg=rexpr, type_name=pgast.TypeName(name=(sql_types[1], )))
return lexpr, rexpr
def top_output_as_value(
stmt: pgast.SelectStmt,
ir_set: irast.Set, *,
env: context.Environment) -> pgast.SelectStmt:
"""Finalize output serialization on the top level."""
if (env.output_format is context.OutputFormat.JSON and
not env.expected_cardinality_one):
# For JSON we just want to aggregate the whole thing
# into a JSON array.
return aggregate_json_output(stmt, ir_set, env=env)
elif (env.output_format is context.OutputFormat.NATIVE and
env.explicit_top_cast is not None):
typecast = pgast.TypeCast(
arg=stmt.target_list[0].val,
type_name=pgast.TypeName(
name=pgtypes.pg_type_from_ir_typeref(
env.explicit_top_cast,
persistent_tuples=True,
),
),
)
stmt.target_list[0] = pgast.ResTarget(
name=env.aliases.get('v'),
val=typecast,
)
return stmt
else:
# JSON_ELEMENTS and BINARY don't require any wrapping
return stmt
def get_path_output_or_null(
rel: pgast.Query, path_id: irast.PathId, *,
aspect: str, env: context.Environment) -> \
typing.Tuple[pgast.OutputVar, bool]:
path_id = map_path_id(path_id, rel.view_path_id_map)
ref = maybe_get_path_output(rel, path_id, aspect=aspect, env=env)
if ref is not None:
return ref, False
alt_aspect = get_less_specific_aspect(path_id, aspect)
if alt_aspect is not None:
ref = maybe_get_path_output(rel, path_id, aspect=alt_aspect, env=env)
if ref is not None:
_put_path_output_var(rel, path_id, aspect, ref, env=env)
return ref, False
alias = env.aliases.get('null')
restarget = pgast.ResTarget(name=alias, val=pgast.NullConstant())
rel.target_list.append(restarget)
ref = pgast.ColumnRef(name=[alias], nullable=True)
_put_path_output_var(rel, path_id, aspect, ref, env=env)
return ref, True
def get_volatility_ref(
path_id: irast.PathId, stmt: pgast.SelectStmt, *,
ctx: context.CompilerContextLevel) -> Optional[pgast.BaseExpr]:
"""Produce an appropriate volatility_ref from a path_id."""
ref: Optional[pgast.BaseExpr] = relctx.maybe_get_path_var(
stmt, path_id, aspect='identity', ctx=ctx)
if not ref:
rvar = relctx.maybe_get_path_rvar(stmt,
path_id,
aspect='value',
ctx=ctx)
if rvar and isinstance(rvar.query, pgast.ReturningQuery):
# If we are selecting from a nontrivial subquery, manually
# add a volatility ref based on row_number. We do it
# manually because the row number isn't /really/ the
# identity of the set.
name = ctx.env.aliases.get('key')
rvar.query.target_list.append(
pgast.ResTarget(name=name,
val=pgast.FuncCall(name=('row_number', ),
args=[],
over=pgast.WindowDef())))
ref = pgast.ColumnRef(name=[rvar.alias.aliasname, name])
else:
ref = relctx.maybe_get_path_var(stmt,
path_id,
aspect='value',
ctx=ctx)
return ref
def unnamed_tuple_as_json_object(expr, *, styperef, env):
vals = []
if styperef.in_schema:
for el_idx, el_type in enumerate(styperef.subtypes):
val = pgast.Indirection(
arg=expr,
indirection=[
pgast.ColumnRef(name=[str(el_idx)], ),
],
)
if irtyputils.is_collection(el_type):
val = coll_as_json_object(val, styperef=el_type, env=env)
vals.append(val)
return pgast.FuncCall(name=_get_json_func('build_array', env=env),
args=vals,
null_safe=True,
ser_safe=True,
nullable=expr.nullable)
else:
coldeflist = []
for el_idx, el_type in enumerate(styperef.subtypes):
coldeflist.append(
pgast.ColumnDef(
name=str(el_idx),
typename=pgast.TypeName(
name=pgtypes.pg_type_from_ir_typeref(el_type), ),
))
val = pgast.ColumnRef(name=[str(el_idx)])
if irtyputils.is_collection(el_type):
val = coll_as_json_object(val, styperef=el_type, env=env)
vals.append(val)
res = pgast.FuncCall(name=_get_json_func('build_array', env=env),
args=vals,
null_safe=True,
ser_safe=True,
nullable=expr.nullable)
return pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=res, ),
],
from_clause=[
pgast.RangeFunction(functions=[
pgast.FuncCall(
name=('unnest', ),
args=[pgast.ArrayExpr(elements=[expr], )],
coldeflist=coldeflist,
)
])
])
def tuple_getattr(tuple_val, tuple_typeref, attr):
ttypes = []
pgtypes = []
for i, st in enumerate(tuple_typeref.subtypes):
pgtype = pg_types.pg_type_from_ir_typeref(st)
pgtypes.append(pgtype)
if st.element_name:
ttypes.append(st.element_name)
else:
ttypes.append(str(i))
index = ttypes.index(attr)
if tuple_typeref.in_schema:
set_expr = pgast.Indirection(
arg=tuple_val,
indirection=[
pgast.ColumnRef(
name=[attr],
),
],
)
else:
set_expr = pgast.SelectStmt(
target_list=[
pgast.ResTarget(
val=pgast.ColumnRef(
name=[str(index)],
),
),
],
from_clause=[
pgast.RangeFunction(
functions=[
pgast.FuncCall(
name=('unnest',),
args=[
pgast.ArrayExpr(
elements=[tuple_val],
)
],
coldeflist=[
pgast.ColumnDef(
name=str(i),
typename=pgast.TypeName(
name=t
)
)
for i, t in enumerate(pgtypes)
]
)
]
)
]
)
return set_expr
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')))
def new_root_rvar(ir_set: irast.Set,
*,
typeref: Optional[irast.TypeRef] = None,
ctx: context.CompilerContextLevel) -> pgast.PathRangeVar:
if not ir_set.path_id.is_objtype_path():
raise ValueError('cannot create root rvar for non-object path')
if typeref is None:
typeref = ir_set.typeref
if typeref.intersection:
wrapper = pgast.SelectStmt()
for component in typeref.intersection:
component_rvar = new_root_rvar(ir_set, typeref=component, ctx=ctx)
pathctx.put_rvar_path_bond(component_rvar, ir_set.path_id)
include_rvar(wrapper, component_rvar, ir_set.path_id, ctx=ctx)
return rvar_for_rel(wrapper, ctx=ctx)
dml_source = irutils.get_nearest_dml_stmt(ir_set)
set_rvar = range_for_typeref(typeref,
ir_set.path_id,
dml_source=dml_source,
ctx=ctx)
pathctx.put_rvar_path_bond(set_rvar, ir_set.path_id)
set_rvar.query.value_scope.add(ir_set.path_id)
if ir_set.rptr and ir_set.rptr.is_inbound:
ptrref = ir_set.rptr.ptrref
ptr_info = pg_types.get_ptrref_storage_info(ptrref,
resolve_type=False,
link_bias=False)
if ptr_info.table_type == 'ObjectType':
# Inline link
prefix_path_id = ir_set.path_id.src_path()
assert prefix_path_id is not None, 'expected a path'
rref = pgast.ColumnRef(name=[ptr_info.column_name],
nullable=not ptrref.required)
pathctx.put_rvar_path_bond(set_rvar, prefix_path_id)
pathctx.put_rvar_path_output(set_rvar,
prefix_path_id,
aspect='identity',
var=rref,
env=ctx.env)
if astutils.is_set_op_query(set_rvar.query):
assert isinstance(set_rvar.query, pgast.SelectStmt)
astutils.for_each_query_in_set(
set_rvar.query, lambda qry: qry.target_list.append(
pgast.ResTarget(
val=rref,
name=ptr_info.column_name,
)))
return set_rvar
def _get_volatility_ref() -> Optional[pgast.BaseExpr]:
nonlocal vol_ref
if vol_ref:
return vol_ref
name = ctx.env.aliases.get('key')
grouprel.target_list.append(
pgast.ResTarget(name=name,
val=pgast.FuncCall(name=('row_number', ),
args=[],
over=pgast.WindowDef())))
vol_ref = pgast.ColumnRef(name=[group_rvar.alias.aliasname, name])
return vol_ref
def _pull_col(comp_qry: pgast.Query) -> None:
rvar = pathctx.get_path_rvar(comp_qry,
path_id,
aspect='source',
env=ctx.env)
typeref = rvar.typeref
assert typeref is not None
comp_ptrref = ptr_ref_map[typeref.id]
comp_pi = pg_types.get_ptrref_storage_info(
comp_ptrref, resolve_type=False, link_bias=False)
comp_qry.target_list.append(
pgast.ResTarget(
val=pgast.ColumnRef(name=[comp_pi.column_name]),
name=ptr_info.column_name,
))
def get_path_output_or_null(
rel: pgast.Query, path_id: irast.PathId, *,
disable_output_fusion: bool=False,
aspect: str, env: context.Environment) -> \
Tuple[pgast.OutputVar, bool]:
path_id = map_path_id(path_id, rel.view_path_id_map)
ref = maybe_get_path_output(
rel, path_id,
disable_output_fusion=disable_output_fusion,
aspect=aspect, env=env)
if ref is not None:
return ref, False
alt_aspect = get_less_specific_aspect(path_id, aspect)
if alt_aspect is not None:
# If disable_output_fusion is true, we need to be careful
# to not reuse an existing column
if disable_output_fusion:
preexisting = rel.path_outputs.pop((path_id, alt_aspect), None)
ref = maybe_get_path_output(
rel, path_id,
disable_output_fusion=disable_output_fusion,
aspect=alt_aspect, env=env)
if disable_output_fusion:
# Put back the path_output to whatever it was before
if not preexisting:
rel.path_outputs.pop((path_id, alt_aspect), None)
else:
rel.path_outputs[(path_id, alt_aspect)] = preexisting
if ref is not None:
_put_path_output_var(rel, path_id, aspect, ref, env=env)
return ref, False
alias = env.aliases.get('null')
restarget = pgast.ResTarget(
name=alias,
val=pgast.NullConstant())
rel.target_list.append(restarget)
ref = pgast.ColumnRef(name=[alias], nullable=True)
_put_path_output_var(rel, path_id, aspect, ref, env=env)
return ref, True
def get_path_serialized_output(
rel: pgast.Query, path_id: irast.PathId, *,
env: context.Environment) -> pgast.OutputVar:
# Serialized output is a special case, we don't
# want this behaviour to be recursive, so it
# must be kept outside of get_path_output() generic.
aspect = 'serialized'
path_id = map_path_id(path_id, rel.view_path_id_map)
result = rel.path_outputs.get((path_id, aspect))
if result is not None:
return result
ref = get_path_serialized_or_value_var(rel, path_id, env=env)
if (
isinstance(ref, pgast.TupleVarBase)
and not isinstance(ref, pgast.TupleVar)
):
elements = []
for el in ref.elements:
assert el.path_id is not None
val = get_path_serialized_or_value_var(rel, el.path_id, env=env)
elements.append(
pgast.TupleElement(
path_id=el.path_id, name=el.name, val=val))
ref = pgast.TupleVar(
elements,
named=ref.named,
typeref=ref.typeref,
)
refexpr = output.serialize_expr(ref, path_id=path_id, env=env)
alias = get_path_output_alias(path_id, aspect, env=env)
restarget = pgast.ResTarget(name=alias, val=refexpr, ser_safe=True)
rel.target_list.append(restarget)
result = pgast.ColumnRef(
name=[alias], nullable=refexpr.nullable, ser_safe=True)
_put_path_output_var(rel, path_id, aspect, result, env=env)
return result
def new_root_rvar(ir_set: irast.Set,
*,
typeref: typing.Optional[irast.TypeRef] = None,
ctx: context.CompilerContextLevel) -> pgast.BaseRangeVar:
if not ir_set.path_id.is_objtype_path():
raise ValueError('cannot create root rvar for non-object path')
if typeref is None:
typeref = ir_set.typeref
set_rvar = dbobj.range_for_typeref(typeref, ir_set.path_id, env=ctx.env)
pathctx.put_rvar_path_bond(set_rvar, ir_set.path_id)
set_rvar.value_scope.add(ir_set.path_id)
if ir_set.rptr and ir_set.rptr.is_inbound:
ptrref = ir_set.rptr.ptrref
ptr_info = pg_types.get_ptrref_storage_info(ptrref,
resolve_type=False,
link_bias=False)
if ptr_info.table_type == 'ObjectType':
# Inline link
rref = pgast.ColumnRef(name=[ptr_info.column_name],
nullable=not ptrref.required)
pathctx.put_rvar_path_bond(set_rvar, ir_set.path_id.src_path())
pathctx.put_rvar_path_output(set_rvar,
ir_set.path_id.src_path(),
aspect='identity',
var=rref,
env=ctx.env)
if astutils.is_set_op_query(set_rvar.query):
astutils.for_each_query_in_set(
set_rvar.query, lambda qry: qry.target_list.append(
pgast.ResTarget(
val=rref,
name=ptr_info.column_name,
)))
return set_rvar
def array_as_json_object(
expr: pgast.BaseExpr,
*,
styperef: irast.TypeRef,
env: context.Environment,
) -> pgast.BaseExpr:
el_type = styperef.subtypes[0]
if irtyputils.is_tuple(el_type):
coldeflist = []
json_args: List[pgast.BaseExpr] = []
is_named = any(st.element_name for st in el_type.subtypes)
for i, st in enumerate(el_type.subtypes):
if is_named:
colname = st.element_name
json_args.append(pgast.StringConstant(val=st.element_name))
else:
colname = str(i)
val: pgast.BaseExpr = pgast.ColumnRef(name=[colname])
if irtyputils.is_collection(st):
val = coll_as_json_object(val, styperef=st, env=env)
json_args.append(val)
if not irtyputils.is_persistent_tuple(el_type):
# Column definition list is only allowed for functions
# returning "record", i.e. an anonymous tuple, which
# would not be the case for schema-persistent tuple types.
coldeflist.append(
pgast.ColumnDef(
name=colname,
typename=pgast.TypeName(
name=pgtypes.pg_type_from_ir_typeref(st)
)
)
)
if is_named:
json_func = _get_json_func('build_object', env=env)
else:
json_func = _get_json_func('build_array', env=env)
return pgast.SelectStmt(
target_list=[
pgast.ResTarget(
val=pgast.CoalesceExpr(
args=[
pgast.FuncCall(
name=_get_json_func('agg', env=env),
args=[
pgast.FuncCall(
name=json_func,
args=json_args,
)
]
),
pgast.StringConstant(val='[]'),
]
),
ser_safe=True,
)
],
from_clause=[
pgast.RangeFunction(
alias=pgast.Alias(
aliasname=env.aliases.get('q'),
),
is_rowsfrom=True,
functions=[
pgast.FuncCall(
name=('unnest',),
args=[expr],
coldeflist=coldeflist,
)
]
)
]
)
else:
return pgast.FuncCall(
name=_get_json_func('to', env=env), args=[expr],
null_safe=True, ser_safe=True)
def named_tuple_as_json_object(
expr: pgast.BaseExpr,
*,
styperef: irast.TypeRef,
env: context.Environment,
) -> pgast.BaseExpr:
keyvals: List[pgast.BaseExpr] = []
if irtyputils.is_persistent_tuple(styperef):
for el_type in styperef.subtypes:
keyvals.append(pgast.StringConstant(val=el_type.element_name))
val: pgast.BaseExpr = pgast.Indirection(
arg=expr,
indirection=[
pgast.ColumnRef(
name=[el_type.element_name]
)
]
)
if irtyputils.is_collection(el_type):
val = coll_as_json_object(val, styperef=el_type, env=env)
keyvals.append(val)
return pgast.FuncCall(
name=_get_json_func('build_object', env=env),
args=keyvals, null_safe=True, ser_safe=True,
nullable=expr.nullable)
else:
coldeflist = []
for el_type in styperef.subtypes:
keyvals.append(pgast.StringConstant(val=el_type.element_name))
coldeflist.append(pgast.ColumnDef(
name=el_type.element_name,
typename=pgast.TypeName(
name=pgtypes.pg_type_from_ir_typeref(el_type),
),
))
val = pgast.ColumnRef(name=[el_type.element_name])
if irtyputils.is_collection(el_type):
val = coll_as_json_object(val, styperef=el_type, env=env)
keyvals.append(val)
res = pgast.FuncCall(
name=_get_json_func('build_object', env=env),
args=keyvals, null_safe=True, ser_safe=True,
nullable=expr.nullable)
return pgast.SelectStmt(
target_list=[
pgast.ResTarget(
val=res,
),
],
from_clause=[
pgast.RangeFunction(
functions=[
pgast.FuncCall(
name=('unnest',),
args=[
pgast.ArrayExpr(
elements=[expr],
)
],
coldeflist=coldeflist,
)
]
)
]
)
def new_primitive_rvar(
ir_set: irast.Set,
*,
path_id: irast.PathId,
ctx: context.CompilerContextLevel,
) -> pgast.PathRangeVar:
if not ir_set.path_id.is_objtype_path():
raise ValueError('cannot create root rvar for non-object path')
typeref = ir_set.typeref
dml_source = irutils.get_nearest_dml_stmt(ir_set)
set_rvar = range_for_typeref(typeref,
path_id,
dml_source=dml_source,
ctx=ctx)
pathctx.put_rvar_path_bond(set_rvar, path_id)
if ir_set.rptr is not None:
ptr_ref_map: Dict[uuid.UUID, irast.BasePointerRef] = {}
p: irast.BasePointerRef
rptrref = ir_set.rptr.ptrref
if isinstance(rptrref, irast.TypeIntersectionPointerRef):
if rptrref.rptr_specialization:
for p in rptrref.rptr_specialization:
ptr_ref_map[p.dir_target.id] = p
src_set = ir_set.rptr.source
if src_set.rptr is not None:
src_rptrref = src_set.rptr.ptrref
if src_rptrref.union_components:
for p in src_rptrref.union_components:
ptr_ref_map[p.dir_target.id] = p
else:
ptr_ref_map[src_rptrref.dir_target.id] = src_rptrref
rptrref = src_rptrref
else:
ptr_ref_map[rptrref.dir_target.id] = rptrref
else:
if rptrref.union_components:
for p in rptrref.union_components:
ptr_ref_map[p.dir_target.id] = p
else:
ptr_ref_map[rptrref.dir_target.id] = rptrref
if (set_rvar.typeref is not None
and (narrow_rptrref := ptr_ref_map.get(set_rvar.typeref.id))):
rptrref = narrow_rptrref
ptr_info = pg_types.get_ptrref_storage_info(rptrref,
resolve_type=False,
link_bias=False)
if ptr_info.table_type == 'ObjectType' and rptrref.is_inbound:
# Inline link
prefix_path_id = path_id.src_path()
assert prefix_path_id is not None, 'expected a path'
rref = pgast.ColumnRef(name=[ptr_info.column_name],
nullable=not rptrref.required)
pathctx.put_rvar_path_bond(set_rvar, prefix_path_id)
pathctx.put_rvar_path_output(set_rvar,
prefix_path_id,
aspect='identity',
var=rref,
env=ctx.env)
if astutils.is_set_op_query(set_rvar.query):
assert isinstance(set_rvar.query, pgast.SelectStmt)
def _pull_col(comp_qry: pgast.Query) -> None:
rvar = pathctx.get_path_rvar(comp_qry,
path_id,
aspect='source',
env=ctx.env)
typeref = rvar.typeref
assert typeref is not None
comp_ptrref = ptr_ref_map[typeref.id]
comp_pi = pg_types.get_ptrref_storage_info(
comp_ptrref, resolve_type=False, link_bias=False)
comp_qry.target_list.append(
pgast.ResTarget(
val=pgast.ColumnRef(name=[comp_pi.column_name]),
name=ptr_info.column_name,
))
astutils.for_each_query_in_set(
set_rvar.query,
_pull_col,
)
elif isinstance(set_rvar, pgast.RangeSubselect):
rvar_path_var = pathctx.maybe_get_path_rvar(
set_rvar.query,
path_id=path_id,
aspect='identity',
env=ctx.env,
)
if isinstance(rvar_path_var, pgast.IntersectionRangeVar):
for comp_rvar in rvar_path_var.component_rvars:
if comp_rvar.typeref is None:
continue
comp_ptrref = ptr_ref_map.get(comp_rvar.typeref.id)
if comp_ptrref is None:
continue
comp_pi = pg_types.get_ptrref_storage_info(
comp_ptrref, resolve_type=False)
set_rvar.query.target_list.append(
pgast.ResTarget(
val=pgast.ColumnRef(name=[
comp_rvar.alias.aliasname,
comp_pi.column_name,
]),
name=ptr_info.column_name,
))
def range_for_ptrref(
ptrref: irast.BasePointerRef,
*,
dml_source: Optional[irast.MutatingStmt] = None,
for_mutation: bool = False,
only_self: bool = False,
ctx: context.CompilerContextLevel,
) -> pgast.PathRangeVar:
""""Return a Range subclass corresponding to a given ptr step.
The return value may potentially be a UNION of all tables
corresponding to a set of specialized links computed from the given
`ptrref` taking source inheritance into account.
"""
tgt_col = pg_types.get_ptrref_storage_info(ptrref,
resolve_type=False,
link_bias=True).column_name
cols = ['source', tgt_col]
set_ops = []
if ptrref.union_components:
refs = ptrref.union_components
if only_self and len(refs) > 1:
raise errors.InternalServerError('unexpected union link')
else:
refs = {ptrref}
assert isinstance(ptrref, irast.PointerRef), \
"expected regular PointerRef"
overlays = get_ptr_rel_overlays(ptrref, dml_source=dml_source, ctx=ctx)
for src_ptrref in refs:
assert isinstance(src_ptrref, irast.PointerRef), \
"expected regular PointerRef"
table = table_from_ptrref(
src_ptrref,
include_descendants=not ptrref.union_is_concrete,
for_mutation=for_mutation,
ctx=ctx,
)
qry = pgast.SelectStmt()
qry.from_clause.append(table)
# Make sure all property references are pulled up properly
for colname in cols:
selexpr = pgast.ColumnRef(name=[table.alias.aliasname, colname])
qry.target_list.append(pgast.ResTarget(val=selexpr, name=colname))
set_ops.append(('union', qry))
overlays = get_ptr_rel_overlays(src_ptrref,
dml_source=dml_source,
ctx=ctx)
if overlays and not for_mutation:
for op, cte in overlays:
rvar = pgast.RelRangeVar(
relation=cte,
alias=pgast.Alias(aliasname=ctx.env.aliases.get(cte.name)))
qry = pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.ColumnRef(name=[col]))
for col in cols
],
from_clause=[rvar],
)
set_ops.append((op, qry))
return range_from_queryset(set_ops, ptrref.shortname, ctx=ctx)
def process_link_values(
ir_stmt, ir_expr, target_tab, col_data,
dml_rvar, sources, props_only, target_is_scalar, iterator_cte, *,
ctx=context.CompilerContext) -> \
typing.Tuple[pgast.CommonTableExpr, typing.List[str]]:
"""Unpack data from an update expression into a series of selects.
:param ir_expr:
IR of the INSERT/UPDATE body element.
:param target_tab:
The link table being updated.
:param col_data:
Expressions used to populate well-known columns of the link
table such as `source` and `__type__`.
:param sources:
A list of relations which must be joined into the data query
to resolve expressions in *col_data*.
:param props_only:
Whether this link update only touches link properties.
:param target_is_scalar:
Whether the link target is an ScalarType.
:param iterator_cte:
CTE representing the iterator range in the FOR clause of the
EdgeQL DML statement.
"""
with ctx.newscope() as newscope, newscope.newrel() as subrelctx:
row_query = subrelctx.rel
relctx.include_rvar(row_query,
dml_rvar,
path_id=ir_stmt.subject.path_id,
ctx=subrelctx)
subrelctx.path_scope[ir_stmt.subject.path_id] = row_query
if iterator_cte is not None:
iterator_rvar = relctx.rvar_for_rel(iterator_cte,
lateral=True,
ctx=subrelctx)
relctx.include_rvar(row_query,
iterator_rvar,
path_id=iterator_cte.query.path_id,
ctx=subrelctx)
with subrelctx.newscope() as sctx, sctx.subrel() as input_rel_ctx:
input_rel = input_rel_ctx.rel
if iterator_cte is not None:
input_rel_ctx.path_scope[iterator_cte.query.path_id] = \
row_query
input_rel_ctx.expr_exposed = False
input_rel_ctx.volatility_ref = pathctx.get_path_identity_var(
row_query, ir_stmt.subject.path_id, env=input_rel_ctx.env)
dispatch.visit(ir_expr, ctx=input_rel_ctx)
shape_tuple = None
if ir_expr.shape:
shape_tuple = shapecomp.compile_shape(ir_expr,
ir_expr.shape,
ctx=input_rel_ctx)
for element in shape_tuple.elements:
pathctx.put_path_var_if_not_exists(input_rel_ctx.rel,
element.path_id,
element.val,
aspect='value',
env=input_rel_ctx.env)
input_stmt = input_rel
input_rvar = pgast.RangeSubselect(
subquery=input_rel,
lateral=True,
alias=pgast.Alias(aliasname=ctx.env.aliases.get('val')))
source_data: typing.Dict[str, pgast.BaseExpr] = {}
if input_stmt.op is not None:
# UNION
input_stmt = input_stmt.rarg
path_id = ir_expr.path_id
if shape_tuple is not None:
for element in shape_tuple.elements:
if not element.path_id.is_linkprop_path():
continue
colname = element.path_id.rptr_name().name
val = pathctx.get_rvar_path_value_var(input_rvar,
element.path_id,
env=ctx.env)
source_data.setdefault(colname, val)
else:
if target_is_scalar:
target_ref = pathctx.get_rvar_path_value_var(input_rvar,
path_id,
env=ctx.env)
else:
target_ref = pathctx.get_rvar_path_identity_var(input_rvar,
path_id,
env=ctx.env)
source_data['target'] = target_ref
if not target_is_scalar and 'target' not in source_data:
target_ref = pathctx.get_rvar_path_identity_var(input_rvar,
path_id,
env=ctx.env)
source_data['target'] = target_ref
specified_cols = []
for col, expr in collections.ChainMap(col_data, source_data).items():
row_query.target_list.append(pgast.ResTarget(val=expr, name=col))
specified_cols.append(col)
row_query.from_clause += list(sources) + [input_rvar]
link_rows = pgast.CommonTableExpr(query=row_query,
name=ctx.env.aliases.get(hint='r'))
return link_rows, specified_cols
def process_link_update(
*, ir_stmt: irast.MutatingStmt, ir_set: irast.Set, props_only: bool,
is_insert: bool, wrapper: pgast.Query, dml_cte: pgast.CommonTableExpr,
iterator_cte: typing.Optional[pgast.CommonTableExpr],
ctx: context.CompilerContextLevel) -> pgast.CommonTableExpr:
"""Perform updates to a link relation as part of a DML statement.
:param ir_stmt:
IR of the statement.
:param ir_set:
IR of the INSERT/UPDATE body element.
:param props_only:
Whether this link update only touches link properties.
:param wrapper:
Top-level SQL query.
:param dml_cte:
CTE representing the SQL INSERT or UPDATE to the main
relation of the Object.
:param iterator_cte:
CTE representing the iterator range in the FOR clause of the
EdgeQL DML statement.
"""
toplevel = ctx.toplevel_stmt
rptr = ir_set.rptr
ptrref = rptr.ptrref
assert isinstance(ptrref, irast.PointerRef)
target_is_scalar = irtyputils.is_scalar(ptrref.dir_target)
path_id = ir_set.path_id
# The links in the dml class shape have been derived,
# but we must use the correct specialized link class for the
# base material type.
if ptrref.material_ptr is not None:
mptrref = ptrref.material_ptr
assert isinstance(mptrref, irast.PointerRef)
else:
mptrref = ptrref
target_rvar = relctx.range_for_ptrref(mptrref,
include_overlays=False,
only_self=True,
ctx=ctx)
assert isinstance(target_rvar, pgast.RelRangeVar)
assert isinstance(target_rvar.relation, pgast.Relation)
target_alias = target_rvar.alias.aliasname
target_tab_name = (target_rvar.relation.schemaname,
target_rvar.relation.name)
dml_cte_rvar = pgast.RelRangeVar(
relation=dml_cte,
alias=pgast.Alias(aliasname=ctx.env.aliases.get('m')))
col_data = {
'ptr_item_id':
pgast.TypeCast(arg=pgast.StringConstant(val=str(mptrref.id)),
type_name=pgast.TypeName(name=('uuid', ))),
'source':
pathctx.get_rvar_path_identity_var(dml_cte_rvar,
ir_stmt.subject.path_id,
env=ctx.env)
}
if not is_insert:
# Drop all previous link records for this source.
delcte = pgast.CommonTableExpr(query=pgast.DeleteStmt(
relation=target_rvar,
where_clause=astutils.new_binop(
lexpr=col_data['source'],
op='=',
rexpr=pgast.ColumnRef(name=[target_alias, 'source'])),
using_clause=[dml_cte_rvar],
returning_list=[
pgast.ResTarget(val=pgast.ColumnRef(
name=[target_alias, pgast.Star()]))
]),
name=ctx.env.aliases.get(hint='d'))
pathctx.put_path_value_rvar(delcte.query,
path_id.ptr_path(),
target_rvar,
env=ctx.env)
# Record the effect of this removal in the relation overlay
# context to ensure that references to the link in the result
# of this DML statement yield the expected results.
dml_stack = get_dml_stmt_stack(ir_stmt, ctx=ctx)
relctx.add_ptr_rel_overlay(ptrref,
'except',
delcte,
dml_stmts=dml_stack,
ctx=ctx)
toplevel.ctes.append(delcte)
# Turn the IR of the expression on the right side of :=
# into a subquery returning records for the link table.
data_cte, specified_cols = process_link_values(ir_stmt,
ir_set,
target_tab_name,
col_data,
dml_cte_rvar, [],
props_only,
target_is_scalar,
iterator_cte,
ctx=ctx)
toplevel.ctes.append(data_cte)
data_select = pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.ColumnRef(
name=[data_cte.name, pgast.Star()]))
],
from_clause=[pgast.RelRangeVar(relation=data_cte)])
cols = [pgast.ColumnRef(name=[col]) for col in specified_cols]
if is_insert:
conflict_clause = None
else:
# Inserting rows into the link table may produce cardinality
# constraint violations, since the INSERT into the link table
# is executed in the snapshot where the above DELETE from
# the link table is not visible. Hence, we need to use
# the ON CONFLICT clause to resolve this.
conflict_cols = ['source', 'target', 'ptr_item_id']
conflict_inference = []
conflict_exc_row = []
for col in conflict_cols:
conflict_inference.append(pgast.ColumnRef(name=[col]))
conflict_exc_row.append(pgast.ColumnRef(name=['excluded', col]))
conflict_data = pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.ColumnRef(
name=[data_cte.name, pgast.Star()]))
],
from_clause=[pgast.RelRangeVar(relation=data_cte)],
where_clause=astutils.new_binop(
lexpr=pgast.ImplicitRowExpr(args=conflict_inference),
rexpr=pgast.ImplicitRowExpr(args=conflict_exc_row),
op='='))
conflict_clause = pgast.OnConflictClause(
action='update',
infer=pgast.InferClause(index_elems=conflict_inference),
target_list=[
pgast.MultiAssignRef(columns=cols, source=conflict_data)
])
updcte = pgast.CommonTableExpr(
name=ctx.env.aliases.get(hint='i'),
query=pgast.InsertStmt(
relation=target_rvar,
select_stmt=data_select,
cols=cols,
on_conflict=conflict_clause,
returning_list=[
pgast.ResTarget(val=pgast.ColumnRef(name=[pgast.Star()]))
]))
pathctx.put_path_value_rvar(updcte.query,
path_id.ptr_path(),
target_rvar,
env=ctx.env)
# Record the effect of this insertion in the relation overlay
# context to ensure that references to the link in the result
# of this DML statement yield the expected results.
dml_stack = get_dml_stmt_stack(ir_stmt, ctx=ctx)
relctx.add_ptr_rel_overlay(ptrref,
'union',
updcte,
dml_stmts=dml_stack,
ctx=ctx)
toplevel.ctes.append(updcte)
return data_cte