feat(lutra): pull database schema into source (#4182)

lutra/lutra/Cargo.toml

@@ -20,14 +20,16 @@ required-features = ["cli"]
 [target.'cfg(not(target_family="wasm"))'.dependencies]
 anyhow = "1.0.79"
 clap = { version = "4.4.18", features = ["derive"], optional = true }
-connector_arrow = { version = "0.1.1", features = ["src_sqlite"] }
+connector_arrow = { version = "0.2.0", features = ["src_sqlite"] }
 rusqlite = { version = "0.30.0", features = ["bundled"] }
-arrow = { version = "49.0", features = ["prettyprint"] }
+arrow = { version = "49.0", features = [
+    "prettyprint",
+], default-features = false }
 env_logger = "0.10.2"
 log = "0.4.20"
 prqlc = { path = "../../prqlc/prqlc", default-features = false }
 walkdir = "2.4.0"
+itertools = "0.12.0"
 
 [target.'cfg(not(target_family="wasm"))'.dev-dependencies]
 insta = { version = "1.34", features = ["colors"] }
-itertools = "0.12.0"

lutra/lutra/src/cli.rs

@@ -15,6 +15,7 @@ fn main() {
     let res = match action.command {
         Action::Discover(cmd) => discover_and_print(cmd),
         Action::Execute(cmd) => execute_and_print(cmd),
+        Action::PullSchema(cmd) => pull_schema_and_print(cmd),
     };
 
     match res {
@@ -31,7 +32,7 @@ fn main() {
 #[cfg(not(target_family = "wasm"))]
 mod inner {
     use clap::{Parser, Subcommand};
-    use lutra::{CompileParams, DiscoverParams, ExecuteParams};
+    use lutra::{CompileParams, DiscoverParams, ExecuteParams, PullSchemaParams};
 
     #[derive(Parser)]
     pub struct Command {
@@ -46,6 +47,9 @@ mod inner {
 
         /// Discover, compile, execute
         Execute(ExecuteCommand),
+
+        /// Pull schema from data sources
+        PullSchema(PullSchemaCommand),
     }
 
     #[derive(clap::Parser)]
@@ -87,4 +91,59 @@ mod inner {
         }
         Ok(())
     }
+
+    #[derive(clap::Parser)]
+    pub struct PullSchemaCommand {
+        #[clap(flatten)]
+        discover: DiscoverParams,
+
+        #[clap(flatten)]
+        compile: CompileParams,
+
+        #[clap(flatten)]
+        execute: PullSchemaParams,
+    }
+
+    pub fn pull_schema_and_print(cmd: PullSchemaCommand) -> anyhow::Result<()> {
+        let project = lutra::discover(cmd.discover)?;
+
+        let project = lutra::compile(project, cmd.compile)?;
+
+        let db_mod_decl_id = project.database_module.def_id.unwrap();
+
+        let stmts = lutra::pull_schema(&project, cmd.execute)?;
+
+        use prqlc::ast::*;
+
+        let db_mod_name = project.database_module.path.last().cloned();
+        let new_module_def = ModuleDef {
+            name: db_mod_name.unwrap_or_default(),
+            stmts,
+        };
+        let mut new_module_stmt = Stmt::new(StmtKind::ModuleDef(new_module_def));
+
+        // pull annotations and other metadata from the resolved module tree
+        // TODO: this is not a good idea, because we have to restrict PL back into AST
+        //   also, we'd have to pull back doc comments
+        //   ideally, we'd edit only the module contents, not the module itself
+        let db_mod_path = &project.database_module.path;
+        let root_mod = &project.root_module.module;
+        if !db_mod_path.is_empty() {
+            let i = Ident::from_path(db_mod_path.clone());
+            let database_mod_decl = root_mod.get(&i).unwrap();
+
+            new_module_stmt.annotations = database_mod_decl
+                .annotations
+                .clone()
+                .into_iter()
+                .map(prqlc::semantic::ast_expand::restrict_annotation)
+                .collect();
+        }
+
+        let new_source = prqlc::pl_to_prql(vec![new_module_stmt])?;
+
+        lutra::editing::edit_source_file(&project, db_mod_decl_id, new_source)?;
+
+        Ok(())
+    }
 }

lutra/lutra/src/compile.rs

@@ -18,17 +18,17 @@ pub fn compile(mut project: ProjectDiscovered, _: CompileParams) -> Result<Proje
     let files = std::mem::take(&mut project.sources);
     let source_tree = SourceTree::new(files, Some(project.root_path.clone()));
 
-    let res = parse_and_compile(&source_tree, project);
+    let res = parse_and_compile(&source_tree);
 
-    Ok(res
+    let mut project = res
         .map_err(prqlc::downcast)
-        .map_err(|err| err.composed(&source_tree))?)
+        .map_err(|err| err.composed(&source_tree))?;
+
+    project.sources = source_tree;
+    Ok(project)
 }
 
-fn parse_and_compile(
-    source_tree: &SourceTree,
-    project: ProjectDiscovered,
-) -> Result<ProjectCompiled> {
+fn parse_and_compile(source_tree: &SourceTree) -> Result<ProjectCompiled> {
     let options = Options::default()
         .with_target(Target::Sql(Some(Dialect::SQLite)))
         .no_format()
@@ -54,9 +54,10 @@ fn parse_and_compile(
         queries.insert(main_ident, sql);
     }
     Ok(ProjectCompiled {
-        inner: project,
+        sources: SourceTree::default(), // placeholder
         queries,
         database_module,
+        root_module,
     })
 }
 
@@ -87,6 +88,8 @@ fn find_database_module(root_module: &mut RootModule) -> Result<DatabaseModule>
         .find(|x| prqlc::semantic::is_ident_or_func_call(&x.expr, &lutra_sqlite))
         .unwrap();
 
+    let def_id = decl.declared_at;
+
     // make sure that there is exactly one arg
     let arg = match &annotation.expr.kind {
         prqlc::ir::pl::ExprKind::Ident(_) => {
@@ -132,6 +135,7 @@ fn find_database_module(root_module: &mut RootModule) -> Result<DatabaseModule>
 
     Ok(DatabaseModule {
         path: db_module_fq.into_iter().collect(),
+        def_id,
         connection_params: SqliteConnectionParams { file_relative },
     })
 }

lutra/lutra/src/connection.rs

@@ -0,0 +1,17 @@
+use std::path::Path;
+
+use anyhow::Result;
+
+use crate::project::DatabaseModule;
+
+pub fn open(db: &DatabaseModule, project_root: &Path) -> Result<rusqlite::Connection> {
+    // convert relative to absolute path
+    let mut sqlite_file_abs = project_root.to_path_buf();
+    sqlite_file_abs.push(&db.connection_params.file_relative);
+    let sqlite_file_abs = sqlite_file_abs.as_os_str().to_str().unwrap();
+
+    // init SQLite
+    let sqlite_conn = rusqlite::Connection::open(sqlite_file_abs)?;
+
+    Ok(sqlite_conn)
+}

lutra/lutra/src/editing.rs

@@ -0,0 +1,41 @@
+use std::fs;
+
+use anyhow::Result;
+use prqlc::Error;
+
+use crate::ProjectCompiled;
+
+/// Edit a source file such that the source of the declaration with id `decl_id` is now `new_source`.
+pub fn edit_source_file(
+    project: &ProjectCompiled,
+    decl_id: usize,
+    new_source: String,
+) -> Result<()> {
+    let span = project.root_module.span_map.get(&decl_id);
+    let Some(span) = span else {
+        // TODO: bad error message, we should not mention decl ids
+        return Err(Error::new_simple(format!(
+            "cannot find where declaration {decl_id} came from"
+        ))
+        .into());
+    };
+
+    // retrieve file path, relative to project root
+    // this is safe, because the source_id must exist, right? It was created during parsing.
+    let file_path = project.sources.get_path(span.source_id).unwrap().clone();
+
+    // find original source
+    let mut source = project.sources.sources.get(&file_path).unwrap().clone();
+
+    // replace the text
+    source.replace_range(span.start..span.end, &new_source);
+
+    // reconstruct full path of the file
+    // sources are loaded from file-system, so root path must always exist
+    let mut file_path_full = project.sources.root.clone().unwrap();
+    file_path_full.extend(&file_path);
+
+    // write the new file contents
+    fs::write(file_path_full, &source)?;
+    Ok(())
+}

lutra/lutra/src/execute.rs

@@ -36,14 +36,9 @@ pub fn execute(project: ProjectCompiled, params: ExecuteParams) -> Result<Vec<(I
 fn execute_one(project: &ProjectCompiled, pipeline_ident: &Ident) -> Result<Relation> {
     log::info!("executing {pipeline_ident}");
     let db = &project.database_module;
+    let project_root = project.sources.root.clone().unwrap();
 
-    // convert relative to absolute path
-    let mut sqlite_file_abs = project.inner.root_path.clone();
-    sqlite_file_abs.push(&db.connection_params.file_relative);
-    let sqlite_file_abs = sqlite_file_abs.as_os_str().to_str().unwrap();
-
-    // init SQLite
-    let mut sqlite_conn = rusqlite::Connection::open(sqlite_file_abs)?;
+    let mut conn = crate::connection::open(db, &project_root)?;
 
     let Some(pipeline) = project.queries.get(pipeline_ident) else {
         return Err(
@@ -52,7 +47,7 @@ fn execute_one(project: &ProjectCompiled, pipeline_ident: &Ident) -> Result<Rela
     };
     log::debug!("executing sql: {pipeline}");
 
-    let batches = connector_arrow::query_one(&mut sqlite_conn, pipeline)?;
+    let batches = connector_arrow::query_one(&mut conn, pipeline)?;
 
     Ok(batches)
 }