def output_as_value(expr: pgast.BaseExpr, *,
env: context.Environment) -> pgast.BaseExpr:
val = expr
if isinstance(expr, pgast.TupleVar):
RowCls: Union[Type[pgast.ImplicitRowExpr], Type[pgast.RowExpr]]
if (env.output_format is context.OutputFormat.NATIVE_INTERNAL
and len(expr.elements) == 1
and (path_id := (el0 := expr.elements[0]).path_id) is not None
and (rptr_name := path_id.rptr_name()) is not None
and (rptr_name.name == 'id')):
# This is is a special mode whereby bare refs to objects
# are serialized to UUID values.
return output_as_value(el0.val, env=env)
elif len(expr.elements) > 1:
RowCls = pgast.ImplicitRowExpr
else:
RowCls = pgast.RowExpr
val = RowCls(
args=[output_as_value(e.val, env=env) for e in expr.elements])
if (expr.typeref is not None and not env.singleton_mode
and irtyputils.is_persistent_tuple(expr.typeref)):
pg_type = pgtypes.pg_type_from_ir_typeref(expr.typeref)
val = pgast.TypeCast(
arg=val,
type_name=pgast.TypeName(name=pg_type, ),
)
def output_as_value(
expr: pgast.BaseExpr, *,
env: context.Environment) -> pgast.BaseExpr:
val = expr
if isinstance(expr, pgast.TupleVar):
RowCls: Union[Type[pgast.ImplicitRowExpr],
Type[pgast.RowExpr]]
if len(expr.elements) > 1:
RowCls = pgast.ImplicitRowExpr
else:
RowCls = pgast.RowExpr
val = RowCls(args=[
output_as_value(e.val, env=env) for e in expr.elements
])
if (expr.typeref is not None
and not env.singleton_mode
and irtyputils.is_persistent_tuple(expr.typeref)):
pg_type = pgtypes.pg_type_from_ir_typeref(expr.typeref)
val = pgast.TypeCast(
arg=val,
type_name=pgast.TypeName(
name=pg_type,
),
)
return val
def tuple_getattr(
tuple_val: pgast.BaseExpr,
tuple_typeref: irast.TypeRef,
attr: str,
) -> pgast.BaseExpr:
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)
set_expr: pgast.BaseExpr
if irtyputils.is_persistent_tuple(tuple_typeref):
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 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 unnamed_tuple_as_json_object(
expr: pgast.BaseExpr,
*,
styperef: irast.TypeRef,
env: context.Environment,
) -> pgast.BaseExpr:
vals: List[pgast.BaseExpr] = []
if irtyputils.is_persistent_tuple(styperef):
for el_idx, el_type in enumerate(styperef.subtypes):
val: pgast.BaseExpr = 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 _build_json(
'build_array',
args=vals,
null_safe=True,
ser_safe=True,
nullable=expr.nullable,
env=env,
)
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 = _build_json(
'build_array',
args=vals,
null_safe=True,
ser_safe=True,
nullable=expr.nullable,
env=env,
)
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 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:
assert el_type.element_name
keyvals.append(pgast.StringConstant(val=el_type.element_name))
val: pgast.BaseExpr = pgast.Indirection(
arg=expr,
indirection=[pgast.ColumnRef(name=[el_type.element_name])])
val = serialize_expr_to_json(val,
styperef=el_type,
nested=True,
env=env)
keyvals.append(val)
obj = _build_json(
'build_object',
args=keyvals,
null_safe=True,
ser_safe=True,
nullable=expr.nullable,
env=env,
)
else:
coldeflist = []
for el_type in styperef.subtypes:
assert el_type.element_name
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])
val = serialize_expr_to_json(val,
styperef=el_type,
nested=True,
env=env)
keyvals.append(val)
obj = _build_json(
'build_object',
args=keyvals,
null_safe=True,
ser_safe=True,
nullable=expr.nullable,
env=env,
)
obj = pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=obj, ),
],
from_clause=[
pgast.RangeFunction(functions=[
pgast.FuncCall(
name=('unnest', ),
args=[pgast.ArrayExpr(elements=[expr], )],
coldeflist=coldeflist,
)
])
] if styperef.subtypes else [])
if expr.nullable:
obj = pgast.SelectStmt(target_list=[pgast.ResTarget(val=obj)],
where_clause=pgast.NullTest(arg=expr,
negated=True))
return obj
def array_as_json_object(
expr: pgast.BaseExpr,
*,
styperef: irast.TypeRef,
env: context.Environment,
) -> pgast.BaseExpr:
el_type = styperef.subtypes[0]
is_tuple = irtyputils.is_tuple(el_type)
# Tuples and bytes might need underlying casts to be done
if is_tuple or irtyputils.is_bytes(el_type):
coldeflist = []
out_alias = env.aliases.get('q')
val: pgast.BaseExpr
if is_tuple:
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
assert colname
json_args.append(pgast.StringConstant(val=colname))
else:
colname = str(i)
val = pgast.ColumnRef(name=[colname])
val = serialize_expr_to_json(val,
styperef=st,
nested=True,
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))))
json_func = 'build_object' if is_named else 'build_array'
agg_arg = _build_json(json_func, json_args, env=env)
needs_unnest = bool(el_type.subtypes)
else:
assert not el_type.subtypes
val = pgast.ColumnRef(name=[out_alias])
agg_arg = serialize_expr_to_json(val,
styperef=el_type,
nested=True,
env=env)
needs_unnest = True
return pgast.SelectStmt(
target_list=[
pgast.ResTarget(
val=pgast.CoalesceExpr(args=[
pgast.FuncCall(
name=_get_json_func('agg', env=env),
args=[agg_arg],
),
pgast.StringConstant(val='[]'),
]),
ser_safe=True,
)
],
from_clause=[
pgast.RangeFunction(alias=pgast.Alias(aliasname=out_alias),
is_rowsfrom=True,
functions=[
pgast.FuncCall(
name=('unnest', ),
args=[expr],
coldeflist=coldeflist,
)
])
] if needs_unnest else [],
)
else:
return pgast.FuncCall(name=_get_json_func('to', env=env),
args=[expr],
null_safe=True,
ser_safe=True)
