dre_core/
query.rs

1//! A single read-only query on a configured connection, for external integrations.
2
3use dre_protocol::host::{Execution, LogSink};
4use serde_json::{Value, json};
5use sqlparser::ast::{Query, SetExpr, Statement, Visit, Visitor, visit_statements};
6use sqlparser::dialect::{DatabricksDialect, DuckDbDialect, GenericDialect, PostgreSqlDialect};
7use sqlparser::parser::Parser;
8use std::sync::Arc;
9
10use crate::profiles::{Entry, Role};
11use crate::project::{PluginKind, Project};
12
13pub fn execute(project: &Project, profile: &str, sql: &str, max_rows: u64) -> Result<Value, String> {
14    if max_rows == 0 || max_rows > 100_000 {
15        return Err("max_rows must be between 1 and 100000".into());
16    }
17    let target = match project.profiles.entry(Role::Connection, profile) {
18        Entry::Use(target) => target,
19        _ => return Err(format!("connection `{profile}` has no entry for this target")),
20    };
21    let dialect: Box<dyn sqlparser::dialect::Dialect> = match target.kind.as_str() {
22        "duckdb" => Box::new(DuckDbDialect {}),
23        "postgres" => Box::new(PostgreSqlDialect {}),
24        "databricks" => Box::new(DatabricksDialect {}),
25        _ => Box::new(GenericDialect {}),
26    };
27    let statements = Parser::parse_sql(dialect.as_ref(), sql)
28        .map_err(|e| format!("expected a single read-only query: {e}"))?;
29    let mut safe = statements.len() == 1;
30    let _ = visit_statements(&statements, |statement| {
31        if !matches!(statement, Statement::Query(_)) {
32            safe = false;
33        }
34        std::ops::ControlFlow::<()>::Continue(())
35    });
36    struct ReadOnly;
37    impl Visitor for ReadOnly {
38        type Break = ();
39        fn pre_visit_query(&mut self, query: &Query) -> std::ops::ControlFlow<()> {
40            fn reads(body: &SetExpr) -> bool {
41                match body {
42                    SetExpr::Select(select) => select.into.is_none(),
43                    SetExpr::SetOperation { left, right, .. } => reads(left) && reads(right),
44                    SetExpr::Query(query) => reads(&query.body),
45                    SetExpr::Values(_) | SetExpr::Table(_) => true,
46                    _ => false,
47                }
48            }
49            if !reads(&query.body) || !query.locks.is_empty() || !query.pipe_operators.is_empty() {
50                std::ops::ControlFlow::Break(())
51            } else {
52                std::ops::ControlFlow::Continue(())
53            }
54        }
55    }
56    safe &= statements.visit(&mut ReadOnly).is_continue();
57    if !safe {
58        return Err("expected a single read-only query; nothing was run".into());
59    }
60    let plugin =
61        crate::run::find_plugin(project, PluginKind::Source, &target.kind).map_err(|e| e.to_string())?;
62    let log: LogSink =
63        Arc::new(|name, message| eprintln!("{name}: {}", crate::secrets::mask_source_error(message, false)));
64    let mut process = plugin
65        .start(log, Some(&project.root))
66        .map_err(|e| e.to_string())?;
67    let session_read_only = process.has(dre_protocol::CAP_READ_ONLY);
68    // Databricks has no read-only session capability. Its single-SELECT guard is
69    // enforced here; database permissions remain the user's enforcement boundary.
70    if !session_read_only && target.kind != "databricks" {
71        return Err(format!(
72            "the {} plugin cannot guarantee a read-only connection",
73            target.kind
74        ));
75    }
76    process
77        .open(crate::run::render_connection(target)?, session_read_only)
78        .map_err(|e| crate::secrets::mask_source_error(&e.to_string(), false).into_owned())?;
79    let mut rows = Vec::new();
80    let mut columns = Vec::new();
81    let result = process
82        .execute(sql, Some(max_rows + 1), |schema, batch| {
83            if columns.is_empty() {
84                columns = schema
85                    .fields()
86                    .iter()
87                    .map(|f| json!({"name": f.name(), "type": f.data_type().to_string()}))
88                    .collect();
89            }
90            let left = (max_rows + 1).saturating_sub(rows.len() as u64) as usize;
91            let batch = batch.slice(0, left.min(batch.num_rows()));
92            rows.extend(
93                crate::values::batch_rows(&batch)
94                    .into_iter()
95                    .map(|row| serde_json::to_value(row).unwrap()),
96            );
97            Ok(())
98        })
99        .map_err(|e| crate::secrets::mask_source_error(&e.to_string(), false).into_owned())?;
100    if let Execution::Result { schema, .. } = result {
101        if columns.is_empty() {
102            columns = schema
103                .fields()
104                .iter()
105                .map(|f| json!({"name": f.name(), "type": f.data_type().to_string()}))
106                .collect();
107        }
108    } else {
109        return Err("the query returned no result set".into());
110    }
111    let truncated = rows.len() as u64 > max_rows;
112    rows.truncate(max_rows as usize);
113    let _ = process.close();
114    Ok(
115        json!({"schema_version": "dre/rows/v1", "columns": columns, "row_count": rows.len(), "rows": rows,
116        "truncated": truncated, "read_only_enforcement": if session_read_only {"session"} else {"SQL validation; use a read-only Databricks database role"}, "connection": profile, "target": project.target_name}),
117    )
118}
dre_core/query.rs

dre_core/
query.rs
