def aggregate_json_output(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,
)
new_val = pgast.CoalesceExpr(args=[
pgast.FuncCall(name=_get_json_func('agg', env=env),
args=[pgast.ColumnRef(name=[stmt_res.name])]),
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 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 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 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
