def populate_argmap(
params: List[irast.Param],
globals: List[irast.Global],
*,
ctx: context.CompilerContextLevel,
) -> None:
next_argument = itertools.count(1)
for param in params:
if ctx.env.use_named_params and not param.name.isdecimal():
continue
if param.name.startswith('__edb_arg_'):
index = int(param.name[10:]) + 1
elif param.name.isdecimal():
index = int(param.name) + 1
else:
index = next(next_argument)
ctx.argmap[param.name] = pgast.Param(
index=index,
required=param.required,
)
for param in globals:
ctx.argmap[param.name] = pgast.Param(
index=len(ctx.argmap) + 1,
required=param.required,
)
if param.has_present_arg:
ctx.argmap[param.name + "present__"] = pgast.Param(
index=len(ctx.argmap) + 1,
required=True,
)
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 compile_Parameter(expr: irast.Parameter, *,
ctx: context.CompilerContextLevel) -> pgast.BaseExpr:
result: pgast.BaseParamRef
is_decimal: bool = expr.name.isdecimal()
if not is_decimal and ctx.env.use_named_params:
result = pgast.NamedParamRef(
name=expr.name,
nullable=not expr.required,
)
else:
try:
index = ctx.argmap[expr.name].index
except KeyError:
if expr.name in ctx.argmap:
index = ctx.argmap[expr.name].index
else:
if expr.name.startswith('__edb_arg_'):
index = int(expr.name[10:]) + 1
elif is_decimal:
index = int(expr.name) + 1
else:
index = next(ctx.next_argument)
ctx.argmap[expr.name] = pgast.Param(
index=index,
required=expr.required,
)
result = pgast.ParamRef(number=index, nullable=not expr.required)
return pgast.TypeCast(
arg=result,
type_name=pgast.TypeName(
name=pg_types.pg_type_from_ir_typeref(expr.typeref)))
