fix: missing main pipeline (#2587)

prql-compiler/prqlc/src/cli.rs

@@ -1,4 +1,4 @@
-use anyhow::{anyhow, Result};
+use anyhow::Result;
 use ariadne::Source;
 use clap::{Parser, Subcommand};
 use clio::Output;
@@ -191,7 +191,7 @@ impl Command {
                 let (source_id, source) = sources.files.clone().into_iter().exactly_one()?;
                 let stmts = prql_to_pl(&source)?;
 
-                let context = semantic::resolve_only(stmts, None)?;
+                let context = semantic::resolve_single(stmts, None)?;
 
                 let source_id = source_id.to_str().unwrap().to_string();
                 let references = label_references(&context, source_id, source);
@@ -207,7 +207,7 @@ impl Command {
                 let stmts = prql_to_pl(&source)?;
 
                 // resolve
-                let ctx = semantic::resolve_only(stmts, None)?;
+                let ctx = semantic::resolve_single(stmts, None)?;
 
                 let frames = if let Some((main, _)) = ctx.find_main(&[]) {
                     collect_frames(main.clone())
@@ -220,7 +220,7 @@ impl Command {
             }
             Command::Resolve { format, .. } => {
                 let ast = prql_to_pl_tree(sources)?;
-                let ir = pl_to_rq_tree(ast, main_path)?;
+                let ir = pl_to_rq_tree(ast, &main_path)?;
 
                 match format {
                     Format::Json => serde_json::to_string_pretty(&ir)?.into_bytes(),
@@ -238,7 +238,7 @@ impl Command {
                     .with_signature_comment(*include_signature_comment);
 
                 prql_to_pl_tree(sources)
-                    .and_then(|pl| pl_to_rq_tree(pl, main_path))
+                    .and_then(|pl| pl_to_rq_tree(pl, &main_path))
                     .and_then(|rq| rq_to_sql(rq, &opts))
                     .map_err(|e| e.composed(sources, opts.color))?
                     .as_bytes()
@@ -247,13 +247,13 @@ impl Command {
 
             Command::SQLPreprocess { .. } => {
                 let ast = prql_to_pl_tree(sources)?;
-                let rq = pl_to_rq_tree(ast, main_path)?;
+                let rq = pl_to_rq_tree(ast, &main_path)?;
                 let srq = prql_compiler::sql::internal::preprocess(rq)?;
                 format!("{srq:#?}").as_bytes().to_vec()
             }
             Command::SQLAnchor { format, .. } => {
                 let ast = prql_to_pl_tree(sources)?;
-                let rq = pl_to_rq_tree(ast, main_path)?;
+                let rq = pl_to_rq_tree(ast, &main_path)?;
                 let srq = prql_compiler::sql::internal::anchor(rq)?;
 
                 let json = serde_json::to_string_pretty(&srq)?;
@@ -297,16 +297,7 @@ impl Command {
 
         let file_tree = input.read_to_tree()?;
 
-        let mut main_path = io_args.main_path.clone();
-        if let Ok(only) = file_tree.files.iter().exactly_one() {
-            if main_path.is_none() {
-                main_path =
-                    Some(prql_compiler::semantic::os_path_to_prql_path(only.0.clone())?.join("."));
-            }
-        }
-
-        let main_path =
-            main_path.ok_or_else(|| anyhow!("Missing the path of the main pipeline"))?;
+        let main_path = io_args.main_path.clone().unwrap_or_default();
 
         Ok((file_tree, main_path))
     }

prql-compiler/src/lib.rs

@@ -127,7 +127,7 @@ pub static COMPILER_VERSION: Lazy<Version> = Lazy::new(|| {
 /// See [`sql::Options`](sql/struct.Options.html) and [`sql::Dialect`](sql/enum.Dialect.html) for options and supported SQL dialects.
 pub fn compile(prql: &str, options: &Options) -> Result<String, ErrorMessages> {
     parser::parse(prql)
-        .and_then(semantic::resolve)
+        .and_then(semantic::resolve_and_lower_single)
         .and_then(|rq| sql::compile(rq, options))
         .map_err(error::downcast)
         .map_err(|e| e.composed(&prql.into(), options.color))
@@ -256,15 +256,15 @@ pub fn prql_to_pl_tree(prql: &FileTree) -> Result<FileTree<Vec<ast::pl::Stmt>>,
 
 /// Perform semantic analysis and convert PL to RQ.
 pub fn pl_to_rq(pl: Vec<ast::pl::Stmt>) -> Result<ast::rq::Query, ErrorMessages> {
-    semantic::resolve(pl).map_err(error::downcast)
+    semantic::resolve_and_lower_single(pl).map_err(error::downcast)
 }
 
 /// Perform semantic analysis and convert PL to RQ.
 pub fn pl_to_rq_tree(
     pl: FileTree<Vec<ast::pl::Stmt>>,
-    main_path: Vec<String>,
+    main_path: &[String],
 ) -> Result<ast::rq::Query, ErrorMessages> {
-    semantic::resolve_tree(pl, main_path).map_err(error::downcast)
+    semantic::resolve_and_lower(pl, main_path).map_err(error::downcast)
 }
 
 /// Generate SQL from RQ.

prql-compiler/src/semantic/context.rs

@@ -489,6 +489,11 @@ impl Context {
         let decl = self.root_mod.get(&ident)?;
         decl.kind.as_query_def()
     }
+
+    /// Finds all main pipelines.
+    pub fn find_mains(&self) -> Vec<Ident> {
+        self.root_mod.find_by_suffix(NS_MAIN)
+    }
 }
 
 impl Default for DeclKind {

prql-compiler/src/semantic/lowering.rs

@@ -25,7 +25,7 @@ use super::NS_DEFAULT_DB;
 /// - transforms are not nested,
 /// - transforms have correct partition, window and sort set,
 /// - make sure there are no unresolved expressions.
-pub fn lower_to_ir(context: Context, main_path: &[String]) -> Result<Query> {
+pub fn lower_to_ir(context: Context, main_path: &[String]) -> Result<(Query, Context)> {
     // find main
     log::debug!("lookup for main pipeline in {main_path:?}");
     let (_, main_ident) = context
@@ -57,11 +57,12 @@ pub fn lower_to_ir(context: Context, main_path: &[String]) -> Result<Query> {
         }
     }
 
-    Ok(Query {
+    let query = Query {
         def,
         tables: l.table_buffer,
         relation: main_relation.unwrap(),
-    })
+    };
+    Ok((query, l.context))
 }
 
 fn extern_ref_to_relation(

prql-compiler/src/semantic/mod.rs

@@ -11,45 +11,53 @@ mod type_resolver;
 
 use anyhow::Result;
 use itertools::Itertools;
+use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 
 pub use self::context::Context;
 pub use self::module::Module;
+pub use lowering::lower_to_ir;
 
 use crate::ast::pl::{Lineage, LineageColumn, Stmt};
 use crate::ast::rq::Query;
 use crate::error::WithErrorInfo;
 use crate::{Error, FileTree};
 
 /// Runs semantic analysis on the query and lowers PL to RQ.
-pub fn resolve(statements: Vec<Stmt>) -> Result<Query> {
+pub fn resolve_and_lower_single(statements: Vec<Stmt>) -> Result<Query> {
     let context = load_std_lib();
 
     let context = resolver::resolve(statements, vec![], context)?;
 
-    let query = lowering::lower_to_ir(context, &[])?;
+    let (query, _) = lowering::lower_to_ir(context, &[])?;
 
     Ok(query)
 }
 
 /// Runs semantic analysis on the query and lowers PL to RQ.
-pub fn resolve_tree(file_tree: FileTree<Vec<Stmt>>, main_path: Vec<String>) -> Result<Query> {
-    let mut context = load_std_lib();
-
-    for (path, stmts) in normalize(file_tree)? {
-        context = resolver::resolve(stmts, path, context)?;
-    }
-
-    let query = lowering::lower_to_ir(context, &main_path)?;
+pub fn resolve_and_lower(file_tree: FileTree<Vec<Stmt>>, main_path: &[String]) -> Result<Query> {
+    let context = resolve(file_tree, None)?;
 
+    let (query, _) = lowering::lower_to_ir(context, main_path)?;
     Ok(query)
 }
 
 /// Runs semantic analysis on the query.
-pub fn resolve_only(statements: Vec<Stmt>, context: Option<Context>) -> Result<Context> {
-    let context = context.unwrap_or_else(load_std_lib);
+pub fn resolve_single(statements: Vec<Stmt>, context: Option<Context>) -> Result<Context> {
+    let tree = FileTree {
+        files: HashMap::from([(PathBuf::from(""), statements)]),
+    };
 
-    resolver::resolve(statements, vec![], context)
+    resolve(tree, context)
+}
+
+/// Runs semantic analysis on the query.
+pub fn resolve(file_tree: FileTree<Vec<Stmt>>, context: Option<Context>) -> Result<Context> {
+    let mut context = context.unwrap_or_else(load_std_lib);
+    for (path, stmts) in normalize(file_tree)? {
+        context = resolver::resolve(stmts, path, context)?;
+    }
+    Ok(context)
 }
 
 pub fn load_std_lib() -> Context {
@@ -144,11 +152,11 @@ mod test {
     use anyhow::Result;
     use insta::assert_yaml_snapshot;
 
-    use super::resolve;
+    use super::resolve_and_lower_single;
     use crate::{ast::rq::Query, parser::parse};
 
     fn parse_and_resolve(query: &str) -> Result<Query> {
-        resolve(parse(query)?)
+        resolve_and_lower_single(parse(query)?)
     }
 
     #[test]

prql-compiler/src/semantic/module.rs

@@ -367,6 +367,25 @@ impl Module {
         }
         r
     }
+
+    /// Recursively finds all declarations that end in suffix.
+    pub fn find_by_suffix(&self, suffix: &str) -> Vec<Ident> {
+        let mut res = Vec::new();
+
+        for (name, decl) in &self.names {
+            if let DeclKind::Module(module) = &decl.kind {
+                let nested = module.find_by_suffix(suffix);
+                res.extend(nested.into_iter().map(|x| x.prepend(name.clone())));
+                continue;
+            }
+
+            if name == suffix {
+                res.push(Ident::from_name(name));
+            }
+        }
+
+        res
+    }
 }
 
 impl std::fmt::Debug for Module {

prql-compiler/src/semantic/resolver.rs

@@ -938,10 +938,10 @@ mod test {
     use insta::assert_yaml_snapshot;
 
     use crate::ast::pl::{Expr, Lineage};
-    use crate::semantic::resolve_only;
+    use crate::semantic::resolve_single;
 
     fn parse_and_resolve(query: &str) -> Result<Expr> {
-        let ctx = resolve_only(crate::parser::parse(query)?, None)?;
+        let ctx = resolve_single(crate::parser::parse(query)?, None)?;
         let (main, _) = ctx.find_main(&[]).unwrap();
         Ok(main.clone())
     }

prql-compiler/src/semantic/transforms.rs

@@ -1036,7 +1036,7 @@ mod tests {
     use insta::assert_yaml_snapshot;
 
     use crate::parser::parse;
-    use crate::semantic::{resolve, resolve_only};
+    use crate::semantic::{resolve_and_lower_single, resolve_single};
 
     #[test]
     fn test_aggregate_positional_arg() {
@@ -1051,7 +1051,7 @@ mod tests {
         ",
         )
         .unwrap();
-        let result = resolve(query).unwrap();
+        let result = resolve_and_lower_single(query).unwrap();
         assert_yaml_snapshot!(result, @r###"
         ---
         def:
@@ -1105,7 +1105,7 @@ mod tests {
         ",
         )
         .unwrap();
-        let result = resolve(query);
+        let result = resolve_and_lower_single(query);
         assert!(result.is_err());
 
         // oops, two arguments
@@ -1116,7 +1116,7 @@ mod tests {
         ",
         )
         .unwrap();
-        let result = resolve(query);
+        let result = resolve_and_lower_single(query);
         assert!(result.is_err());
 
         // correct function call
@@ -1129,7 +1129,7 @@ mod tests {
         ",
         )
         .unwrap();
-        let ctx = resolve_only(query, None).unwrap();
+        let ctx = resolve_single(query, None).unwrap();
         let res = ctx.find_main(&[]).unwrap().clone();
         assert_yaml_snapshot!(res, @r###"
         ---
@@ -1221,7 +1221,7 @@ mod tests {
         )
         .unwrap();
 
-        let result = resolve(query).unwrap();
+        let result = resolve_and_lower_single(query).unwrap();
         assert_yaml_snapshot!(result, @r###"
         ---
         def:

prql-compiler/src/sql/srq/mod.rs

@@ -25,10 +25,10 @@ mod test {
     use super::*;
 
     use crate::sql::{Context, Dialect};
-    use crate::{parser::parse, semantic::resolve};
+    use crate::{parser::parse, semantic::resolve_and_lower_single};
 
     fn parse_and_resolve(prql: &str) -> Result<(SqlQuery, Context)> {
-        let query = resolve(parse(prql)?)?;
+        let query = resolve_and_lower_single(parse(prql)?)?;
 
         compile_query(query, Some(Dialect::Generic))
     }

prql-compiler/src/sql/std.rs

@@ -22,7 +22,7 @@ fn load_std_impl() -> semantic::Module {
         ..semantic::Context::default()
     };
 
-    let context = semantic::resolve_only(statements, Some(context)).unwrap();
+    let context = semantic::resolve_single(statements, Some(context)).unwrap();
     context.root_mod
 }
 

prql-compiler/src/tests/test.rs

@@ -1528,7 +1528,7 @@ derive mng_name = s"managers.first_name || ' ' || managers.last_name"
 select [mng_name, managers.gender, salary_avg, salary_sd]"#;
 
     let sql_from_prql = parse(original_prql)
-        .and_then(crate::semantic::resolve)
+        .and_then(crate::semantic::resolve_and_lower_single)
         .and_then(|rq| sql::compile(rq, &Options::default()))
         .unwrap();