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::{
6    Expr, ObjectName, Query, SetExpr, Statement, TableFactor, Visit, Visitor, visit_statements,
7};
8use sqlparser::dialect::{DatabricksDialect, DuckDbDialect, GenericDialect, PostgreSqlDialect};
9use sqlparser::parser::Parser;
10use std::sync::Arc;
11
12use crate::profiles::{Entry, Role};
13use crate::project::{PluginKind, Project};
14
15pub fn execute(project: &Project, profile: &str, sql: &str, max_rows: u64) -> Result<Value, String> {
16    if max_rows == 0 || max_rows > 100_000 {
17        return Err("max_rows must be between 1 and 100000".into());
18    }
19    let target = match project.profiles.entry(Role::Connection, profile) {
20        Entry::Use(target) => target,
21        _ => return Err(format!("connection `{profile}` has no entry for this target")),
22    };
23    let dialect: Box<dyn sqlparser::dialect::Dialect> = match target.kind.as_str() {
24        "duckdb" => Box::new(DuckDbDialect {}),
25        "postgres" => Box::new(PostgreSqlDialect {}),
26        "databricks" => Box::new(DatabricksDialect {}),
27        _ => Box::new(GenericDialect {}),
28    };
29    let statements = Parser::parse_sql(dialect.as_ref(), sql)
30        .map_err(|e| format!("expected a single read-only query: {e}"))?;
31    let mut safe = statements.len() == 1;
32    let _ = visit_statements(&statements, |statement| {
33        if !matches!(statement, Statement::Query(_)) {
34            safe = false;
35        }
36        std::ops::ControlFlow::<()>::Continue(())
37    });
38    /// Refuses anything but reads, and functions that reach files, programs or the network.
39    struct ReadOnly {
40        dialect: Option<&'static crate::dialect::Dialect>,
41        refused: Option<String>,
42    }
43    impl ReadOnly {
44        fn function(&mut self, name: &ObjectName) -> std::ops::ControlFlow<()> {
45            let name = name.to_string();
46            if self.dialect.is_some_and(|d| d.denies(&name)) {
47                self.refused = Some(name);
48                return std::ops::ControlFlow::Break(());
49            }
50            std::ops::ControlFlow::Continue(())
51        }
52    }
53    impl Visitor for ReadOnly {
54        type Break = ();
55        fn pre_visit_query(&mut self, query: &Query) -> std::ops::ControlFlow<()> {
56            fn reads(body: &SetExpr) -> bool {
57                match body {
58                    SetExpr::Select(select) => select.into.is_none(),
59                    SetExpr::SetOperation { left, right, .. } => reads(left) && reads(right),
60                    SetExpr::Query(query) => reads(&query.body),
61                    SetExpr::Values(_) | SetExpr::Table(_) => true,
62                    _ => false,
63                }
64            }
65            if !reads(&query.body) || !query.locks.is_empty() || !query.pipe_operators.is_empty() {
66                std::ops::ControlFlow::Break(())
67            } else {
68                std::ops::ControlFlow::Continue(())
69            }
70        }
71        fn pre_visit_table_factor(&mut self, factor: &TableFactor) -> std::ops::ControlFlow<()> {
72            match factor {
73                // `FROM 'file.csv'`: a path, not a table.
74                TableFactor::Table { name, .. }
75                    if name
76                        .0
77                        .iter()
78                        .any(|p| p.as_ident().is_some_and(|i| i.quote_style == Some('\''))) =>
79                {
80                    self.refused = Some(name.to_string());
81                    std::ops::ControlFlow::Break(())
82                }
83                // A table function, not a table that happens to share a function's name.
84                TableFactor::Table {
85                    name, args: Some(_), ..
86                }
87                | TableFactor::Function { name, .. } => self.function(name),
88                _ => std::ops::ControlFlow::Continue(()),
89            }
90        }
91        fn pre_visit_expr(&mut self, expr: &Expr) -> std::ops::ControlFlow<()> {
92            match expr {
93                Expr::Function(f) => self.function(&f.name),
94                _ => std::ops::ControlFlow::Continue(()),
95            }
96        }
97    }
98    let mut guard = ReadOnly {
99        dialect: crate::dialect::of(&target.kind),
100        refused: None,
101    };
102    safe &= statements.visit(&mut guard).is_continue();
103    if let Some(name) = guard.refused {
104        return Err(format!(
105            "`{name}` reads files, runs programs or reaches the network; `dre query` refuses it and nothing was run"
106        ));
107    }
108    if !safe {
109        return Err("expected a single read-only query; nothing was run".into());
110    }
111    let plugin =
112        crate::run::find_plugin(project, PluginKind::Source, &target.kind).map_err(|e| e.to_string())?;
113    let log: LogSink =
114        Arc::new(|name, message| eprintln!("{name}: {}", crate::secrets::mask_source_error(message, false)));
115    let mut process = plugin
116        .start(log, Some(&project.root))
117        .map_err(|e| e.to_string())?;
118    let session_read_only = process.has(dre_protocol::CAP_READ_ONLY);
119    // Databricks has no read-only session capability. Its single-SELECT guard is
120    // enforced here; database permissions remain the user's enforcement boundary.
121    if !session_read_only && target.kind != "databricks" {
122        return Err(format!(
123            "the {} plugin cannot guarantee a read-only connection",
124            target.kind
125        ));
126    }
127    let connection = crate::run::render_connection(target)?;
128    if session_read_only {
129        process.open_sandboxed(connection)
130    } else {
131        process.open(connection, false)
132    }
133    .map_err(|e| crate::secrets::mask_source_error(&e.to_string(), false).into_owned())?;
134    let mut rows = Vec::new();
135    let mut columns = Vec::new();
136    let result = process
137        .execute(sql, Some(max_rows + 1), |schema, batch| {
138            if columns.is_empty() {
139                columns = crate::rows::column_types(schema);
140            }
141            let left = (max_rows + 1).saturating_sub(rows.len() as u64) as usize;
142            let batch = batch.slice(0, left.min(batch.num_rows()));
143            rows.extend(
144                crate::values::batch_rows(&batch)
145                    .into_iter()
146                    .map(|row| serde_json::to_value(row).unwrap()),
147            );
148            Ok(())
149        })
150        .map_err(|e| crate::secrets::mask_source_error(&e.to_string(), false).into_owned())?;
151    if let Execution::Result { schema, .. } = result {
152        if columns.is_empty() {
153            columns = crate::rows::column_types(&schema);
154        }
155    } else {
156        return Err("the query returned no result set".into());
157    }
158    let truncated = rows.len() as u64 > max_rows;
159    rows.truncate(max_rows as usize);
160    let _ = process.close();
161    Ok(
162        json!({"schema_version": "dre/rows/v1", "columns": columns, "row_count": rows.len(), "rows": rows,
163        "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}),
164    )
165}
dre_core/query.rs

dre_core/
query.rs
