def _cast_array(ir_set: irast.Set, orig_stype: s_types.Type, new_stype: s_types.Type, *, srcctx: parsing.ParserContext, ctx: context.ContextLevel) -> irast.Base: direct_cast = _find_cast(orig_stype, new_stype, srcctx=srcctx, ctx=ctx) if direct_cast is None: if not new_stype.is_array(): raise errors.QueryError( f'cannot cast {orig_stype.get_displayname(ctx.env.schema)!r} ' f'to {new_stype.get_displayname(ctx.env.schema)!r}', context=srcctx) el_type = new_stype.get_subtypes(ctx.env.schema)[0] else: el_type = new_stype orig_el_type = orig_stype.get_subtypes(ctx.env.schema)[0] el_cast = _find_cast(orig_el_type, el_type, srcctx=srcctx, ctx=ctx) if el_cast is not None and el_cast.get_from_cast(ctx.env.schema): # Simple cast return _cast_to_ir(ir_set, el_cast, orig_stype, new_stype, ctx=ctx) else: pathctx.register_set_in_scope(ir_set, ctx=ctx) with ctx.new() as subctx: subctx.anchors = subctx.anchors.copy() source_alias = subctx.aliases.get('a') subctx.anchors[source_alias] = ir_set unpacked = qlast.FunctionCall( func=('std', 'array_unpack'), args=[ qlast.Path(steps=[qlast.ObjectRef(name=source_alias)], ), ], ) elements = qlast.FunctionCall( func=('std', 'array_agg'), args=[ qlast.TypeCast( expr=unpacked, type=typegen.type_to_ql_typeref(el_type, ctx=subctx), ), ], ) array_ir = dispatch.compile(elements, ctx=subctx) if direct_cast is not None: array_stype = s_types.Array.from_subtypes( ctx.env.schema, [el_type]) return _cast_to_ir(array_ir, direct_cast, array_stype, new_stype, ctx=ctx) else: return array_ir
def _cast_json_to_tuple( ir_set: irast.Set, orig_stype: s_types.Type, new_stype: s_types.Tuple, cardinality_mod: Optional[qlast.CardinalityModifier], *, srcctx: Optional[parsing.ParserContext], ctx: context.ContextLevel) -> irast.Set: with ctx.new() as subctx: subctx.anchors = subctx.anchors.copy() source_path = subctx.create_anchor(ir_set, 'a') # Top-level json->tuple casts should produce an empty set on # null inputs, but error on missing fields or null # subelements, so filter out json nulls directly here to # distinguish those cases. if cardinality_mod != qlast.CardinalityModifier.Required: pathctx.register_set_in_scope(ir_set, ctx=subctx) check = qlast.FunctionCall( func=('__std__', 'json_typeof'), args=[source_path] ) filtered = qlast.SelectQuery( result=source_path, where=qlast.BinOp( left=check, op='!=', right=qlast.StringConstant(value='null'), ) ) filtered_ir = dispatch.compile(filtered, ctx=subctx) source_path = subctx.create_anchor(filtered_ir, 'a') # TODO: try using jsonb_to_record instead of a bunch of # json_get calls and see if that is faster. elements = [] for new_el_name, new_st in new_stype.iter_subtypes(ctx.env.schema): val_e = qlast.FunctionCall( func=('__std__', 'json_get'), args=[ source_path, qlast.StringConstant(value=new_el_name), ], ) val = dispatch.compile(val_e, ctx=subctx) val = compile_cast( val, new_st, cardinality_mod=qlast.CardinalityModifier.Required, ctx=subctx, srcctx=srcctx) elements.append(irast.TupleElement(name=new_el_name, val=val)) return setgen.new_tuple_set( elements, named=new_stype.is_named(ctx.env.schema), ctx=subctx, )
def reduce_NodeName_LPAREN_OptFuncArgList_RPAREN(self, *kids): module = kids[0].val.module func_name = kids[0].val.name name = func_name if not module else (module, func_name) last_named_seen = None args = [] kwargs = {} for argname, argname_ctx, arg in kids[2].val: if argname is not None: if argname in kwargs: raise EdgeQLSyntaxError( f"duplicate named argument `{argname}`", context=argname_ctx) last_named_seen = argname kwargs[argname] = arg else: if last_named_seen is not None: raise EdgeQLSyntaxError( f"positional argument after named " f"argument `{last_named_seen}`", context=arg.context) args.append(arg) self.val = qlast.FunctionCall(func=name, args=args, kwargs=kwargs)
def eta_expand_ordered( expr: qlast.Expr, stype: s_types.Type, *, ctx: context.ContextLevel, ) -> qlast.Expr: """Do an order-preserving η-expansion Unlike in the lambda calculus, edgeql is a set-based language with a notion of ordering, which we need to preserve. We do this by using enumerate and ORDER BY on it: EXPAND_ORDERED(t, e) = WITH enum := enumerate(e) SELECT EXPAND(t, enum.1) ORDER BY enum.0 """ enumerated = qlast.FunctionCall(func=('__std__', 'enumerate'), args=[expr]) enumerated_alias, enumerated_path = _get_alias('enum', ctx=ctx) element_path = astutils.extend_path(enumerated_path, '1') result_expr = eta_expand(element_path, stype, ctx=ctx) return qlast.SelectQuery( result=result_expr, orderby=[ qlast.SortExpr(path=astutils.extend_path(enumerated_path, '0')) ], aliases=[qlast.AliasedExpr(alias=enumerated_alias, expr=enumerated)], )
def visit_FunctionCall(self, node): args = node.args args = [qlast.FuncArg(arg=self.visit(arg.expr)) for arg in args] result = qlast.FunctionCall( func=(node.func_shortname.module, node.func_shortname.name), args=args, ) return result
def get_params_symtable( params: FuncParameterList, schema: s_schema.Schema, *, inlined_defaults: bool, ) -> Dict[str, qlast.Expr]: anchors: Dict[str, qlast.Expr] = {} defaults_mask = qlast.TypeCast( expr=qlast.Parameter(name='__defaults_mask__', optional=False), type=qlast.TypeName( maintype=qlast.ObjectRef( module='std', name='bytes', ), ), ) for pi, p in enumerate(params.get_in_canonical_order(schema)): p_shortname = p.get_parameter_name(schema) p_is_optional = p.get_typemod(schema) is not ft.TypeModifier.SINGLETON anchors[p_shortname] = qlast.TypeCast( expr=qlast.Parameter( name=p_shortname, optional=p_is_optional, ), type=utils.typeref_to_ast(schema, p.get_type(schema)), ) p_default = p.get_default(schema) if p_default is None: continue if not inlined_defaults: continue anchors[p_shortname] = qlast.IfElse( condition=qlast.BinOp( left=qlast.FunctionCall( func=('std', 'bytes_get_bit'), args=[ defaults_mask, qlast.IntegerConstant(value=str(pi)), ]), op='=', right=qlast.IntegerConstant(value='0'), ), if_expr=anchors[p_shortname], else_expr=qlast._Optional(expr=p_default.qlast), ) return anchors
def get_param_anchors_for_callable( params: s_func.ParameterLikeList, schema: s_schema.Schema, *, inlined_defaults: bool, ) -> Tuple[Dict[str, irast.Parameter], List[qlast.AliasedExpr], ]: anchors = {} aliases = [] if inlined_defaults: anchors['__defaults_mask__'] = irast.Parameter( name='__defaults_mask__', typeref=irtyputils.type_to_typeref( # note: no cache schema, cast(s_scalars.ScalarType, schema.get('std::bytes')), ), ) pg_params = s_func.PgParams.from_params(schema, params) for pi, p in enumerate(pg_params.params): p_shortname = p.get_shortname(schema) anchors[p_shortname] = irast.Parameter( name=p_shortname, typeref=irtyputils.type_to_typeref(schema, p.get_type(schema))) if p.get_default(schema) is None: continue if not inlined_defaults: continue aliases.append( qlast.AliasedExpr( alias=p_shortname, expr=qlast. IfElse(condition=qlast.BinOp(left=qlast.FunctionCall( func=('std', 'bytes_get_bit'), args=[ qlast.Path( steps=[qlast.ObjectRef(name='__defaults_mask__')]), qlast.IntegerConstant(value=str(pi)), ]), right=qlast.IntegerConstant( value='0'), op='='), if_expr=qlast.Path( steps=[qlast.ObjectRef(name=p_shortname)]), else_expr=qlast._Optional( expr=p.get_ql_default(schema))))) return anchors, aliases
def eta_expand_array( path: qlast.Path, stype: s_types.Array, *, ctx: context.ContextLevel, ) -> qlast.Expr: """η-expansion of arrays η-expansion of array types is is a little peculiar to edgeql and less grounded in typed lambda calculi: EXPAND(array<t>, p) = (p, array_agg(EXPAND_ORDERED(t, array_unpack(p)))).1 We use a similar approach for compiling casts. The tuple projection trick serves to make sure that we iterate over `p` *outside* of the array_agg (or else all the arrays would get aggregated together) as well as ensuring that `p` appears in the expansion in a non-fenced position (or else sorting it from outside wouldn't work). (If it wasn't for the latter requirement, we could just use a FOR. I find it a little unsatisfying that our η-expansion needs to use this trick, and the pgsql compiler needed to be hacked to make it work.) """ unpacked = qlast.FunctionCall(func=('__std__', 'array_unpack'), args=[path]) expanded = eta_expand_ordered(unpacked, stype.get_element_type(ctx.env.schema), ctx=ctx) agg_expr = qlast.FunctionCall(func=('__std__', 'array_agg'), args=[expanded]) return astutils.extend_path(qlast.Tuple(elements=[path, agg_expr]), '1')
def _cast_json_to_tuple(ir_set: irast.Set, orig_stype: s_types.Type, new_stype: s_types.Tuple, *, srcctx: Optional[parsing.ParserContext], ctx: context.ContextLevel) -> irast.Set: with ctx.new() as subctx: subctx.anchors = subctx.anchors.copy() source_alias = subctx.aliases.get('a') subctx.anchors[source_alias] = ir_set # TODO: try using jsonb_to_record instead of a bunch of # json_get calls and see if that is faster. elements = [] for new_el_name, new_st in new_stype.iter_subtypes(ctx.env.schema): val_e = qlast.FunctionCall( func=('__std__', 'json_get'), args=[ qlast.Path(steps=[qlast.ObjectRef(name=source_alias)]), qlast.StringConstant(value=new_el_name), ], ) val = dispatch.compile(val_e, ctx=subctx) val = compile_cast( val, new_st, cardinality_mod=qlast.CardinalityModifier.Required, ctx=subctx, srcctx=srcctx) elements.append(irast.TupleElement(name=new_el_name, val=val)) return setgen.new_tuple_set( elements, named=new_stype.is_named(ctx.env.schema), ctx=subctx, )
def get_param_anchors_for_callable(params, schema): anchors = {} aliases = [] anchors['__defaults_mask__'] = irast.Parameter( name='__defaults_mask__', typeref=irtyputils.type_to_typeref(schema, schema.get('std::bytes'))) pg_params = s_func.PgParams.from_params(schema, params) for pi, p in enumerate(pg_params.params): p_shortname = p.get_shortname(schema) anchors[p_shortname] = irast.Parameter( name=p_shortname, typeref=irtyputils.type_to_typeref(schema, p.get_type(schema))) if p.get_default(schema) is None: continue aliases.append( qlast.AliasedExpr( alias=p_shortname, expr=qlast. IfElse(condition=qlast.BinOp(left=qlast.FunctionCall( func=('std', 'bytes_get_bit'), args=[ qlast.Path( steps=[qlast.ObjectRef(name='__defaults_mask__')]), qlast.IntegerConstant(value=str(pi)), ]), right=qlast.IntegerConstant( value='0'), op='='), if_expr=qlast.Path( steps=[qlast.ObjectRef(name=p_shortname)]), else_expr=qlast._Optional( expr=p.get_ql_default(schema))))) return anchors, aliases
def compile_func_to_ir(func, schema, *, anchors=None, security_context=None, modaliases=None, implicit_id_in_shapes=False, implicit_tid_in_shapes=False): """Compile an EdgeQL function into EdgeDB IR.""" if debug.flags.edgeql_compile: debug.header('EdgeQL Function') debug.print(func.get_code(schema)) trees = ql_parser.parse_block(func.get_code(schema) + ';') if len(trees) != 1: raise errors.InvalidFunctionDefinitionError( 'functions can only contain one statement') tree = trees[0] if modaliases: ql_parser.append_module_aliases(tree, modaliases) if anchors is None: anchors = {} anchors['__defaults_mask__'] = irast.Parameter( name='__defaults_mask__', typeref=irtyputils.type_to_typeref(schema, schema.get('std::bytes'))) func_params = func.get_params(schema) pg_params = s_func.PgParams.from_params(schema, func_params) for pi, p in enumerate(pg_params.params): p_shortname = p.get_shortname(schema) anchors[p_shortname] = irast.Parameter( name=p_shortname, typeref=irtyputils.type_to_typeref(schema, p.get_type(schema))) if p.get_default(schema) is None: continue tree.aliases.append( qlast.AliasedExpr( alias=p_shortname, expr=qlast.IfElse( condition=qlast.BinOp( left=qlast.FunctionCall( func=('std', 'bytes_get_bit'), args=[ qlast.FuncArg( arg=qlast.Path(steps=[ qlast.ObjectRef( name='__defaults_mask__') ])), qlast.FuncArg( arg=qlast.IntegerConstant(value=str(pi))) ]), right=qlast.IntegerConstant(value='0'), op='='), if_expr=qlast.Path( steps=[qlast.ObjectRef(name=p_shortname)]), else_expr=qlast._Optional(expr=p.get_ql_default(schema))))) ir = compile_ast_to_ir( tree, schema, anchors=anchors, func=func, security_context=security_context, modaliases=modaliases, implicit_id_in_shapes=implicit_id_in_shapes, implicit_tid_in_shapes=implicit_tid_in_shapes) return ir
def _cast_array(ir_set: irast.Set, orig_stype: s_types.Type, new_stype: s_types.Type, *, srcctx: Optional[parsing.ParserContext], ctx: context.ContextLevel) -> irast.Set: assert isinstance(orig_stype, s_types.Array) direct_cast = _find_cast(orig_stype, new_stype, srcctx=srcctx, ctx=ctx) if direct_cast is None: if not new_stype.is_array(): raise errors.QueryError( f'cannot cast {orig_stype.get_displayname(ctx.env.schema)!r} ' f'to {new_stype.get_displayname(ctx.env.schema)!r}', context=srcctx) assert isinstance(new_stype, s_types.Array) el_type = new_stype.get_subtypes(ctx.env.schema)[0] else: el_type = new_stype orig_el_type = orig_stype.get_subtypes(ctx.env.schema)[0] el_cast = _find_cast(orig_el_type, el_type, srcctx=srcctx, ctx=ctx) if el_cast is not None and el_cast.get_from_cast(ctx.env.schema): # Simple cast return _cast_to_ir(ir_set, el_cast, orig_stype, new_stype, ctx=ctx) else: pathctx.register_set_in_scope(ir_set, ctx=ctx) with ctx.new() as subctx: subctx.anchors = subctx.anchors.copy() source_alias = subctx.aliases.get('a') subctx.anchors[source_alias] = ir_set unpacked = qlast.FunctionCall( func=('__std__', 'array_unpack'), args=[ qlast.Path(steps=[qlast.ObjectRef(name=source_alias)], ), ], ) enumerated = setgen.ensure_set( dispatch.compile( qlast.FunctionCall( func=('__std__', 'enumerate'), args=[unpacked], ), ctx=subctx, ), ctx=subctx, ) enumerated_alias = subctx.aliases.get('e') subctx.anchors[enumerated_alias] = enumerated enumerated_ref = qlast.Path( steps=[qlast.ObjectRef(name=enumerated_alias)], ) elements = qlast.FunctionCall( func=('__std__', 'array_agg'), args=[ qlast.SelectQuery( result=qlast.TypeCast( expr=qlast.Path(steps=[ enumerated_ref, qlast.Ptr(ptr=qlast.ObjectRef( name='1', direction='>', ), ), ], ), type=typegen.type_to_ql_typeref( el_type, ctx=subctx, ), cardinality_mod=qlast.CardinalityModifier.Required, ), orderby=[ qlast.SortExpr( path=qlast.Path(steps=[ enumerated_ref, qlast.Ptr(ptr=qlast.ObjectRef( name='0', direction='>', ), ), ], ), direction=qlast.SortOrder.Asc, ), ], ), ], ) array_ir = dispatch.compile(elements, ctx=subctx) assert isinstance(array_ir, irast.Set) if direct_cast is not None: ctx.env.schema, array_stype = s_types.Array.from_subtypes( ctx.env.schema, [el_type]) return _cast_to_ir(array_ir, direct_cast, array_stype, new_stype, ctx=ctx) else: return array_ir
subjectexpr = constr.get_subjectexpr(ctx.env.schema) assert subjectexpr and isinstance(subjectexpr.qlast, qlast.Expr) lhs = qlutils.subject_paths_substitute(subjectexpr.qlast, ptr_anchors) rhs = qlutils.subject_substitute(subjectexpr.qlast, insert_subject) conds.append(qlast.BinOp(op='=', left=lhs, right=rhs)) if not conds: return None # We use `any` to compute the disjunction here because some might # be empty. if len(conds) == 1: cond = conds[0] else: cond = qlast.FunctionCall( func='any', args=[qlast.Set(elements=conds)], ) # For the result filtering we need to *ignore* the same object if fake_dml_set: anchor = qlutils.subject_paths_substitute( ptr_anchors['id'], ptr_anchors) ptr_val = qlast.Path(partial=True, steps=[ qlast.Ptr(ptr=qlast.ObjectRef(name='id')) ]) cond = qlast.BinOp( op='AND', left=cond, right=qlast.BinOp(op='!=', left=anchor, right=ptr_val), )
def compile_DescribeStmt(ql: qlast.DescribeStmt, *, ctx: context.ContextLevel) -> irast.Set: with ctx.subquery() as ictx: stmt = irast.SelectStmt() init_stmt(stmt, ql, ctx=ictx, parent_ctx=ctx) if ql.object == qlast.DescribeGlobal.Schema: if ql.language is qltypes.DescribeLanguage.DDL: # DESCRIBE SCHEMA text = s_ddl.ddl_text_from_schema(ctx.env.schema, ) else: raise errors.QueryError( f'cannot describe full schema as {ql.language}') ct = typegen.type_to_typeref( ctx.env.get_track_schema_type('std::str'), env=ctx.env, ) stmt.result = setgen.ensure_set( irast.StringConstant(value=text, typeref=ct), ctx=ictx, ) elif ql.object == qlast.DescribeGlobal.SystemConfig: if ql.language is qltypes.DescribeLanguage.DDL: function_call = dispatch.compile(qlast.FunctionCall( func=('cfg', '_describe_system_config_as_ddl'), ), ctx=ictx) assert isinstance(function_call, irast.Set), function_call stmt.result = function_call else: raise errors.QueryError( f'cannot describe config as {ql.language}') elif ql.object == qlast.DescribeGlobal.Roles: if ql.language is qltypes.DescribeLanguage.DDL: function_call = dispatch.compile(qlast.FunctionCall( func=('sys', '_describe_roles_as_ddl'), ), ctx=ictx) assert isinstance(function_call, irast.Set), function_call stmt.result = function_call else: raise errors.QueryError( f'cannot describe roles as {ql.language}') else: assert isinstance(ql.object, qlast.ObjectRef), ql.object modules = [] items: DefaultDict[str, List[str]] = defaultdict(list) referenced_classes: List[s_obj.ObjectMeta] = [] objref = ql.object itemclass = objref.itemclass if itemclass is qltypes.SchemaObjectClass.MODULE: modules.append(objref.name) else: itemtype: Optional[Type[s_obj.Object]] = None name: str if objref.module: name = s_name.Name(module=objref.module, name=objref.name) else: name = objref.name if itemclass is not None: if itemclass is qltypes.SchemaObjectClass.ALIAS: # Look for underlying derived type. itemtype = s_types.Type else: itemtype = ( s_obj.ObjectMeta.get_schema_metaclass_for_ql_class( itemclass)) last_exc = None # Search in the current namespace AND in std. We do # this to avoid masking a `std` object/function by one # in a default module. search_ns = [ictx.modaliases] # Only check 'std' separately if the current # modaliases don't already include it. if ictx.modaliases.get(None, 'std') != 'std': search_ns.append({None: 'std'}) # Search in the current namespace AND in std. for aliases in search_ns: # Use the specific modaliases instead of the # context ones. with ictx.subquery() as newctx: newctx.modaliases = aliases # Get the default module name modname = aliases[None] # Is the current item a function is_function = (itemclass is qltypes.SchemaObjectClass.FUNCTION) # We need to check functions if we're looking for them # specifically or if this is a broad search. They are # handled separately because they allow multiple # matches for the same name. if (itemclass is None or is_function): try: funcs: Tuple[s_func.Function, ...] = ( newctx.env.schema.get_functions( name, module_aliases=aliases)) except errors.InvalidReferenceError: pass else: for func in funcs: items[f'function_{modname}'].append( func.get_name(newctx.env.schema)) # Also find an object matching the name as long as # it's not a function we're looking for specifically. if not is_function: try: if itemclass is not \ qltypes.SchemaObjectClass.ALIAS: condition = None label = None else: condition = (lambda obj: obj. get_alias_is_persistent( ctx.env.schema)) label = 'alias' obj = schemactx.get_schema_object( objref, item_type=itemtype, condition=condition, label=label, ctx=newctx, ) items[f'other_{modname}'].append( obj.get_name(newctx.env.schema)) except errors.InvalidReferenceError as exc: # Record the exception to be possibly # raised if no matches are found last_exc = exc # If we already have some results, suppress the exception, # otherwise raise the recorded exception. if not items and last_exc: raise last_exc verbose = ql.options.get_flag('VERBOSE') if ql.language is qltypes.DescribeLanguage.DDL: method = s_ddl.ddl_text_from_schema elif ql.language is qltypes.DescribeLanguage.SDL: method = s_ddl.sdl_text_from_schema elif ql.language is qltypes.DescribeLanguage.TEXT: method = s_ddl.descriptive_text_from_schema if not verbose.val: referenced_classes = [s_links.Link, s_lprops.Property] else: raise errors.InternalServerError( f'cannot handle describe language {ql.language}') # Based on the items found generate main text and a # potential comment about masked items. defmod = ictx.modaliases.get(None, 'std') default_items = [] masked_items = set() for objtype in ['function', 'other']: defkey = f'{objtype}_{defmod}' mskkey = f'{objtype}_std' default_items += items.get(defkey, []) if defkey in items and mskkey in items: # We have a match in default module and some masked. masked_items.update(items.get(mskkey, [])) else: default_items += items.get(mskkey, []) # Throw out anything in the masked set that's already in # the default. masked_items.difference_update(default_items) text = method( ctx.env.schema, included_modules=modules, included_items=default_items, included_ref_classes=referenced_classes, include_module_ddl=False, include_std_ddl=True, ) if masked_items: text += ('\n\n' '# The following builtins are masked by the above:' '\n\n') masked = method( ctx.env.schema, included_modules=modules, included_items=masked_items, included_ref_classes=referenced_classes, include_module_ddl=False, include_std_ddl=True, ) masked = textwrap.indent(masked, '# ') text += masked ct = typegen.type_to_typeref( ctx.env.get_track_schema_type('std::str'), env=ctx.env, ) stmt.result = setgen.ensure_set( irast.StringConstant(value=text, typeref=ct), ctx=ictx, ) result = fini_stmt(stmt, ql, ctx=ictx, parent_ctx=ctx) return result
def _cast_array( ir_set: irast.Set, orig_stype: s_types.Type, new_stype: s_types.Type, *, srcctx: Optional[parsing.ParserContext], ctx: context.ContextLevel) -> irast.Set: assert isinstance(orig_stype, s_types.Array) direct_cast = _find_cast(orig_stype, new_stype, srcctx=srcctx, ctx=ctx) if direct_cast is None: if not new_stype.is_array(): raise errors.QueryError( f'cannot cast {orig_stype.get_displayname(ctx.env.schema)!r} ' f'to {new_stype.get_displayname(ctx.env.schema)!r}', context=srcctx) assert isinstance(new_stype, s_types.Array) el_type = new_stype.get_subtypes(ctx.env.schema)[0] else: el_type = new_stype orig_el_type = orig_stype.get_subtypes(ctx.env.schema)[0] el_cast = _find_cast(orig_el_type, el_type, srcctx=srcctx, ctx=ctx) if el_cast is not None and el_cast.get_from_cast(ctx.env.schema): # Simple cast return _cast_to_ir( ir_set, el_cast, orig_stype, new_stype, ctx=ctx) else: with ctx.new() as subctx: subctx.anchors = subctx.anchors.copy() source_path = subctx.create_anchor(ir_set, 'a') unpacked = qlast.FunctionCall( func=('__std__', 'array_unpack'), args=[source_path], ) enumerated = dispatch.compile( qlast.FunctionCall( func=('__std__', 'enumerate'), args=[unpacked], ), ctx=subctx, ) enumerated_ref = subctx.create_anchor(enumerated, 'e') elements = qlast.FunctionCall( func=('__std__', 'array_agg'), args=[ qlast.SelectQuery( result=qlast.TypeCast( expr=astutils.extend_path(enumerated_ref, '1'), type=typegen.type_to_ql_typeref( el_type, ctx=subctx, ), cardinality_mod=qlast.CardinalityModifier.Required, ), orderby=[ qlast.SortExpr( path=astutils.extend_path(enumerated_ref, '0'), direction=qlast.SortOrder.Asc, ), ], ), ], ) # Force the elements to be correlated with whatever the # anchor was. (Doing it this way ensures a NULL check, # and just registering it in the scope would not.) correlated_elements = astutils.extend_path( qlast.Tuple(elements=[source_path, elements]), '1' ) if el_type.contains_json(subctx.env.schema): subctx.inhibit_implicit_limit = True array_ir = dispatch.compile(correlated_elements, ctx=subctx) assert isinstance(array_ir, irast.Set) if direct_cast is not None: ctx.env.schema, array_stype = s_types.Array.from_subtypes( ctx.env.schema, [el_type]) return _cast_to_ir( array_ir, direct_cast, array_stype, new_stype, ctx=ctx ) else: return array_ir
def desugar_group( node: qlast.GroupQuery, aliases: AliasGenerator, ) -> qlast.InternalGroupQuery: assert not isinstance(node, qlast.InternalGroupQuery) alias_map: Dict[str, Tuple[str, qlast.Expr]] = {} def rewrite_atom(el: qlast.GroupingAtom) -> qlast.GroupingAtom: if isinstance(el, qlast.ObjectRef): return el elif isinstance(el, qlast.Path): assert isinstance(el.steps[0], qlast.Ptr) ptrname = el.steps[0].ptr.name if ptrname not in alias_map: alias = aliases.get(ptrname) alias_map[ptrname] = (alias, el) alias = alias_map[ptrname][0] return qlast.ObjectRef(name=alias) else: return qlast.GroupingIdentList( context=el.context, elements=tuple(rewrite_atom(at) for at in el.elements), ) def rewrite(el: qlast.GroupingElement) -> qlast.GroupingElement: if isinstance(el, qlast.GroupingSimple): return qlast.GroupingSimple( context=el.context, element=rewrite_atom(el.element)) elif isinstance(el, qlast.GroupingSets): return qlast.GroupingSets( context=el.context, sets=[rewrite(s) for s in el.sets]) elif isinstance(el, qlast.GroupingOperation): return qlast.GroupingOperation( context=el.context, oper=el.oper, elements=[rewrite_atom(a) for a in el.elements]) raise AssertionError for using_clause in (node.using or ()): alias_map[using_clause.alias] = (using_clause.alias, using_clause.expr) using = node.using[:] if node.using else [] by = [rewrite(by_el) for by_el in node.by] for alias, path in alias_map.values(): using.append(qlast.AliasedExpr(alias=alias, expr=path)) actual_keys = collect_grouping_atoms(by) g_alias = aliases.get('g') grouping_alias = aliases.get('grouping') output_dict = { 'key': make_free_object({ name: name_path(alias) for name, (alias, _) in alias_map.items() if alias in actual_keys }), 'grouping': qlast.FunctionCall( func='array_unpack', args=[name_path(grouping_alias)], ), 'elements': name_path(g_alias), } output_shape = make_free_object(output_dict) return qlast.InternalGroupQuery( context=node.context, aliases=node.aliases, subject_alias=node.subject_alias, subject=node.subject, # rewritten parts! using=using, by=by, group_alias=g_alias, grouping_alias=grouping_alias, result=output_shape, from_desugaring=True, )