use anyhow::{Context, Result, bail};
use colored::Colorize;
use cosq_core::config::Config;
use cosq_core::stored_query::{
StoredQuery, StoredQueryMetadata, find_stored_query, list_stored_queries, query_file_path,
user_queries_dir,
};
use crate::cli::QueriesCommands;
pub async fn run(cmd: QueriesCommands, quiet: bool) -> Result<()> {
match cmd {
QueriesCommands::List => list(),
QueriesCommands::Create { name, project } => create(&name, project),
QueriesCommands::Edit { name } => edit(&name),
QueriesCommands::Delete { name, yes } => delete(&name, yes),
QueriesCommands::Show { name } => show(&name),
QueriesCommands::Generate {
description,
db,
container,
project,
} => generate(description, db, container, project, quiet).await,
}
}
fn list() -> Result<()> {
let queries = list_stored_queries().unwrap_or_default();
if queries.is_empty() {
println!("No stored queries found.");
println!(
"\n Create one with: {}",
"cosq queries create <name>".cyan()
);
return Ok(());
}
println!(
"{} ({}):\n",
"Stored queries".bold(),
"~/.cosq/queries/".dimmed()
);
let max_name_len = queries.iter().map(|q| q.name.len()).max().unwrap_or(0);
for query in &queries {
let ai_badge = if query.metadata.generated_by.is_some() {
" (AI)".dimmed().to_string()
} else {
String::new()
};
println!(
" {:<width$} {}{}",
query.name.green().bold(),
query.metadata.description.dimmed(),
ai_badge,
width = max_name_len,
);
}
println!("\n{} queries found.", queries.len());
Ok(())
}
fn create(name: &str, project: bool) -> Result<()> {
let path = query_file_path(name, project)?;
if path.exists() {
bail!(
"Query '{}' already exists at {}. Use `cosq queries edit {}` to modify it.",
name,
path.display(),
name
);
}
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let template = StoredQueryMetadata {
description: "TODO: describe what this query does".to_string(),
database: None,
container: None,
steps: None,
params: Vec::new(),
template: None,
template_file: None,
generated_by: None,
generated_from: None,
};
let yaml = serde_yaml::to_string(&template)?;
let contents =
format!("---\n{yaml}---\n-- Write your Cosmos DB SQL query below\nSELECT * FROM c\n");
std::fs::write(&path, &contents)?;
println!("{} Created {}", "OK".green().bold(), path.display());
open_in_editor(&path)?;
Ok(())
}
fn edit(name: &str) -> Result<()> {
let path = find_query_path(name)?;
let _ = StoredQuery::load(&path).map_err(|e| anyhow::anyhow!("Query file has errors: {e}"))?;
open_in_editor(&path)?;
Ok(())
}
fn delete(name: &str, yes: bool) -> Result<()> {
let path = find_query_path(name)?;
if !yes {
let confirm =
inquire::Confirm::new(&format!("Delete query '{name}' at {}?", path.display()))
.with_default(false)
.prompt()
.context("confirmation cancelled")?;
if !confirm {
println!("Cancelled.");
return Ok(());
}
}
std::fs::remove_file(&path)?;
println!("{} Deleted query '{name}'.", "OK".green().bold());
Ok(())
}
fn show(name: &str) -> Result<()> {
let query =
find_stored_query(name).map_err(|e| anyhow::anyhow!("Query '{name}' not found: {e}"))?;
println!("{}", query.name.green().bold());
println!(" {} {}", "Description:".bold(), query.metadata.description);
if let Some(ref db) = query.metadata.database {
println!(" {} {}", "Database:".bold(), db);
}
if let Some(ref ctr) = query.metadata.container {
println!(" {} {}", "Container:".bold(), ctr);
}
let has_template = query.metadata.template.is_some() || query.metadata.template_file.is_some();
println!(
" {} {}",
"Template:".bold(),
if has_template { "yes" } else { "(none)" }
);
if let Some(ref generated) = query.metadata.generated_by {
println!(" {} {}", "Generated by:".bold(), generated);
}
if let Some(ref prompt) = query.metadata.generated_from {
println!(" {} \"{}\"", "Prompt:".bold(), prompt);
}
if !query.metadata.params.is_empty() {
println!("\n {}:", "Parameters".bold());
for param in &query.metadata.params {
let type_str = param.param_type.to_string();
let desc = param.description.as_deref().unwrap_or("");
let default_str = param
.default
.as_ref()
.map(|d| format!(" (default: {})", d))
.unwrap_or_default();
println!(
" {} <{}> {}{}",
format!("--{}", param.name).cyan(),
type_str,
desc,
default_str.dimmed(),
);
if let Some(ref choices) = param.choices {
let choices_str: Vec<String> = choices
.iter()
.map(|c| match c {
serde_json::Value::String(s) => s.clone(),
other => other.to_string(),
})
.collect();
println!(" {}: {}", "choices".dimmed(), choices_str.join(", "));
}
if let Some(min) = param.min {
print!(" {}: {min}", "min".dimmed());
}
if let Some(max) = param.max {
print!(" {}: {max}", "max".dimmed());
}
if param.min.is_some() || param.max.is_some() {
println!();
}
}
}
println!("\n {}:", "Query".bold());
for line in query.sql.lines() {
println!(" {}", line.dimmed());
}
if let Some(ref tmpl) = query.metadata.template {
println!("\n {}:", "Template".bold());
for line in tmpl.lines() {
println!(" {}", line.dimmed());
}
}
Ok(())
}
async fn generate(
description: Option<String>,
cli_db: Option<String>,
cli_container: Option<String>,
project: bool,
quiet: bool,
) -> Result<()> {
let mut config = Config::load()?;
let ai_config = config.ai.clone().ok_or_else(|| {
anyhow::anyhow!("AI is not configured. Run `cosq ai init` to set up an AI provider.")
})?;
let client = cosq_client::cosmos::CosmosClient::new(&config.account.endpoint).await?;
let (database, db_changed) =
super::common::resolve_database(&client, &mut config, cli_db, None).await?;
if db_changed {
config.save()?;
}
let containers = if let Some(ctr) = cli_container {
vec![ctr]
} else {
pick_containers_interactive(&client, &database).await?
};
let mut container_samples: Vec<(String, String)> = Vec::new();
for ctr in &containers {
if !quiet {
eprintln!("{}", format!("Sampling documents from {ctr}...").dimmed());
}
let sample_result = client
.query(&database, ctr, "SELECT TOP 3 * FROM c")
.await
.with_context(|| format!("failed to sample documents from {ctr}"))?;
let sample_json = if sample_result.documents.is_empty() {
"(container is empty)".to_string()
} else {
format_sample_documents(&sample_result.documents)
};
container_samples.push((ctr.clone(), sample_json));
}
let description = if let Some(desc) = description {
desc
} else {
eprintln!();
inquire::Text::new("Describe the query you want to generate:")
.prompt()
.context("input cancelled")?
};
let system_prompt = build_system_prompt(&database, &container_samples);
let user_prompt = format!("Generate a .cosq stored query for: {description}");
if !quiet {
eprintln!(
"{}",
format!("Generating via {}...", ai_config.provider.display_name()).dimmed()
);
}
let mut conversation_prompt = user_prompt;
let mut query = None;
let max_rounds = 3;
for round in 0..max_rounds {
let response =
cosq_client::ai::generate_text(&ai_config, &system_prompt, &conversation_prompt)
.await
.context("failed to generate query")?;
let content = strip_markdown_fences(&response);
match StoredQuery::parse("generated", &content) {
Ok(parsed) => {
query = Some(parsed);
break;
}
Err(parse_err) if round < max_rounds - 1 => {
let looks_like_query = content.contains("---")
&& (content.contains("SELECT") || content.contains("select"));
if looks_like_query {
eprintln!(
"\n{} {}",
"Generated query has errors, retrying:".yellow().bold(),
format!("{parse_err}").dimmed()
);
conversation_prompt = format!(
"Original request: {description}\n\n\
Your previous response had a parse error: {parse_err}\n\n\
Here was your response:\n{content}\n\n\
Fix the error and generate a valid .cosq file."
);
} else {
eprintln!();
for line in content.lines() {
if !line.trim().is_empty() {
eprintln!(" {}", line);
}
}
eprintln!();
let answer: String = inquire::Text::new("Your answer:")
.prompt()
.context("input cancelled")?;
conversation_prompt = format!(
"Original request: {description}\n\n\
You asked:\n{content}\n\n\
User answered: {answer}\n\n\
Now generate the .cosq file."
);
}
if !quiet {
eprintln!(
"{}",
format!("Generating via {}...", ai_config.provider.display_name()).dimmed()
);
}
}
Err(_) => {
bail!(
"Could not generate a valid query after {max_rounds} attempts. \
Try rephrasing your description or adding more detail."
);
}
}
}
let mut query = query.ok_or_else(|| {
anyhow::anyhow!("Could not generate a valid query. Try rephrasing your description.")
})?;
query.metadata.database = Some(database);
if !query.is_multi_step() && containers.len() == 1 {
query.metadata.container = Some(containers[0].clone());
}
let provider_info = match ai_config.effective_model() {
Some(model) => format!("{} ({})", ai_config.provider.display_name(), model),
None => ai_config.provider.display_name().to_string(),
};
query.metadata.generated_by = Some(provider_info);
query.metadata.generated_from = Some(description.clone());
let suggested_name = generate_filename(&description);
show_query_preview(&query, &suggested_name);
let name: String = inquire::Text::new("Query name:")
.with_default(&suggested_name)
.prompt()
.context("input cancelled")?;
query.name = name.clone();
let path = query_file_path(&name, project)?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let contents = query.to_file_contents()?;
std::fs::write(&path, &contents)?;
println!("{} Saved to {}", "OK".green().bold(), path.display());
let options = vec!["Run it now", "Open in editor", "Done"];
let action_str = inquire::Select::new("What next?", options.clone())
.prompt()
.context("selection cancelled")?;
let action = options.iter().position(|o| *o == action_str).unwrap();
match action {
0 => {
eprintln!();
super::run::run(super::run::RunArgs {
name: Some(name),
params: Vec::new(),
output: None,
db: None,
container: None,
template: None,
quiet,
})
.await?;
}
1 => {
open_in_editor(&path)?;
}
_ => {}
}
Ok(())
}
async fn pick_containers_interactive(
client: &cosq_client::cosmos::CosmosClient,
database: &str,
) -> Result<Vec<String>> {
let all_containers = client.list_containers(database).await?;
if all_containers.is_empty() {
bail!("No containers found in database '{database}'.");
}
if all_containers.len() == 1 {
eprintln!(
"{} {}",
"Using container:".bold(),
all_containers[0].green()
);
return Ok(all_containers);
}
let scope_options = vec!["Single container", "Multiple containers (multi-step query)"];
let mode = inquire::Select::new("Query scope:", scope_options)
.prompt()
.context("selection cancelled")?;
if mode == "Single container" {
let selection = inquire::Select::new("Select a container:", all_containers.clone())
.prompt()
.context("container selection cancelled")?;
Ok(vec![selection])
} else {
let selections = inquire::MultiSelect::new("Select containers:", all_containers.clone())
.prompt()
.context("container selection cancelled")?;
if selections.is_empty() {
bail!("No containers selected.");
}
let selected: Vec<String> = selections;
for ctr in &selected {
eprintln!(" {} {}", "▸".dimmed(), ctr.green());
}
Ok(selected)
}
}
fn show_query_preview(query: &StoredQuery, suggested_name: &str) {
eprintln!("\n{}", "Generated query:".bold());
eprintln!(" {} {}", "Name:".dimmed(), suggested_name.green());
eprintln!(
" {} {}",
"Description:".dimmed(),
query.metadata.description
);
if query.is_multi_step() {
eprintln!(" {}:", "Steps".dimmed());
for step in query.metadata.steps.as_ref().unwrap() {
eprintln!(" {} → {}", step.name.cyan(), step.container.dimmed());
}
}
if !query.metadata.params.is_empty() {
eprintln!(" {}:", "Parameters".dimmed());
for p in &query.metadata.params {
let default_str = p
.default
.as_ref()
.map(|d| format!(" (default: {d})"))
.unwrap_or_default();
eprintln!(
" --{} <{}>{}",
p.name,
p.param_type,
default_str.dimmed()
);
}
}
if query.is_multi_step() {
for step in query.metadata.steps.as_ref().unwrap() {
if let Some(sql) = query.step_queries.get(&step.name) {
eprintln!("\n {} {}:", "SQL".dimmed(), step.name.cyan());
for line in sql.lines() {
eprintln!(" {}", line.cyan());
}
}
}
} else {
eprintln!("\n {}:", "SQL".dimmed());
for line in query.sql.lines() {
eprintln!(" {}", line.cyan());
}
}
if let Some(ref tmpl) = query.metadata.template {
eprintln!("\n {}:", "Template".dimmed());
for line in tmpl.lines() {
eprintln!(" {}", line);
}
}
}
fn build_system_prompt(database: &str, container_samples: &[(String, String)]) -> String {
let is_multi = container_samples.len() > 1;
let mut container_section = String::new();
for (name, sample_json) in container_samples {
container_section.push_str(&format!(
"\nContainer: \"{name}\"\nSample documents:\n{sample_json}\n"
));
}
let multi_step_rules = if is_multi {
r#"
MULTI-STEP QUERY RULES — when the user's request involves data from multiple containers:
- Use a `steps:` section in the YAML metadata listing each step with a name and container
- The SQL body uses `-- step: <name>` markers to separate each step's SQL
- Steps that share the same @param inputs run in PARALLEL automatically
- A step can reference another step's result using @step.field syntax (e.g., @customer.id)
This creates a DEPENDENCY — the referenced step runs first, then the value from its first result row is injected
- Do NOT create fan-out queries (one step running per row of another). This is NOT supported.
- Each step's results are available in the template as a top-level array by step name
Multi-step example (parallel — same input):
---
description: Order with line items
database: mydb
params:
- name: orderId
type: string
steps:
- name: header
container: order-headers
- name: lines
container: order-lines
template: |
Order: {{ header[0].orderId }}
Customer: {{ header[0].customerName }}
{% for line in lines %}
{{ line.productName }} {{ line.quantity }} ${{ line.price }}
{% endfor %}
---
-- step: header
SELECT * FROM c WHERE c.orderId = @orderId
-- step: lines
SELECT * FROM c WHERE c.orderId = @orderId ORDER BY c.lineNumber
Multi-step example (chain — step 2 depends on step 1):
---
description: Orders for customer by name
params:
- name: customerName
type: string
steps:
- name: customer
container: customers
- name: orders
container: orders
template: |
Customer: {{ customer[0].name }} ({{ customer[0].id }})
{% for order in orders %}
#{{ order.orderId }} {{ order.date }} ${{ order.total }}
{% endfor %}
---
-- step: customer
SELECT TOP 1 * FROM c WHERE c.name = @customerName
-- step: orders
SELECT * FROM c WHERE c.customerId = @customer.id ORDER BY c.date DESC
"#
} else {
""
};
let single_container_note = if !is_multi {
"\n- Set database and container in the YAML metadata"
} else {
"\n- Set database in the YAML metadata (containers are set per-step)"
};
let template_var_note = if is_multi {
"- For multi-step queries: each step's results are available as a top-level array by step name (e.g., {{ header[0].field }}, {% for line in lines %})"
} else {
"- Templates have access to 'documents' (array of results) and all parameter values"
};
format!(
r#"You are a Cosmos DB SQL query generator. You create .cosq stored query files.
TARGET:
Database: "{database}"
{container_section}
FORMAT — .cosq files use YAML front matter between --- delimiters, followed by the SQL query.
SQL RULES:
- ONLY reference fields that exist in the sample documents above
- Use 'c' as the container alias (e.g., SELECT * FROM c)
- Use Cosmos DB SQL syntax (TOP not LIMIT, no OFFSET, use DateTimeAdd/GetCurrentDateTime for dates)
- Extract variable parts as @param parameters
- Parameters: define in params section with name, type (string/number/bool), description, and optional default/choices/min/max{single_container_note}
{multi_step_rules}
OUTPUT TEMPLATE RULES — ALWAYS include a MiniJinja template in the .cosq file:
- Templates use {{{{ variable }}}} syntax and {{% for doc in documents %}} loops
- {template_var_note}
Choose the template style based on the query:
- MULTIPLE results (lists, searches, TOP N where N > 1): use a readable table layout
Example:
{{% for doc in documents %}}
{{{{ doc.displayName | truncate(25) }}}} {{{{ doc.email }}}} {{{{ doc.status }}}}
{{% endfor %}}
Total: {{{{ documents | length }}}}
- SINGLE result (TOP 1, lookup by ID, "latest", "last"): use a property-value layout
Example:
{{% for doc in documents %}}
Name: {{{{ doc.name }}}}
Email: {{{{ doc.email }}}}
Status: {{{{ doc.status }}}}
{{% endfor %}}
- If the user mentions "JSON" output: do NOT include a template field at all
- If the user mentions "CSV" output: create a template with a header row and comma-separated values
- For complex nested documents: format nested objects and arrays readably (indented sub-properties, bullet lists for arrays)
- When in doubt, pick the format that makes the data most readable
CONVERSATION RULES:
- Be CONFIDENT. If you can make a reasonable assumption, make it and generate the query.
- Only ask clarifying questions if the description is genuinely ambiguous (e.g., which field to filter on, or the user mentions something not in the schema)
- When asking questions, ask 1-3 short questions. Do NOT generate a .cosq file in the same response.
- When generating, respond with ONLY the .cosq file content — no explanation, no markdown fences."#
)
}
fn strip_markdown_fences(response: &str) -> String {
let trimmed = response.trim();
let stripped = trimmed
.strip_prefix("```yaml")
.or_else(|| trimmed.strip_prefix("```cosq"))
.or_else(|| trimmed.strip_prefix("```"))
.unwrap_or(trimmed);
let stripped = stripped.strip_suffix("```").unwrap_or(stripped);
stripped.trim().to_string()
}
fn format_sample_documents(docs: &[serde_json::Value]) -> String {
let truncated: Vec<serde_json::Value> = docs.iter().map(truncate_for_prompt).collect();
if let Ok(json) = serde_json::to_string_pretty(&truncated) {
if json.len() <= 4000 {
return json;
}
}
for n in (1..truncated.len()).rev() {
if let Ok(json) = serde_json::to_string_pretty(&truncated[..n]) {
if json.len() <= 4000 {
return format!("{json}\n(showing {n} of {} sampled documents)", docs.len());
}
}
}
serde_json::to_string_pretty(&truncated[..1])
.unwrap_or_else(|_| "(documents too large to display)".to_string())
}
fn truncate_for_prompt(value: &serde_json::Value) -> serde_json::Value {
use serde_json::Value;
match value {
Value::String(s) if s.len() > 100 => Value::String(format!("{}...", &s[..80])),
Value::Array(arr) if arr.len() > 3 => {
let mut truncated: Vec<Value> = arr[..3].iter().map(truncate_for_prompt).collect();
truncated.push(Value::String(format!("... ({} items total)", arr.len())));
Value::Array(truncated)
}
Value::Array(arr) => Value::Array(arr.iter().map(truncate_for_prompt).collect()),
Value::Object(map) => Value::Object(
map.iter()
.map(|(k, v)| (k.clone(), truncate_for_prompt(v)))
.collect(),
),
other => other.clone(),
}
}
fn generate_filename(description: &str) -> String {
let words: Vec<&str> = description
.split_whitespace()
.filter(|w| {
!matches!(
w.to_lowercase().as_str(),
"find"
| "get"
| "list"
| "show"
| "select"
| "all"
| "the"
| "a"
| "an"
| "from"
| "in"
| "of"
| "by"
| "with"
| "who"
| "that"
| "are"
| "is"
| "and"
| "or"
| "for"
| "to"
)
})
.take(4)
.collect();
let name = words.join("-").to_lowercase();
name.chars()
.filter(|c| c.is_alphanumeric() || *c == '-')
.collect()
}
fn open_in_editor(path: &std::path::Path) -> Result<()> {
let editor = std::env::var("VISUAL")
.or_else(|_| std::env::var("EDITOR"))
.unwrap_or_else(|_| {
if cfg!(target_os = "macos") {
"open".to_string()
} else if cfg!(target_os = "windows") {
"notepad".to_string()
} else {
"xdg-open".to_string()
}
});
eprintln!("{} Opening in {editor}...", ">>".dimmed());
std::process::Command::new(&editor)
.arg(path)
.status()
.with_context(|| format!("failed to open editor: {editor}"))?;
Ok(())
}
fn find_query_path(name: &str) -> Result<std::path::PathBuf> {
let filename = if name.ends_with(".cosq") {
name.to_string()
} else {
format!("{name}.cosq")
};
if let Some(project_dir) = cosq_core::stored_query::project_queries_dir() {
let path = project_dir.join(&filename);
if path.exists() {
return Ok(path);
}
}
let user_dir = user_queries_dir().map_err(|e| anyhow::anyhow!("{e}"))?;
let path = user_dir.join(&filename);
if path.exists() {
return Ok(path);
}
bail!("Query '{name}' not found. Run `cosq queries list` to see available queries.")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_generate_filename() {
assert_eq!(
generate_filename("Find all users who signed up in the last 30 days"),
"users-signed-up-last"
);
assert_eq!(generate_filename("orders over $100"), "orders-over-100");
assert_eq!(
generate_filename("active subscriptions"),
"active-subscriptions"
);
}
#[test]
fn test_generate_filename_empty_input() {
let name = generate_filename("users");
assert_eq!(name, "users");
}
#[test]
fn test_strip_markdown_fences_yaml() {
let input = "```yaml\n---\ndescription: test\n---\nSELECT * FROM c\n```";
assert_eq!(
strip_markdown_fences(input),
"---\ndescription: test\n---\nSELECT * FROM c"
);
}
#[test]
fn test_strip_markdown_fences_plain() {
let input = "```\nsome content\n```";
assert_eq!(strip_markdown_fences(input), "some content");
}
#[test]
fn test_strip_markdown_fences_no_fences() {
let input = "---\ndescription: test\n---\nSELECT * FROM c";
assert_eq!(strip_markdown_fences(input), input);
}
#[test]
fn test_truncate_for_prompt_short_values() {
use serde_json::json;
let value = json!({"name": "short", "count": 42});
let result = truncate_for_prompt(&value);
assert_eq!(result, value);
}
#[test]
fn test_truncate_for_prompt_long_string() {
use serde_json::json;
let long = "x".repeat(200);
let value = json!({"content": long});
let result = truncate_for_prompt(&value);
let content = result["content"].as_str().unwrap();
assert!(content.len() < 200);
assert!(content.ends_with("..."));
}
#[test]
fn test_truncate_for_prompt_large_array() {
use serde_json::json;
let value = json!({"items": [1, 2, 3, 4, 5, 6, 7]});
let result = truncate_for_prompt(&value);
let items = result["items"].as_array().unwrap();
assert_eq!(items.len(), 4); }
#[test]
fn test_format_sample_documents() {
use serde_json::json;
let docs = vec![json!({"id": "1", "name": "test"})];
let formatted = format_sample_documents(&docs);
assert!(formatted.contains("\"id\": \"1\""));
assert!(formatted.contains("\"name\": \"test\""));
}
}