Run Data pre-query scripts before dispatch

crates/core/src/data.rs

@@ -29,6 +29,7 @@ pub enum DataFileFormat {
 pub struct DataConfig {
     pub source: DataSource,
     pub sql: String,
+    pub query_pre_script: String,
     pub file_name: String,
     pub file_format: DataFileFormat,
     pub sheet_name: String,
@@ -44,6 +45,7 @@ impl Default for DataConfig {
         Self {
             source: DataSource::Postgresql,
             sql: "SELECT 1 AS value".into(),
+            query_pre_script: String::new(),
             file_name: String::new(),
             file_format: DataFileFormat::Csv,
             sheet_name: String::new(),
@@ -57,6 +59,7 @@ impl Default for DataConfig {
     }
 }
 pub fn validate_data_config(config: &DataConfig) -> Result<()> {
+    crate::validate_script(&config.query_pre_script)?;
     ensure!(
         config.sql.len() <= MAX_SQL_SOURCE && !config.sql.contains('\0'),
         "Data SQL source exceeds limits"
@@ -221,6 +224,20 @@ pub fn data_url(raw: &str, source: DataSource) -> Result<Url> {
 mod tests {
     use super::*;
     #[test]
+    fn query_pre_script_defaults_and_obeys_shared_script_limit() {
+        let legacy: DataConfig =
+            serde_json::from_value(serde_json::json!({"sql":"SELECT 1"})).unwrap();
+        assert!(legacy.query_pre_script.is_empty());
+        let mut configured = legacy;
+        configured.query_pre_script = "x".repeat(crate::MAX_SCRIPT_BYTES);
+        validate_data_config(&configured).unwrap();
+        let copied: DataConfig =
+            serde_json::from_value(serde_json::to_value(&configured).unwrap()).unwrap();
+        assert_eq!(copied.query_pre_script, configured.query_pre_script);
+        configured.query_pre_script.push('x');
+        assert!(validate_data_config(&configured).is_err());
+    }
+    #[test]
     fn retains_selected_file_and_sql_but_rejects_hidden_network_options() {
         let config = DataConfig {
             source: DataSource::LocalFile,

crates/core/src/interpolation.rs

@@ -118,6 +118,9 @@ pub fn resolve_request(
     // JavaScript is source code, not an interpolated request field.
     value["pre_request_script"] = "".into();
     value["post_response_script"] = "".into();
+    if request.protocol.is_data() {
+        value["protocol"]["query_pre_script"] = "".into();
+    }
     let mqtt_will = if let Protocol::Mqtt { config } = &request.protocol {
         config.will.as_ref().map(|w| w.message.clone())
     } else {
@@ -499,6 +502,11 @@ pub fn resolve_request(
         config.message.payload_source = original.message.payload_source.clone();
         resolved.body = request.body.clone();
     }
+    if let (Protocol::Data { config }, Protocol::Data { config: original }) =
+        (&mut resolved.protocol, &request.protocol)
+    {
+        config.query_pre_script = original.query_pre_script.clone();
+    }
     resolved.pre_request_script = request.pre_request_script.clone();
     resolved.post_response_script = request.post_response_script.clone();
     reconcile_graphql(&mut resolved)?;

crates/formats/src/redact.rs

@@ -522,6 +522,8 @@ pub(super) fn workspace(source: &Workspace) -> Workspace {
                     request.url = data_privacy.screen_text(&redact_url(&request.url));
                     config.file_base64.clear();
                     config.sql = data_privacy.screen_generation_text(&config.sql);
+                    config.query_pre_script = data_privacy
+                        .screen_bounded(&config.query_pre_script, moleapi_core::MAX_SCRIPT_BYTES);
                     config.file_name = data_privacy.screen_text(&config.file_name);
                     config.table_name = data_privacy.screen_text(&config.table_name);
                 }

crates/formats/tests/interchange.rs

@@ -761,7 +761,7 @@ fn data_request_exports_screen_credentials_copied_sql_and_keep_explicit_original
         sql: "SELECT 'it''s-private' AS copied".into(),
         ..Default::default()
     };
-    let request:moleapi_core::RequestSpec=serde_json::from_value(json!({"id":"data","name":"Data request","protocol":{"kind":"data","source":"postgresql","sql":config.sql},"method":"GET","url":"postgresql://user:it%27s-private@db.example/main","description":"","query":[],"headers":[],"body_kind":"none","body":"","auth":{"kind":"none","token":"","username":"","password":""},"timeout_ms":1000,"verify_tls":true,"follow_redirects":false,"assertions":[],"examples":[]})).unwrap();
+    let request:moleapi_core::RequestSpec=serde_json::from_value(json!({"id":"data","name":"Data request","protocol":{"kind":"data","source":"postgresql","sql":config.sql,"query_pre_script":"pm.request.body.update(\"SELECT 'it''s-private'\")"},"method":"GET","url":"postgresql://user:it%27s-private@db.example/main","description":"","query":[],"headers":[],"body_kind":"none","body":"","auth":{"kind":"none","token":"","username":"","password":""},"timeout_ms":1000,"verify_tls":true,"follow_redirects":false,"assertions":[],"examples":[]})).unwrap();
     let mut w=workspace(serde_json::from_value(json!({"schema_version":1,"collections":[{"id":"c","name":"Data","description":"","requests":[request]}],"environments":[],"active_environment_id":null})).unwrap());
     let format = "moleapi";
     {

crates/protocols/src/data.rs

@@ -1,8 +1,9 @@
 use crate::{Command, EventMessage, SendMessage, Session, SessionState};
 use anyhow::{Context, Result, ensure};
-use moleapi_core::{Environment, NetworkPolicy, Protocol, RequestSpec};
+use moleapi_core::{DataSource, Environment, NetworkPolicy, Protocol, RequestSpec};
 use moleapi_data::{Connection, QueryResult};
 use std::{
+    collections::HashSet,
     path::PathBuf,
     sync::{Arc, OnceLock},
     time::{Duration, Instant},
@@ -15,6 +16,7 @@ pub(crate) struct Control {
     pub worker: Option<PathBuf>,
     active: Option<(String, CancellationToken)>,
     commands: usize,
+    used_query_ids: HashSet<String>,
     pub completions: Option<mpsc::Sender<DataCompletion>>,
 }
 pub struct DataCompletion {
@@ -29,6 +31,119 @@ pub enum DataOutcome {
     Error(String),
     Cancelled { write_outcome_unknown: bool },
 }
+pub struct DataReservation {
+    session: Arc<Session>,
+    query_id: String,
+    stop: CancellationToken,
+    dispatched: bool,
+}
+impl DataReservation {
+    pub fn cancelled_now(&self) -> bool {
+        self.stop.is_cancelled() || self.session.cancel.is_cancelled()
+    }
+    pub async fn cancelled(&self) {
+        tokio::select! { _ = self.stop.cancelled() => {}, _ = self.session.cancel.cancelled() => {} }
+    }
+    pub fn dispatch(mut self, sql: String, read_only: bool) -> Result<()> {
+        let mut record = self.session.record.lock().unwrap();
+        ensure!(
+            record.summary.state == SessionState::Open
+                && !self.session.cancel.is_cancelled()
+                && !self.stop.is_cancelled(),
+            "Data query was cancelled before dispatch"
+        );
+        ensure!(
+            !sql.trim().is_empty() && sql.len() <= moleapi_core::MAX_SQL_SOURCE,
+            "Data SQL query exceeds limits"
+        );
+        let mut control = self.session.data.lock().unwrap();
+        ensure!(
+            control
+                .active
+                .as_ref()
+                .is_some_and(|(id, _)| id == &self.query_id),
+            "Data query reservation expired"
+        );
+        ensure!(control.commands < 512, "Data command budget exceeded");
+        ensure!(
+            !control.used_query_ids.contains(&self.query_id),
+            "Data query ID was already dispatched"
+        );
+        ensure!(
+            record.input.saturating_add(sql.len()) <= 20 * 1024 * 1024,
+            "Data command input budget exceeded"
+        );
+        let size = sql.len();
+        self.session
+            .commands
+            .try_send(Command::DataQuery {
+                query_id: self.query_id.clone(),
+                sql,
+                read_only,
+                stop: self.stop.clone(),
+            })
+            .context("Data command queue is full")?;
+        record.input += size;
+        control.commands += 1;
+        control.used_query_ids.insert(self.query_id.clone());
+        self.dispatched = true;
+        Ok(())
+    }
+}
+impl Drop for DataReservation {
+    fn drop(&mut self) {
+        if !self.dispatched {
+            let mut control = self.session.data.lock().unwrap();
+            if control
+                .active
+                .as_ref()
+                .is_some_and(|(id, _)| id == &self.query_id)
+            {
+                control.active = None;
+            }
+        }
+    }
+}
+pub(crate) fn reserve(
+    session: &Arc<Session>,
+    query_id: String,
+    sql: &str,
+) -> Result<DataReservation> {
+    let record = session.record.lock().unwrap();
+    ensure!(
+        record.summary.protocol == "data"
+            && record.summary.state == SessionState::Open
+            && !session.cancel.is_cancelled(),
+        "Data connection is not open"
+    );
+    ensure!(
+        !query_id.is_empty() && query_id.len() <= 128 && !query_id.chars().any(char::is_control),
+        "Invalid Data query ID"
+    );
+    ensure!(
+        !sql.trim().is_empty() && sql.len() <= moleapi_core::MAX_SQL_SOURCE,
+        "Data SQL query exceeds limits"
+    );
+    let mut control = session.data.lock().unwrap();
+    ensure!(control.active.is_none(), "A Data query is already running");
+    ensure!(
+        !control.used_query_ids.contains(&query_id),
+        "Data query ID was already dispatched"
+    );
+    ensure!(control.commands < 512, "Data command budget exceeded");
+    ensure!(
+        record.input.saturating_add(sql.len()) <= 20 * 1024 * 1024,
+        "Data command input budget exceeded"
+    );
+    let stop = CancellationToken::new();
+    control.active = Some((query_id.clone(), stop.clone()));
+    Ok(DataReservation {
+        session: session.clone(),
+        query_id,
+        stop,
+        dispatched: false,
+    })
+}
 fn complete(
     session: &Session,
     query_id: String,
@@ -61,7 +176,15 @@ fn slots() -> Arc<Semaphore> {
     SLOTS.get_or_init(|| Arc::new(Semaphore::new(4))).clone()
 }
 pub(crate) fn send(session: &Arc<Session>, message: SendMessage) -> Result<()> {
-    let mut record = session.record.lock().unwrap();
+    if let SendMessage::DataQuery {
+        query_id,
+        sql,
+        read_only,
+    } = &message
+    {
+        return reserve(session, query_id.clone(), sql)?.dispatch(sql.clone(), *read_only);
+    }
+    let record = session.record.lock().unwrap();
     ensure!(
         record.summary.protocol == "data"
             && record.summary.state == SessionState::Open
@@ -77,43 +200,6 @@ pub(crate) fn send(session: &Arc<Session>, message: SendMessage) -> Result<()> {
             }
             Ok(())
         }
-        SendMessage::DataQuery {
-            query_id,
-            sql,
-            read_only,
-        } => {
-            ensure!(
-                !query_id.is_empty()
-                    && query_id.len() <= 128
-                    && !query_id.chars().any(char::is_control),
-                "Invalid Data query ID"
-            );
-            ensure!(
-                !sql.trim().is_empty() && sql.len() <= moleapi_core::MAX_SQL_SOURCE,
-                "Data SQL query exceeds limits"
-            );
-            ensure!(control.active.is_none(), "A Data query is already running");
-            ensure!(control.commands < 512, "Data command budget exceeded");
-            ensure!(
-                record.input.saturating_add(sql.len()) <= 20 * 1024 * 1024,
-                "Data command input budget exceeded"
-            );
-            let input_bytes = sql.len();
-            let stop = CancellationToken::new();
-            session
-                .commands
-                .try_send(Command::DataQuery {
-                    query_id: query_id.clone(),
-                    sql,
-                    read_only,
-                    stop: stop.clone(),
-                })
-                .context("Data command queue is full")?;
-            record.input += input_bytes;
-            control.active = Some((query_id, stop));
-            control.commands += 1;
-            Ok(())
-        }
         SendMessage::DataSchemaRefresh => {
             ensure!(control.active.is_none(), "A Data query is already running");
             ensure!(control.commands < 512, "Data command budget exceeded");
@@ -297,6 +383,12 @@ pub(crate) async fn run(
                 read_only,
                 stop,
             } => {
+                let read_only = read_only
+                    || config.read_only
+                    || matches!(
+                        config.source,
+                        DataSource::LocalFile | DataSource::RemoteFile
+                    );
                 let started = Instant::now();
                 let permit = match slots().try_acquire_owned() {
                     Ok(value) => value,
@@ -448,3 +540,56 @@ pub(crate) async fn run(
         }
     }
 }
+
+#[cfg(test)]
+mod reservation_tests {
+    use super::*;
+    use crate::{PreparedFeedback, SessionManager};
+    use serde_json::json;
+
+    #[tokio::test]
+    async fn cancelled_reservation_releases_capacity_without_dispatching_sql() {
+        let request:RequestSpec = serde_json::from_value(json!({"id":"data","name":"Data","method":"GET","url":"","description":"","query":[],"headers":[],"body_kind":"none","body":"","auth":{"kind":"none","token":"","username":"","password":""},"timeout_ms":1000,"verify_tls":true,"follow_redirects":false,"assertions":[],"examples":[],"protocol":{"kind":"data","source":"local_file"}})).unwrap();
+        let manager = SessionManager::default();
+        let summary = manager
+            .register(
+                "owner",
+                "workspace",
+                &request,
+                "safe".into(),
+                PreparedFeedback::default(),
+            )
+            .unwrap();
+        let session = manager.owned("owner", &summary.id).unwrap();
+        session.record.lock().unwrap().summary.state = SessionState::Open;
+        let reserved = manager
+            .reserve_data_query("owner", &summary.id, "first".into(), "SELECT 1")
+            .unwrap();
+        assert!(
+            manager
+                .reserve_data_query("owner", &summary.id, "second".into(), "SELECT 2")
+                .is_err()
+        );
+        send(
+            &session,
+            SendMessage::DataCancel {
+                query_id: "first".into(),
+            },
+        )
+        .unwrap();
+        tokio::time::timeout(Duration::from_millis(100), reserved.cancelled())
+            .await
+            .unwrap();
+        assert!(reserved.dispatch("SELECT 1".into(), true).is_err());
+        let next = manager
+            .reserve_data_query("owner", &summary.id, "second".into(), "SELECT 2")
+            .unwrap();
+        next.dispatch("SELECT 3".into(), true).unwrap();
+        let mut receiver = session.receiver.lock().unwrap().take().unwrap();
+        let command = receiver.try_recv().unwrap();
+        assert!(
+            matches!(command, Command::DataQuery { query_id, sql, .. } if query_id == "second" && sql == "SELECT 3")
+        );
+        assert!(receiver.try_recv().is_err());
+    }
+}

crates/protocols/src/lib.rs

@@ -7,7 +7,7 @@ mod mqtt;
 mod tcp;
 pub use grpc::{ReflectionResult, ReflectionStatus, reflect};
 mod data;
-pub use data::{DataCompletion, DataOutcome};
+pub use data::{DataCompletion, DataOutcome, DataReservation};
 mod engine;
 mod graphql;
 mod models;
@@ -86,6 +86,10 @@ impl Session {
             EventMessage::DataReady { .. }
             | EventMessage::DataSchema { .. }
             | EventMessage::DataStarted { .. }
+            | EventMessage::DataPreparing { .. }
+            | EventMessage::DataPrepared { .. }
+            | EventMessage::DataPreprocessError { .. }
+            | EventMessage::DataSkipped { .. }
             | EventMessage::DataColumns { .. }
             | EventMessage::DataRows { .. }
             | EventMessage::DataFinished { .. }
@@ -524,6 +528,15 @@ impl SessionManager {
         );
         session.event("system", message)
     }
+    pub fn reserve_data_query(
+        &self,
+        owner: &str,
+        id: &str,
+        query_id: String,
+        sql: &str,
+    ) -> Result<DataReservation> {
+        data::reserve(&self.owned(owner, id)?, query_id, sql)
+    }
     pub fn configure_tcp(
         &self,
         owner: &str,