Search std for module name when using with x as module. (#7836)

- Allows queries such as `with http as net::http select http::Response`. The aliased module `net::http` will resolve to `std::net::http`.
- Apply module aliases to function names during normalization, even if no function was found.
    - This fixes an issue where given the following modules:
        - `module A { function foo() -> int64 using (1) }`
        - `module B {}`
        - `module C { alias query := (with A as B select A::foo()}`
    - `C::query` would normalize to `select A::foo()` with the expectation that `A::Foo` does not exist.

edb/edgeql/compiler/normalization.py

@@ -419,6 +419,30 @@ def normalize_FunctionCall(
             sname = funcs[0].get_shortname(schema)
             node.func = (sname.module, sname.name)
 
+        elif modaliases and isinstance(name, sn.QualName):
+            # Even if no function was found, apply the modaliases.
+            # It is possible that a function without the modalias exists but
+            # we don't want to find that.
+            #
+            # Eg.
+            # module A {
+            #     function foo() -> int64 using (1);
+            # }
+            # module B {}
+            # module default {
+            #     alias query := (with A as module B select A::foo() );
+            # }
+            current_module = (
+                modaliases[None]
+                if modaliases and None in modaliases else
+                None
+            )
+            _, module = s_schema.apply_module_aliases(
+                name.module, modaliases, current_module,
+            )
+            if module is not None:
+                node.func = (module, name.name)
+
         # It's odd we don't find a function, but this will be picked up
         # by the compiler with a more appropriate error message.
 

edb/edgeql/tracer.py

@@ -381,42 +381,63 @@ def resolve_name(
     """Resolve a name into a fully-qualified one.
 
     This takes into account the current module and modaliases.
+
+    This function mostly mirrors schema.FlatSchema._search_with_getter
+    except:
+    - If no module and no default module was set, try the current module
+    - When searching in std, ensure module is not a local module
+    - If no result found, return a name with the best modname available
     """
+
+    def exists(name: sn.QualName) -> bool:
+        return (
+            objects.get(name) is not None
+            or schema.get(name, default=None, type=so.Object) is not None
+        )
+
     module = ref.module
+    orig_module = module
 
-    no_std = declaration
-    if module and module.startswith('__current__::'):
-        no_std = True
-        module = f'{current_module}::{module.removeprefix("__current__::")}'
-    elif not module:
-        module = current_module
-    elif modaliases:
-        if module:
-            first, sep, rest = module.partition('::')
-        else:
-            first, sep, rest = module, '', ''
-
-        fq_module = modaliases.get(first)
-        if fq_module is not None:
-            no_std = True
-            module = fq_module + sep + rest
-
-    qname = sn.QualName(module=module, name=ref.name)
-
-    # check if there's a name in default module
-    # that matches
-    if not no_std and not (
-        ref.module and ref.module in local_modules
-    ) and not (
-        objects.get(qname)
-        or schema.get(
-            qname, default=None, type=so.Object) is not None
-    ):
-        std_name = sn.QualName(
-            f'std::{ref.module}' if ref.module else 'std', ref.name)
-        if schema.get(std_name, default=None) is not None:
-            return std_name
-    return qname
+    # Apply module aliases
+    is_current, module = s_schema.apply_module_aliases(
+        module, modaliases, current_module,
+    )
+    no_std = declaration or is_current
+
+    # Check if something matches the name
+    if module is not None:
+        fqname = sn.QualName(module=module, name=ref.name)
+        if exists(fqname):
+            return fqname
+
+    elif orig_module is None:
+        # Look for name in current module
+        fqname = sn.QualName(module=current_module, name=ref.name)
+        if exists(fqname):
+            return fqname
+
+    # Try something in std if __current__ was not specified
+    if not no_std:
+        # If module == None, look in std
+        if orig_module is None:
+            mod_name = 'std'
+            fqname = sn.QualName(mod_name, ref.name)
+            if exists(fqname):
+                return fqname
+
+        # Ensure module is not a local module.
+        # Then try the module as part of std.
+        if module and module not in local_modules:
+            mod_name = f'std::{module}'
+            fqname = sn.QualName(mod_name, ref.name)
+            if exists(fqname):
+                return fqname
+
+    # Just pick the best module name available
+    return sn.QualName(
+        module=module or orig_module or current_module,
+        name=ref.name,
+    )
 
 
 class TracerContext:

edb/schema/schema.py

@@ -1088,6 +1088,14 @@ def _search_with_getter(
         module_aliases: Optional[Mapping[Optional[str], str]],
         disallow_module: Optional[Callable[[str], bool]],
     ) -> Any:
+        """
+        Find something in the schema with a given name.
+
+        This function mostly mirrors edgeql.tracer.resolve_name
+        except:
+        - When searching in std, disallow some modules (often the base modules)
+        - If no result found, return default
+        """
         if isinstance(name, str):
             name = sn.name_from_string(name)
         shortname = name.name
@@ -1102,60 +1110,44 @@ def _search_with_getter(
             else:
                 return default
 
-        alias_hit = local = False
-        if module and module.startswith('__current__::'):
-            local = True
-            if not module_aliases or None not in module_aliases:
-                return default
-            cur_module = module_aliases[None]
-            module = f'{cur_module}::{module.removeprefix("__current__::")}'
-        elif module_aliases is not None:
-            first: Optional[str]
-            if module:
-                first, sep, rest = module.partition('::')
-            else:
-                first, sep, rest = module, '', ''
+        # Apply module aliases
+        current_module = (
+            module_aliases[None]
+            if module_aliases and None in module_aliases else
+            None
+        )
+        is_current, module = apply_module_aliases(
+            module, module_aliases, current_module,
+        )
+        if is_current and current_module is None:
+            return default
 
-            fq_module = module_aliases.get(first)
-            if fq_module is not None:
-                alias_hit = True
-                module = fq_module + sep + rest
+        no_std = is_current
 
+        # Check if something matches the name
         if module is not None:
             fqname = sn.QualName(module, shortname)
             result = getter(self, fqname)
             if result is not None:
                 return result
 
-        # Try something in std, but only if there isn't a module clash
-        if not local and (
-            orig_module is None
-            or (
-                not alias_hit and module
-            )
-        ):
-            # If no module was specified, look in std
+        # Try something in std if __current__ was not specified
+        if not no_std:
+            # If module == None, look in std
             if orig_module is None:
                 mod_name = 'std'
                 fqname = sn.QualName(mod_name, shortname)
                 result = getter(self, fqname)
                 if result is not None:
                     return result
 
-            # If a module was specified in the name, ensure that no base module
-            # of the same name exists.
-            #
-            # If no module was specified, try the default module name as a part
-            # of std. The same condition applies.
+            # Ensure module is not a base module.
+            # Then try the module as part of std.
             if module and not (
                 self.has_module(fmod := module.split('::')[0])
                 or (disallow_module and disallow_module(fmod))
             ):
-                mod_name = (
-                    f'std::{module}'
-                    if orig_module is None
-                    else f'std::{orig_module}'
-                )
+                mod_name = f'std::{module}'
                 fqname = sn.QualName(mod_name, shortname)
                 result = getter(self, fqname)
                 if result is not None:
@@ -1544,6 +1536,34 @@ def __repr__(self) -> str:
             f'<{type(self).__name__} gen:{self._generation} at {id(self):#x}>')
 
 
+def apply_module_aliases(
+    module: Optional[str],
+    module_aliases: Optional[Mapping[Optional[str], str]],
+    current_module: Optional[str],
+) -> tuple[bool, Optional[str]]:
+    is_current = False
+    if module and module.startswith('__current__::'):
+        # Replace __current__ with default module
+        is_current = True
+        if current_module is not None:
+            module = f'{current_module}::{module.removeprefix("__current__::")}'
+        else:
+            module = None
+    elif module_aliases is not None:
+        # Apply modalias
+        first: Optional[str]
+        if module:
+            first, sep, rest = module.partition('::')
+        else:
+            first, sep, rest = module, '', ''
+
+        fq_module = module_aliases.get(first)
+        if fq_module is not None:
+            module = fq_module + sep + rest
+
+    return is_current, module
+
+
 EMPTY_SCHEMA = FlatSchema()