add support for variables

lovasoa · lovasoa · commit 4c00694e9902 · 2023-09-19T02:07:33.000+02:00
see #49 #69
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,15 @@
 # CHANGELOG.md
 
+## 0.12.0 (unreleased)
+
+ - **variables** . SQLPage now support setting and reusing variables between statements. This allows you to write more complex SQL queries, and to reuse the result of a query in multiple places.
+   ```sql
+   -- Set a variable
+   SET $person = 'Alice';
+   -- Use it in a query
+   SELECT 'text' AS component, 'Hello ' || $person AS contents;
+   ```
+
 ## 0.11.0 (2023-09-17)
  - Support for **environment variables** ! You can now read environment variables from sql code using `sqlpage.environment_variable('VAR_NAME')`.
  - Better support for connection options in mssql.
diff --git a/src/webserver/database/mod.rs b/src/webserver/database/mod.rs
@@ -7,14 +7,15 @@ use futures_util::stream::Stream;
 use futures_util::StreamExt;
 use serde_json::Value;
 use std::borrow::Cow;
+use std::collections::HashMap;
 use std::fmt::{Display, Formatter};
 use std::time::Duration;
 
 use crate::app_config::AppConfig;
 pub use crate::file_cache::FileCache;
 
 use crate::webserver::database::sql_pseudofunctions::extract_req_param;
-use crate::webserver::http::RequestInfo;
+use crate::webserver::http::{RequestInfo, SingleOrVec};
 use crate::MIGRATIONS_DIR;
 pub use sql::make_placeholder;
 pub use sql::ParsedSqlFile;
@@ -24,9 +25,12 @@ use sqlx::any::{
 use sqlx::migrate::Migrator;
 use sqlx::pool::{PoolConnection, PoolOptions};
 use sqlx::query::Query;
-use sqlx::{Any, AnyConnection, AnyPool, Arguments, ConnectOptions, Either, Executor, Statement};
+use sqlx::{
+    Any, AnyConnection, AnyPool, Arguments, ConnectOptions, Either, Executor, Row, Statement,
+};
 
 use self::sql::ParsedSQLStatement;
+use sql_pseudofunctions::StmtParam;
 
 pub struct Database {
     pub(crate) connection: AnyPool,
@@ -97,27 +101,15 @@ fn migration_err(operation: &'static str) -> String {
 pub fn stream_query_results<'a>(
     db: &'a Database,
     sql_file: &'a ParsedSqlFile,
-    request: &'a RequestInfo,
+    request: &'a mut RequestInfo,
 ) -> impl Stream<Item = DbItem> + 'a {
-    async_stream::stream! {
+    async_stream::try_stream! {
         let mut connection_opt = None;
         for res in &sql_file.statements {
             match res {
-                ParsedSQLStatement::Statement(stmt)=>{
-                    let query = match bind_parameters(stmt, request) {
-                        Ok(q) => q,
-                        Err(e) => {
-                            yield DbItem::Error(e);
-                            continue;
-                        }
-                    };
-                    let connection = match take_connection(db, &mut connection_opt).await {
-                        Ok(c) => c,
-                        Err(e) => {
-                            yield DbItem::Error(e);
-                            return;
-                        }
-                    };
+                ParsedSQLStatement::Statement(stmt) => {
+                    let query = bind_parameters(stmt, request)?;
+                    let connection = take_connection(db, &mut connection_opt).await?;
                     let mut stream = query.fetch_many(connection);
                     while let Some(elem) = stream.next().await {
                         let is_err = elem.is_err();
@@ -127,13 +119,49 @@ pub fn stream_query_results<'a>(
                         }
                     }
                 },
+                ParsedSQLStatement::SetVariable { variable, value} => {
+                    let query = bind_parameters(value, request)?;
+                    let connection = take_connection(db, &mut connection_opt).await?;
+                    let row = query.fetch_optional(connection).await?;
+                    let (vars, name) = vars_and_name(request, variable)?;
+                    if let Some(row) = row {
+                        vars.insert(name.clone(), row_to_varvalue(&row));
+                    } else {
+                        vars.remove(&name);
+                    }
+                },
                 ParsedSQLStatement::StaticSimpleSelect(value) => {
                     yield DbItem::Row(value.clone().into())
                 }
                 ParsedSQLStatement::Error(e) => yield DbItem::Error(clone_anyhow_err(e)),
             }
         }
     }
+    .map(|res| res.unwrap_or_else(DbItem::Error))
+}
+
+fn vars_and_name<'a>(
+    request: &'a mut RequestInfo,
+    variable: &StmtParam,
+) -> anyhow::Result<(&'a mut HashMap<String, SingleOrVec>, String)> {
+    match variable {
+        StmtParam::Get(name) => {
+            let vars = &mut request.get_variables;
+            Ok((vars, name.clone()))
+        }
+        StmtParam::Post(name) => {
+            let vars = &mut request.post_variables;
+            Ok((vars, name.clone()))
+        }
+        _ => Err(anyhow!(
+            "Only GET and POST variables can be set, not {variable:?}"
+        )),
+    }
+}
+
+fn row_to_varvalue(row: &AnyRow) -> SingleOrVec {
+    row.try_get::<String, usize>(0)
+        .map_or_else(|_| SingleOrVec::Vec(vec![]), SingleOrVec::Single)
 }
 
 async fn take_connection<'a, 'b>(
@@ -293,7 +321,7 @@ fn set_custom_connect_options(options: &mut AnyConnectOptions, config: &AppConfi
 }
 struct PreparedStatement {
     statement: AnyStatement<'static>,
-    parameters: Vec<sql_pseudofunctions::StmtParam>,
+    parameters: Vec<StmtParam>,
 }
 
 impl Display for PreparedStatement {
diff --git a/src/webserver/database/sql.rs b/src/webserver/database/sql.rs
@@ -26,6 +26,10 @@ pub(super) enum ParsedSQLStatement {
     Statement(PreparedStatement),
     StaticSimpleSelect(serde_json::Map<String, serde_json::Value>),
     Error(anyhow::Error),
+    SetVariable {
+        variable: StmtParam,
+        value: PreparedStatement,
+    },
 }
 
 impl ParsedSqlFile {
@@ -40,20 +44,15 @@ impl ParsedSqlFile {
             statements.push(match parsed {
                 ParsedStatement::StaticSimpleSelect(s) => ParsedSQLStatement::StaticSimpleSelect(s),
                 ParsedStatement::Error(e) => ParsedSQLStatement::Error(e),
-                ParsedStatement::StmtWithParams { query, params } => {
-                    let param_types = get_param_types(&params);
-                    match db.prepare_with(&query, &param_types).await {
-                        Ok(statement) => {
-                            log::debug!("Successfully prepared SQL statement '{query}'");
-                            ParsedSQLStatement::Statement(PreparedStatement {
-                                statement,
-                                parameters: params,
-                            })
-                        }
-                        Err(err) => {
-                            log::warn!("Failed to prepare {query:?}: {err:#}");
-                            ParsedSQLStatement::Error(err)
+                ParsedStatement::StmtWithParams(stmt_with_params) => {
+                    prepare_query_with_params(db, stmt_with_params).await
+                }
+                ParsedStatement::SetVariable { variable, value } => {
+                    match prepare_query_with_params(db, value).await {
+                        ParsedSQLStatement::Statement(value) => {
+                            ParsedSQLStatement::SetVariable { variable, value }
                         }
+                        err => err,
                     }
                 }
             });
@@ -71,19 +70,45 @@ impl ParsedSqlFile {
     }
 }
 
+async fn prepare_query_with_params(
+    db: &Database,
+    StmtWithParams { query, params }: StmtWithParams,
+) -> ParsedSQLStatement {
+    let param_types = get_param_types(&params);
+    match db.prepare_with(&query, &param_types).await {
+        Ok(statement) => {
+            log::debug!("Successfully prepared SQL statement '{query}'");
+            ParsedSQLStatement::Statement(PreparedStatement {
+                statement,
+                parameters: params,
+            })
+        }
+        Err(err) => {
+            log::warn!("Failed to prepare {query:?}: {err:#}");
+            ParsedSQLStatement::Error(err)
+        }
+    }
+}
+
 #[async_trait(? Send)]
 impl AsyncFromStrWithState for ParsedSqlFile {
     async fn from_str_with_state(app_state: &AppState, source: &str) -> anyhow::Result<Self> {
         Ok(ParsedSqlFile::new(&app_state.db, source).await)
     }
 }
 
+struct StmtWithParams {
+    query: String,
+    params: Vec<StmtParam>,
+}
+
 enum ParsedStatement {
-    StmtWithParams {
-        query: String,
-        params: Vec<StmtParam>,
-    },
+    StmtWithParams(StmtWithParams),
     StaticSimpleSelect(serde_json::Map<String, serde_json::Value>),
+    SetVariable {
+        variable: StmtParam,
+        value: StmtWithParams,
+    },
     Error(anyhow::Error),
 }
 
@@ -113,8 +138,16 @@ fn parse_single_statement(parser: &mut Parser<'_>, db_kind: AnyKind) -> Option<P
         return Some(ParsedStatement::StaticSimpleSelect(static_statement));
     }
     let params = ParameterExtractor::extract_parameters(&mut stmt, db_kind);
-    let query = stmt.to_string();
-    Some(ParsedStatement::StmtWithParams { query, params })
+    if let Some((variable, query)) = extract_set_variable(&mut stmt) {
+        return Some(ParsedStatement::SetVariable {
+            variable,
+            value: StmtWithParams { query, params },
+        });
+    }
+    Some(ParsedStatement::StmtWithParams(StmtWithParams {
+        query: stmt.to_string(),
+        params,
+    }))
 }
 
 fn syntax_error(err: ParserError, parser: &mut Parser) -> ParsedStatement {
@@ -226,6 +259,24 @@ fn extract_static_simple_select(
     Some(map)
 }
 
+fn extract_set_variable(stmt: &mut Statement) -> Option<(StmtParam, String)> {
+    if let Statement::SetVariable {
+        variable: ObjectName(name),
+        value,
+        local: false,
+        hivevar: false,
+    } = stmt
+    {
+        if let ([ident], [value]) = (name.as_mut_slice(), value.as_mut_slice()) {
+            if let Some(variable) = extract_ident_param(ident) {
+                let query = format!("SELECT {value}");
+                return Some((variable, query));
+            }
+        }
+    }
+    None
+}
+
 struct ParameterExtractor {
     db_kind: AnyKind,
     parameters: Vec<StmtParam>,
@@ -378,17 +429,24 @@ pub fn make_placeholder(db_kind: AnyKind, arg_number: usize) -> String {
     DEFAULT_PLACEHOLDER.to_string()
 }
 
+fn extract_ident_param(Ident { value, quote_style }: &mut Ident) -> Option<StmtParam> {
+    if quote_style.is_none() && value.starts_with('$') || value.starts_with(':') {
+        let name = std::mem::take(value);
+        Some(map_param(name))
+    } else {
+        None
+    }
+}
+
 impl VisitorMut for ParameterExtractor {
     type Break = ();
     fn pre_visit_expr(&mut self, value: &mut Expr) -> ControlFlow<Self::Break> {
         match value {
-            Expr::Identifier(Ident {
-                value: var_name,
-                quote_style: None,
-            }) if var_name.starts_with('$') || var_name.starts_with(':') => {
-                let name = std::mem::take(var_name);
-                *value = self.make_placeholder();
-                self.parameters.push(map_param(name));
+            Expr::Identifier(ident) => {
+                if let Some(param) = extract_ident_param(ident) {
+                    *value = self.make_placeholder();
+                    self.parameters.push(param);
+                }
             }
             Expr::Value(Value::Placeholder(param)) if !self.is_own_placeholder(param) =>
             // this check is to avoid recursively replacing placeholders in the form of '?', or '$1', '$2', which we emit ourselves
diff --git a/src/webserver/http.rs b/src/webserver/http.rs
@@ -210,7 +210,7 @@ async fn render_sql(
     srv_req: &mut ServiceRequest,
     sql_file: Arc<ParsedSqlFile>,
 ) -> actix_web::Result<HttpResponse> {
-    let req_param = extract_request_info(srv_req).await;
+    let mut req_param = extract_request_info(srv_req).await;
     log::debug!("Received a request with the following parameters: {req_param:?}");
     let app_state = srv_req
         .app_data::<web::Data<AppState>>()
@@ -220,7 +220,7 @@ async fn render_sql(
 
     let (resp_send, resp_recv) = tokio::sync::oneshot::channel::<HttpResponse>();
     actix_web::rt::spawn(async move {
-        let database_entries_stream = stream_query_results(&app_state.db, &sql_file, &req_param);
+        let database_entries_stream = stream_query_results(&app_state.db, &sql_file, &mut req_param);
         let response_with_writer =
             build_response_header_and_stream(Arc::clone(&app_state), database_entries_stream).await;
         match response_with_writer {
