From 42e2222ee43adbfab094d6ebf012350334e49bd7 Mon Sep 17 00:00:00 2001 From: Chris Agostino Date: Wed, 2 Sep 2026 22:04:16 -0500 Subject: [PATCH] updating tests, fixing bigquery --- .gitignore | 2 + Cargo.lock | 51 +++++++++- Cargo.toml | 5 +- README.md | 14 ++- src/compiler.rs | 200 ++++++++++++++++++++++++++++------------ src/lib.rs | 4 +- src/main.rs | 66 ++++++------- src/parser.rs | 25 +++-- src/runner.rs | 100 ++++++++++++++++---- src/sqlite_udf.rs | 143 ++++++++++++++++++++-------- tests/compiler_tests.rs | 56 +++++++++++ 11 files changed, 503 insertions(+), 163 deletions(-) create mode 100644 tests/compiler_tests.rs diff --git a/.gitignore b/.gitignore index 2f7896d..b9f381a 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ target/ +example/ +*.db diff --git a/Cargo.lock b/Cargo.lock index 98d8550..2d88769 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -188,6 +188,16 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -200,6 +210,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -588,6 +604,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.1" @@ -625,7 +647,7 @@ dependencies = [ [[package]] name = "nql" -version = "0.1.0" +version = "0.1.1" dependencies = [ "clap", "regex", @@ -634,6 +656,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", + "tempfile", "thiserror 1.0.69", ] @@ -902,6 +925,19 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.37" @@ -1094,6 +1130,19 @@ dependencies = [ "syn", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.3.4", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "1.0.69" diff --git a/Cargo.toml b/Cargo.toml index b2a6e31..ba112ff 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "nql" -version = "0.1.0" +version = "0.1.1" edition = "2021" authors = ["NPC Worldwide "] description = "NQL — SQL compiler with AI function calls targeting multiple databases" @@ -25,3 +25,6 @@ serde_yaml = "0.9" reqwest = { version = "0.12", features = ["blocking", "json", "rustls-tls"], default-features = false } clap = { version = "4", features = ["derive"] } thiserror = "1" + +[dev-dependencies] +tempfile = "3" diff --git a/README.md b/README.md index dab19a6..e99bc7e 100644 --- a/README.md +++ b/README.md @@ -6,22 +6,20 @@ ## Install -### Homebrew - -``` -brew install npc-worldwide/tap/nql -``` - ### Cargo ``` cargo install nql ``` -### pip +### Pre-built binaries + +Download the latest release from [GitHub releases](https://github.com/npc-worldwide/nql/releases). + +### Homebrew (planned) ``` -pip install nql +brew install npc-worldwide/tap/nql ``` ## Use diff --git a/src/compiler.rs b/src/compiler.rs index dc7f716..6d5ec27 100644 --- a/src/compiler.rs +++ b/src/compiler.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use crate::parser::{NqlCall, NqlModel}; +use crate::parser::NqlModel; use regex::Regex; /// Supported compilation targets. @@ -13,8 +13,10 @@ pub enum Target { Databricks, } -impl Target { - pub fn from_str(s: &str) -> Result { +impl std::str::FromStr for Target { + type Err = String; + + fn from_str(s: &str) -> Result { match s.to_lowercase().as_str() { "sqlite" => Ok(Target::Sqlite), "postgresql" | "postgres" | "pg" => Ok(Target::Postgresql), @@ -52,7 +54,8 @@ impl Compiler { /// Register a ref mapping: model_name -> fully qualified table name. pub fn register_ref(&mut self, model_name: &str, table_name: &str) { - self.ref_map.insert(model_name.to_string(), table_name.to_string()); + self.ref_map + .insert(model_name.to_string(), table_name.to_string()); } /// Compile a raw SQL string (with nql.* calls and {{ ref() }}) to target SQL. @@ -76,25 +79,37 @@ impl Compiler { let target_name = self.resolve_model_name(model); match model.config.materialized.as_str() { "table" => { - sql = format!( - "CREATE TABLE IF NOT EXISTS {} AS\n{}", - target_name, sql - ); + sql = format!("CREATE TABLE IF NOT EXISTS {} AS\n{}", target_name, sql); } "view" => { - sql = format!( - "CREATE OR REPLACE VIEW {} AS\n{}", - target_name, sql - ); + sql = format!("CREATE OR REPLACE VIEW {} AS\n{}", target_name, sql); + } + "incremental" => { + match model + .config + .extra + .get("unique_key") + .and_then(|v| v.as_str()) + { + Some(key) => { + sql = format!( + "CREATE TABLE IF NOT EXISTS {} AS\n(SELECT * FROM (\n{}\n) WHERE FALSE);\nINSERT INTO {}\nSELECT s.* FROM (\n{}\n) s\nWHERE NOT EXISTS (SELECT 1 FROM {} t WHERE t.{} = s.{});", + target_name, sql, target_name, sql, target_name, key, key + ); + } + None => { + sql = format!( + "-- incremental materialization requires unique_key in config\nCREATE TABLE IF NOT EXISTS {} AS\n{}", + target_name, sql + ); + } + } } "ephemeral" => { // No wrapping; used as CTE in downstream models } other => { - sql = format!( - "-- Unknown materialization: {}\n{}", - other, sql - ); + sql = format!("-- Unknown materialization: {}\n{}", other, sql); } } @@ -122,15 +137,17 @@ impl Compiler { let prefixed = format!(r"nql\.({})\(", funcs); let re_strip = Regex::new(&prefixed).expect("Invalid NQL prefix regex"); - let stripped = re_strip.replace_all(sql, |caps: ®ex::Captures| { - format!("{}(", &caps[1]) - }).to_string(); + let stripped = re_strip + .replace_all(sql, |caps: ®ex::Captures| format!("{}(", &caps[1])) + .to_string(); let bare = format!(r"\b({})\(([^)]*)\)", funcs); let re_bare = Regex::new(&bare).expect("Invalid bare NQL regex"); - re_bare.replace_all(&stripped, |caps: ®ex::Captures| { - self.translate_function(&caps[1], &caps[2], target) - }).to_string() + re_bare + .replace_all(&stripped, |caps: ®ex::Captures| { + self.translate_function(&caps[1], &caps[2], target) + }) + .to_string() } /// Translate a single nql.* function call to target SQL. @@ -158,11 +175,12 @@ impl Compiler { match func_name { "generate_text" => { let col = arg_parts.first().copied().unwrap_or("''"); - let prompt = if arg_parts.len() > 1 { arg_parts[1] } else { "''" }; - format!( - "pgai.generate_text(prompt => {}, text => {})", - prompt, col - ) + let prompt = if arg_parts.len() > 1 { + arg_parts[1] + } else { + "''" + }; + format!("pgai.generate_text(prompt => {}, text => {})", prompt, col) } "summarize" => { let col = arg_parts.first().copied().unwrap_or("''"); @@ -174,7 +192,11 @@ impl Compiler { } "translate" => { let col = arg_parts.first().copied().unwrap_or("''"); - let lang = if arg_parts.len() > 1 { arg_parts[1] } else { "'en'" }; + let lang = if arg_parts.len() > 1 { + arg_parts[1] + } else { + "'en'" + }; format!( "pgai.translate(text => {}, target_language => {})", col, lang @@ -202,63 +224,110 @@ impl Compiler { "translate" => { let arg_parts: Vec<&str> = args.split(',').map(|a| a.trim()).collect(); let col = arg_parts.first().copied().unwrap_or("''"); - let lang = if arg_parts.len() > 1 { arg_parts[1] } else { "'en'" }; + let lang = if arg_parts.len() > 1 { + arg_parts[1] + } else { + "'en'" + }; format!("SNOWFLAKE.CORTEX.TRANSLATE({}, {}, 'en')", col, lang) } - "extract_entities" => format!("SNOWFLAKE.CORTEX.EXTRACT_ANSWER({}, 'Extract all entities')", args), - "generate_embedding" => format!("SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m-v1.5', {})", args), + "extract_entities" => format!( + "SNOWFLAKE.CORTEX.EXTRACT_ANSWER({}, 'Extract all entities')", + args + ), + "generate_embedding" => format!( + "SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m-v1.5', {})", + args + ), _ => format!("SNOWFLAKE.CORTEX.COMPLETE('llama3.1-8b', {})", args), } } - // ── BigQuery: ML functions ───────────────────────────────────────── + // ── BigQuery: built-in AI functions ──────────────────────────────── fn translate_bigquery(&self, func_name: &str, args: &str) -> String { match func_name { "generate_text" => { - format!( - "ML.GENERATE_TEXT(MODEL `nql_model`, (SELECT {} AS prompt), STRUCT(256 AS max_output_tokens))", - args - ) + format!("(AI.GENERATE({})).result", args) } "summarize" => { - format!( - "ML.GENERATE_TEXT(MODEL `nql_model`, (SELECT CONCAT('Summarize: ', {}) AS prompt), STRUCT(512 AS max_output_tokens))", - args - ) + format!("(AI.GENERATE(CONCAT('Summarize: ', {}))).result", args) } "analyze_sentiment" => { format!( - "ML.GENERATE_TEXT(MODEL `nql_model`, (SELECT CONCAT('Analyze sentiment: ', {}) AS prompt), STRUCT(64 AS max_output_tokens))", + "(AI.GENERATE(CONCAT('Analyze sentiment: ', {}))).result", args ) } "translate" => { let arg_parts: Vec<&str> = args.split(',').map(|a| a.trim()).collect(); let col = arg_parts.first().copied().unwrap_or("''"); - let lang = if arg_parts.len() > 1 { arg_parts[1] } else { "'en'" }; + let lang = if arg_parts.len() > 1 { + arg_parts[1] + } else { + "'en'" + }; format!( - "ML.GENERATE_TEXT(MODEL `nql_model`, (SELECT CONCAT('Translate to ', {}, ': ', {}) AS prompt), STRUCT(512 AS max_output_tokens))", + "(AI.GENERATE(CONCAT('Translate to ', {}, ': ', {}))).result", lang, col ) } "extract_entities" => { format!( - "ML.GENERATE_TEXT(MODEL `nql_model`, (SELECT CONCAT('Extract entities from: ', {}) AS prompt), STRUCT(256 AS max_output_tokens))", + "(AI.GENERATE(CONCAT('Extract entities from: ', {}))).result", args ) } "generate_embedding" => { - format!( - "ML.GENERATE_TEXT_EMBEDDING(MODEL `nql_embedding_model`, (SELECT {} AS content), STRUCT(TRUE AS flatten_json_output))", - args - ) + format!("(AI.GENERATE_EMBEDDING({})).result", args) + } + "sentiment" => { + format!("(AI.GENERATE(CONCAT('Analyze the sentiment of the following text. Respond with exactly one word: positive, negative, or neutral.\\n\\n', {}))).result", args) + } + "get_facts" => { + format!("(AI.GENERATE(CONCAT('Extract facts from this text. A fact is a specific statement that can be sourced from the text. Return as JSON array of objects with \"statement\", \"source_text\", and \"type\" (explicit or inferred) fields.\\n\\nText: ', {}))).result", args) + } + "identify_groups" => { + format!("(AI.GENERATE(CONCAT('What are the main groups these items could be organized into? Return as JSON array of group names.\\n\\nItems: ', {}))).result", args) + } + "classify" => { + format!("(AI.GENERATE(CONCAT('Classify the following text into a category. Return only the category name.\\n\\n', {}))).result", args) + } + "classify_into" => { + let arg_parts: Vec<&str> = args.split(',').map(|a| a.trim()).collect(); + let col = arg_parts.first().copied().unwrap_or("''"); + let categories = arg_parts.get(1).copied().unwrap_or("''"); + format!("(AI.GENERATE(CONCAT('Classify the following text into one of these categories: ', {}, '. Return only the category name.\\n\\n', {}))).result", categories, col) + } + "extract_json" => { + format!("(AI.GENERATE(CONCAT('Extract structured data from this text and return as valid JSON:\\n\\n', {}))).result", args) + } + "detect_language" => { + format!("(AI.GENERATE(CONCAT('Detect the language of this text. Return only the ISO 639-1 language code.\\n\\n', {}))).result", args) + } + "answer_question" => { + let arg_parts: Vec<&str> = args.split(',').map(|a| a.trim()).collect(); + let context = arg_parts.first().copied().unwrap_or("''"); + let question = arg_parts.get(1).copied().unwrap_or("''"); + format!("(AI.GENERATE(CONCAT('Based on the following context, answer the question.\\n\\nContext: ', {}, '\\n\\nQuestion: ', {}))).result", context, question) + } + "generate_code" => { + format!("(AI.GENERATE(CONCAT('Generate code for the following task. Return only the code, no explanation.\\n\\n', {}))).result", args) + } + "criticize" => { + format!("(AI.GENERATE(CONCAT('Provide a critical analysis and constructive criticism of the following, focused on weaknesses, improvements, and alternatives:\\n\\n', {}))).result", args) + } + "synthesize" => { + format!("(AI.GENERATE(CONCAT('Synthesize this content into a clear, concise summary that captures the essence:\\n\\n', {}))).result", args) + } + "breathe" => { + format!("(AI.GENERATE(CONCAT('Read the following and identify the high level objective, most recent task, accomplishments, and failures. Return as JSON with keys: high_level_objective, most_recent_task, accomplishments, failures.\\n\\n', {}))).result", args) + } + "zoom_in" => { + format!("(AI.GENERATE(CONCAT('Look at these facts and infer new implied facts. Return as JSON array of objects with \"statement\" and \"inferred_from\" fields.\\n\\n', {}))).result", args) } _ => { - format!( - "ML.GENERATE_TEXT(MODEL `nql_model`, (SELECT {} AS prompt), STRUCT(256 AS max_output_tokens))", - args - ) + format!("(AI.GENERATE({})).result", args) } } } @@ -269,12 +338,22 @@ impl Compiler { match func_name { "generate_text" => { let col = arg_parts.first().copied().unwrap_or("''"); - let prompt = if arg_parts.len() > 1 { arg_parts[1] } else { "''" }; - format!("ai_query('databricks-meta-llama-3-1-70b-instruct', CONCAT({}, ' ', {}))", prompt, col) + let prompt = if arg_parts.len() > 1 { + arg_parts[1] + } else { + "''" + }; + format!( + "ai_query('databricks-meta-llama-3-1-70b-instruct', CONCAT({}, ' ', {}))", + prompt, col + ) } "summarize" => { let col = arg_parts.first().copied().unwrap_or("''"); - format!("ai_query('databricks-meta-llama-3-1-70b-instruct', CONCAT('Summarize: ', {}))", col) + format!( + "ai_query('databricks-meta-llama-3-1-70b-instruct', CONCAT('Summarize: ', {}))", + col + ) } "analyze_sentiment" => { let col = arg_parts.first().copied().unwrap_or("''"); @@ -282,7 +361,11 @@ impl Compiler { } "translate" => { let col = arg_parts.first().copied().unwrap_or("''"); - let lang = if arg_parts.len() > 1 { arg_parts[1] } else { "'en'" }; + let lang = if arg_parts.len() > 1 { + arg_parts[1] + } else { + "'en'" + }; format!("ai_query('databricks-meta-llama-3-1-70b-instruct', CONCAT('Translate to ', {}, ': ', {}))", lang, col) } "extract_entities" => { @@ -293,7 +376,10 @@ impl Compiler { let col = arg_parts.first().copied().unwrap_or("''"); format!("ai_query('databricks-bge-large-en', {})", col) } - _ => format!("ai_query('databricks-meta-llama-3-1-70b-instruct', {})", args), + _ => format!( + "ai_query('databricks-meta-llama-3-1-70b-instruct', {})", + args + ), } } diff --git a/src/lib.rs b/src/lib.rs index f27f5d4..9b34d8b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,4 @@ -pub mod parser; pub mod compiler; -pub mod sqlite_udf; +pub mod parser; pub mod runner; +pub mod sqlite_udf; diff --git a/src/main.rs b/src/main.rs index b9defe4..91e2162 100644 --- a/src/main.rs +++ b/src/main.rs @@ -58,7 +58,7 @@ fn main() { models_dir, target, } => { - let target = match Target::from_str(&target) { + let target = match target.parse::() { Ok(t) => t, Err(e) => { eprintln!("Error: {}", e); @@ -85,17 +85,15 @@ fn main() { ); match model_name { - Some(name) => { - match runner.run_single(&name, db.as_deref()) { - Ok(sql) => { - println!("{}", sql); - } - Err(e) => { - eprintln!("Error running model '{}': {}", name, e); - process::exit(1); - } + Some(name) => match runner.run_single(&name, db.as_deref()) { + Ok(sql) => { + println!("{}", sql); } - } + Err(e) => { + eprintln!("Error running model '{}': {}", name, e); + process::exit(1); + } + }, None => { if target == Target::Sqlite { let db_path = db.as_deref().unwrap_or(":memory:"); @@ -104,6 +102,12 @@ fn main() { process::exit(1); } eprintln!("[nql] All models executed successfully."); + } else if target == Target::Bigquery { + if let Err(e) = runner.execute_bigquery() { + eprintln!("Error executing models on BigQuery: {}", e); + process::exit(1); + } + eprintln!("[nql] All models executed successfully on BigQuery."); } else { match runner.compile_all() { Ok(results) => { @@ -128,7 +132,7 @@ fn main() { models_dir, target, } => { - let target = match Target::from_str(&target) { + let target = match target.parse::() { Ok(t) => t, Err(e) => { eprintln!("Error: {}", e); @@ -149,30 +153,26 @@ fn main() { } match model_name { - Some(name) => { - match runner.run_single(&name, None) { - Ok(sql) => println!("{}", sql), - Err(e) => { - eprintln!("Error compiling model '{}': {}", name, e); - process::exit(1); - } + Some(name) => match runner.run_single(&name, None) { + Ok(sql) => println!("{}", sql), + Err(e) => { + eprintln!("Error compiling model '{}': {}", name, e); + process::exit(1); } - } - None => { - match runner.compile_all() { - Ok(results) => { - for (name, sql) in results { - println!("-- ═══ Model: {} ═══", name); - println!("{}", sql); - println!(); - } - } - Err(e) => { - eprintln!("Error compiling models: {}", e); - process::exit(1); + }, + None => match runner.compile_all() { + Ok(results) => { + for (name, sql) in results { + println!("-- ═══ Model: {} ═══", name); + println!("{}", sql); + println!(); } } - } + Err(e) => { + eprintln!("Error compiling models: {}", e); + process::exit(1); + } + }, } } } diff --git a/src/parser.rs b/src/parser.rs index 7fcbf3d..f78dd0d 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -54,6 +54,19 @@ pub const NQL_FUNCTIONS: &[&str] = &[ "translate", "extract_entities", "generate_embedding", + "sentiment", + "get_facts", + "identify_groups", + "classify", + "classify_into", + "extract_json", + "detect_language", + "answer_question", + "generate_code", + "criticize", + "synthesize", + "breathe", + "zoom_in", ]; /// Parse a .sql model file from a path. @@ -107,8 +120,7 @@ fn extract_frontmatter(content: &str) -> Result { } if in_config { - if trimmed.starts_with("--") { - let after_dashes = &trimmed[2..]; + if let Some(after_dashes) = trimmed.strip_prefix("--") { // Check if this is still an indented config line if after_dashes.starts_with(" ") || after_dashes.starts_with('\t') { yaml_lines.push(after_dashes.to_string()); @@ -161,8 +173,7 @@ fn strip_frontmatter(content: &str) -> String { } if in_config { - if trimmed.starts_with("--") { - let after = &trimmed[2..]; + if let Some(after) = trimmed.strip_prefix("--") { if after.starts_with(" ") || after.starts_with('\t') { continue; } @@ -218,7 +229,7 @@ pub fn extract_nql_calls(sql: &str) -> Vec { for cap in re_bare.captures_iter(sql) { let full = cap[0].to_string(); - let prefixed_form = format!("nql.{}", &full); + let prefixed_form = format!("nql.{}", full); if !seen.contains(&prefixed_form) && seen.insert(full.clone()) { calls.push(NqlCall { full_match: full, @@ -233,8 +244,8 @@ pub fn extract_nql_calls(sql: &str) -> Vec { /// Extract all {{ ref('...') }} references from SQL text. pub fn extract_refs(sql: &str) -> Vec { - let re = Regex::new(r#"\{\{\s*ref\(\s*['"]([^'"]+)['"]\s*\)\s*\}\}"#) - .expect("Invalid ref regex"); + let re = + Regex::new(r#"\{\{\s*ref\(\s*['"]([^'"]+)['"]\s*\)\s*\}\}"#).expect("Invalid ref regex"); re.captures_iter(sql) .map(|cap| cap[1].to_string()) diff --git a/src/runner.rs b/src/runner.rs index a0e2162..e652c80 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -71,7 +71,13 @@ impl Runner { for name in &model_names { if !visited.contains(name) { - self.topo_sort(name, &model_names, &mut visited, &mut in_progress, &mut order)?; + self.topo_sort( + name, + &model_names, + &mut visited, + &mut in_progress, + &mut order, + )?; } } @@ -142,10 +148,57 @@ impl Runner { let compiled_sql = compiler.compile_model(model); eprintln!("[nql] Executing model: {}", model_name); eprintln!("[nql] SQL:\n{}\n", compiled_sql); - conn.execute_batch(&compiled_sql) - .map_err(|e| RunnerError::Compilation(format!( - "Failed to execute model '{}': {}", model_name, e - )))?; + conn.execute_batch(&compiled_sql).map_err(|e| { + RunnerError::Compilation(format!( + "Failed to execute model '{}': {}", + model_name, e + )) + })?; + } + } + + Ok(()) + } + + /// Compile all models and execute them against BigQuery via the `bq` CLI, + /// in dependency order. + pub fn execute_bigquery(&self) -> Result<(), RunnerError> { + let results = self.compile_all()?; + + let mut bq_args: Vec = vec![ + "query".to_string(), + "--use_legacy_sql=false".to_string(), + "--format=none".to_string(), + ]; + if let Ok(project) = std::env::var("NQL_BQ_PROJECT") { + bq_args.push(format!("--project_id={}", project)); + } + if let Ok(location) = std::env::var("NQL_BQ_LOCATION") { + bq_args.push(format!("--location={}", location)); + } + + for (name, sql) in results { + eprintln!("[nql] Executing model on BigQuery: {}", name); + let mut child = std::process::Command::new("bq") + .args(&bq_args) + .stdin(std::process::Stdio::piped()) + .spawn() + .map_err(|e| RunnerError::Compilation(format!("Failed to spawn bq CLI: {}", e)))?; + if let Some(mut stdin) = child.stdin.take() { + use std::io::Write; + stdin.write_all(sql.as_bytes()).map_err(|e| { + RunnerError::Compilation(format!("Failed to write SQL to bq stdin: {}", e)) + })?; + } + let status = child + .wait_with_output() + .map_err(|e| RunnerError::Compilation(format!("Failed waiting for bq: {}", e)))?; + if !status.status.success() { + return Err(RunnerError::Compilation(format!( + "bq query failed for model '{}': {}", + name, + String::from_utf8_lossy(&status.stderr) + ))); } } @@ -179,13 +232,18 @@ impl Runner { } /// Execute or compile a single named model. - pub fn run_single(&self, model_name: &str, db_path: Option<&str>) -> Result { - let model = self.models.get(model_name).ok_or_else(|| { - RunnerError::Parse { + pub fn run_single( + &self, + model_name: &str, + db_path: Option<&str>, + ) -> Result { + let model = self + .models + .get(model_name) + .ok_or_else(|| RunnerError::Parse { model: model_name.to_string(), message: "Model not found".to_string(), - } - })?; + })?; let mut compiler = Compiler::new(self.target); @@ -210,10 +268,12 @@ impl Runner { rusqlite::Connection::open(path)? }; sqlite_udf::register_nql_functions(&conn)?; - conn.execute_batch(&compiled) - .map_err(|e| RunnerError::Compilation(format!( - "Failed to execute model '{}': {}", model_name, e - )))?; + conn.execute_batch(&compiled).map_err(|e| { + RunnerError::Compilation(format!( + "Failed to execute model '{}': {}", + model_name, e + )) + })?; } } @@ -267,8 +327,16 @@ mod tests { runner.resolve_dependencies().unwrap(); // A should come before B - let a_pos = runner.execution_order.iter().position(|n| n == "a").unwrap(); - let b_pos = runner.execution_order.iter().position(|n| n == "b").unwrap(); + let a_pos = runner + .execution_order + .iter() + .position(|n| n == "a") + .unwrap(); + let b_pos = runner + .execution_order + .iter() + .position(|n| n == "b") + .unwrap(); assert!(a_pos < b_pos); let _ = fs::remove_dir_all(&tmp); diff --git a/src/sqlite_udf.rs b/src/sqlite_udf.rs index 87ed8ac..89f2385 100644 --- a/src/sqlite_udf.rs +++ b/src/sqlite_udf.rs @@ -27,13 +27,15 @@ impl Default for LlmConfig { impl LlmConfig { pub fn from_env() -> Self { let backend = std::env::var("NQL_LLM_BACKEND").unwrap_or_else(|_| "ollama".to_string()); - let api_key = std::env::var("OPENAI_API_KEY").ok() + let api_key = std::env::var("OPENAI_API_KEY") + .ok() .or_else(|| std::env::var("ANTHROPIC_API_KEY").ok()) .or_else(|| std::env::var("GEMINI_API_KEY").ok()); let (api_url, model) = match backend.as_str() { "openai" => ( - std::env::var("NQL_API_URL").unwrap_or_else(|_| "https://api.openai.com/v1/chat/completions".to_string()), + std::env::var("NQL_API_URL") + .unwrap_or_else(|_| "https://api.openai.com/v1/chat/completions".to_string()), std::env::var("NQL_MODEL").unwrap_or_else(|_| "gpt-4o-mini".to_string()), ), _ => ( @@ -42,7 +44,12 @@ impl LlmConfig { ), }; - LlmConfig { backend, api_url, model, api_key } + LlmConfig { + backend, + api_url, + model, + api_key, + } } } @@ -54,46 +61,93 @@ fn call_llm(config: &LlmConfig, prompt: &str) -> Result { } } -fn call_ollama(client: &reqwest::blocking::Client, config: &LlmConfig, prompt: &str) -> Result { +fn call_ollama( + client: &reqwest::blocking::Client, + config: &LlmConfig, + prompt: &str, +) -> Result { let body = serde_json::json!({"model": config.model, "prompt": prompt, "stream": false}); - let resp = client.post(&config.api_url).json(&body).send().map_err(|e| format!("Ollama: {}", e))?; + let resp = client + .post(&config.api_url) + .json(&body) + .send() + .map_err(|e| format!("Ollama: {}", e))?; let json: serde_json::Value = resp.json().map_err(|e| format!("Parse: {}", e))?; - json["response"].as_str().map(|s| s.trim().to_string()).ok_or_else(|| "No response".to_string()) + json["response"] + .as_str() + .map(|s| s.trim().to_string()) + .ok_or_else(|| "No response".to_string()) } -fn call_openai(client: &reqwest::blocking::Client, config: &LlmConfig, prompt: &str) -> Result { +fn call_openai( + client: &reqwest::blocking::Client, + config: &LlmConfig, + prompt: &str, +) -> Result { let api_key = config.api_key.as_deref().ok_or("API key not set")?; let body = serde_json::json!({"model": config.model, "messages": [{"role": "user", "content": prompt}], "max_tokens": 512}); - let resp = client.post(&config.api_url).header("Authorization", format!("Bearer {}", api_key)).json(&body).send().map_err(|e| format!("OpenAI: {}", e))?; + let resp = client + .post(&config.api_url) + .header("Authorization", format!("Bearer {}", api_key)) + .json(&body) + .send() + .map_err(|e| format!("OpenAI: {}", e))?; let json: serde_json::Value = resp.json().map_err(|e| format!("Parse: {}", e))?; - json["choices"][0]["message"]["content"].as_str().map(|s| s.trim().to_string()).ok_or_else(|| "No content".to_string()) + json["choices"][0]["message"]["content"] + .as_str() + .map(|s| s.trim().to_string()) + .ok_or_else(|| "No content".to_string()) } -fn register_udf(conn: &Connection, name: &str, config: &LlmConfig, prompt_builder: F) -> rusqlite::Result<()> -where F: Fn(&str) -> String + Send + 'static { +fn register_udf( + conn: &Connection, + name: &str, + config: &LlmConfig, + prompt_builder: F, +) -> rusqlite::Result<()> +where + F: Fn(&str) -> String + Send + 'static, +{ let cfg = config.clone(); - conn.create_scalar_function(name, 1, FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC, move |ctx| { - let input: String = ctx.get(0)?; - let prompt = prompt_builder(&input); - match call_llm(&cfg, &prompt) { - Ok(result) => Ok(ToSqlOutput::Owned(Value::Text(result))), - Err(e) => Ok(ToSqlOutput::Owned(Value::Text(format!("ERROR: {}", e)))), - } - }) + conn.create_scalar_function( + name, + 1, + FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC, + move |ctx| { + let input: String = ctx.get(0)?; + let prompt = prompt_builder(&input); + match call_llm(&cfg, &prompt) { + Ok(result) => Ok(ToSqlOutput::Owned(Value::Text(result))), + Err(e) => Ok(ToSqlOutput::Owned(Value::Text(format!("ERROR: {}", e)))), + } + }, + ) } -fn register_udf2(conn: &Connection, name: &str, config: &LlmConfig, prompt_builder: F) -> rusqlite::Result<()> -where F: Fn(&str, &str) -> String + Send + 'static { +fn register_udf2( + conn: &Connection, + name: &str, + config: &LlmConfig, + prompt_builder: F, +) -> rusqlite::Result<()> +where + F: Fn(&str, &str) -> String + Send + 'static, +{ let cfg = config.clone(); - conn.create_scalar_function(name, 2, FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC, move |ctx| { - let input1: String = ctx.get(0)?; - let input2: String = ctx.get(1)?; - let prompt = prompt_builder(&input1, &input2); - match call_llm(&cfg, &prompt) { - Ok(result) => Ok(ToSqlOutput::Owned(Value::Text(result))), - Err(e) => Ok(ToSqlOutput::Owned(Value::Text(format!("ERROR: {}", e)))), - } - }) + conn.create_scalar_function( + name, + 2, + FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC, + move |ctx| { + let input1: String = ctx.get(0)?; + let input2: String = ctx.get(1)?; + let prompt = prompt_builder(&input1, &input2); + match call_llm(&cfg, &prompt) { + Ok(result) => Ok(ToSqlOutput::Owned(Value::Text(result))), + Err(e) => Ok(ToSqlOutput::Owned(Value::Text(format!("ERROR: {}", e)))), + } + }, + ) } pub fn register_all_functions(conn: &Connection) -> rusqlite::Result<()> { @@ -115,9 +169,7 @@ pub fn register_all_functions(conn: &Connection) -> rusqlite::Result<()> { format!("Extract all named entities (people, places, organizations) from the following text. Return as JSON array.\n\n{}", text) })?; - register_udf(conn, "generate_text", &config, |prompt| { - prompt.to_string() - })?; + register_udf(conn, "generate_text", &config, |prompt| prompt.to_string())?; register_udf(conn, "generate_embedding", &config, |text| { format!("Generate a semantic embedding description for: {}", text) @@ -137,7 +189,10 @@ pub fn register_all_functions(conn: &Connection) -> rusqlite::Result<()> { })?; register_udf(conn, "classify", &config, |text| { - format!("Classify the following text into a category. Return only the category name.\n\n{}", text) + format!( + "Classify the following text into a category. Return only the category name.\n\n{}", + text + ) })?; register_udf2(conn, "classify_into", &config, |text, categories| { @@ -145,19 +200,31 @@ pub fn register_all_functions(conn: &Connection) -> rusqlite::Result<()> { })?; register_udf(conn, "extract_json", &config, |text| { - format!("Extract structured data from this text and return as valid JSON:\n\n{}", text) + format!( + "Extract structured data from this text and return as valid JSON:\n\n{}", + text + ) })?; register_udf(conn, "detect_language", &config, |text| { - format!("Detect the language of this text. Return only the ISO 639-1 language code.\n\n{}", text) + format!( + "Detect the language of this text. Return only the ISO 639-1 language code.\n\n{}", + text + ) })?; register_udf2(conn, "answer_question", &config, |context, question| { - format!("Based on the following context, answer the question.\n\nContext: {}\n\nQuestion: {}", context, question) + format!( + "Based on the following context, answer the question.\n\nContext: {}\n\nQuestion: {}", + context, question + ) })?; register_udf(conn, "generate_code", &config, |prompt| { - format!("Generate code for the following task. Return only the code, no explanation.\n\n{}", prompt) + format!( + "Generate code for the following task. Return only the code, no explanation.\n\n{}", + prompt + ) })?; register_udf(conn, "criticize", &config, |text| { diff --git a/tests/compiler_tests.rs b/tests/compiler_tests.rs new file mode 100644 index 0000000..a06b1a6 --- /dev/null +++ b/tests/compiler_tests.rs @@ -0,0 +1,56 @@ +use nql::compiler::{Compiler, Target}; +use nql::parser::parse_model_file; +use std::io::Write; +use tempfile::NamedTempFile; + +#[test] +fn bigquery_get_facts_emits_ai_generate_result() { + let sql = "SELECT nql.get_facts(bio) AS facts FROM users"; + let compiler = Compiler::new(Target::Bigquery); + let out = compiler.compile(sql, Target::Bigquery); + assert!(out.contains("AI.GENERATE")); + assert!(out.contains(").result")); + assert!(out.contains("Extract facts")); +} + +#[test] +fn bigquery_summarize_emits_ai_generate() { + let sql = "SELECT nql.summarize(bio) AS s FROM users"; + let compiler = Compiler::new(Target::Bigquery); + let out = compiler.compile(sql, Target::Bigquery); + assert!(out.contains("AI.GENERATE(CONCAT('Summarize: ', bio))")); +} + +#[test] +fn bigquery_generate_embedding_emits_ai_generate_embedding() { + let sql = "SELECT nql.generate_embedding(bio) AS e FROM users"; + let compiler = Compiler::new(Target::Bigquery); + let out = compiler.compile(sql, Target::Bigquery); + assert!(out.contains("AI.GENERATE_EMBEDDING(bio)).result")); +} + +#[test] +fn incremental_materialization_emits_insert_not_exists() { + let mut f = NamedTempFile::new().unwrap(); + write!( + f, + "-- config:\n-- materialized: incremental\n-- unique_key: user_id\n\nSELECT user_id, name FROM users\n" + ) + .unwrap(); + let model = parse_model_file(f.path()).unwrap(); + let compiler = Compiler::new(Target::Bigquery); + let out = compiler.compile_model(&model); + assert!(out.contains("CREATE TABLE IF NOT EXISTS")); + assert!(out.contains("INSERT INTO")); + assert!(out.contains("NOT EXISTS")); + assert!(out.contains("t.user_id = s.user_id")); +} + +#[test] +fn ref_replaced_with_table_name() { + let sql = "SELECT * FROM {{ ref('users') }}"; + let mut compiler = Compiler::new(Target::Bigquery); + compiler.register_ref("users", "analytics.users"); + let out = compiler.compile(sql, Target::Bigquery); + assert!(out.contains("FROM analytics.users"), "got: {}", out); +}