Run sqlite statements

Signed-off-by: Ryan Levick <[email protected]>

Author: rylev

build.rs

@@ -68,16 +68,6 @@ error: the `wasm32-wasi` target is not installed
     std::fs::create_dir_all("target/test-programs").unwrap();
 
     build_wasm_test_program("core-wasi-test.wasm", "crates/core/tests/core-wasi-test");
-    // build_wasm_test_program("redis-rust.wasm", "crates/trigger-redis/tests/rust");
-    // build_wasm_test_program(
-    //     "spin-http-benchmark.wasm",
-    //     "crates/trigger-http/benches/spin-http-benchmark",
-    // );
-    // build_wasm_test_program(
-    //     "wagi-benchmark.wasm",
-    //     "crates/trigger-http/benches/wagi-benchmark",
-    // );
-    // build_wasm_test_program("timer_app_example.wasm", "examples/spin-timer/app-example");
 
     cargo_build(TIMER_TRIGGER_INTEGRATION_TEST);
 }

crates/factor-sqlite/src/lib.rs

@@ -149,13 +149,29 @@ pub trait DefaultLabelResolver: Send + Sync {
     fn default(&self, label: &str) -> Option<Arc<dyn ConnectionCreator>>;
 }
 
+#[derive(Clone)]
 pub struct AppState {
     /// A map from component id to a set of allowed database labels.
     allowed_databases: HashMap<String, Arc<HashSet<String>>>,
     /// A function for mapping from database name to a connection creator.
     get_connection_creator: host::ConnectionCreatorGetter,
 }
 
+impl AppState {
+    /// Get a connection for a given database label.
+    ///
+    /// Returns `None` if there is no connection creator for the given label.
+    pub async fn get_connection(
+        &self,
+        label: &str,
+    ) -> Option<Result<Box<dyn Connection>, v2::Error>> {
+        let connection = (self.get_connection_creator)(label)?
+            .create_connection()
+            .await;
+        Some(connection)
+    }
+}
+
 /// A creator of a connections for a particular SQLite database.
 #[async_trait]
 pub trait ConnectionCreator: Send + Sync {

crates/trigger/Cargo.toml

@@ -42,7 +42,7 @@ spin-factors-executor = { path = "../factors-executor" }
 spin-runtime-config = { path = "../runtime-config" }
 spin-telemetry = { path = "../telemetry" }
 terminal = { path = "../terminal" }
-tokio = { version = "1.23", features = ["fs"] }
+tokio = { version = "1.23", features = ["fs", "rt"] }
 tracing = { workspace = true }
 
 [lints]

crates/trigger/src/cli.rs

@@ -1,4 +1,5 @@
 mod launch_metadata;
+mod sqlite_statements;
 mod summary;
 
 use std::future::Future;
@@ -13,6 +14,7 @@ use spin_common::{arg_parser::parse_kv, sloth};
 use spin_core::async_trait;
 use spin_factors_executor::{ComponentLoader, FactorsExecutor};
 use spin_runtime_config::{ResolvedRuntimeConfig, UserProvidedPath};
+use sqlite_statements::SqlStatementExecutorHook;
 use summary::KeyValueDefaultStoreSummaryHook;
 
 use crate::factors::{TriggerFactors, TriggerFactorsRuntimeConfig};
@@ -198,6 +200,7 @@ impl<T: Trigger> FactorsTriggerCommand<T> {
                     state_dir: self.state_dir.as_deref(),
                     local_app_dir: local_app_dir.as_deref(),
                     initial_key_values: self.key_values,
+                    sqlite_statements: self.sqlite_statements,
                     allow_transient_write: self.allow_transient_write,
                     follow_components,
                     log_dir: self.log,
@@ -277,6 +280,8 @@ pub struct TriggerAppOptions<'a> {
     local_app_dir: Option<&'a str>,
     /// Initial key/value pairs to set in the app's default store.
     initial_key_values: Vec<(String, String)>,
+    /// SQLite statements to run.
+    sqlite_statements: Vec<String>,
     /// Whether to allow transient writes to mounted files
     allow_transient_write: bool,
     /// Which components should have their logs followed.
@@ -351,8 +356,6 @@ impl<T: Trigger> TriggerAppBuilder<T> {
         )
         .context("failed to create factors")?;
 
-        // TODO(factors): handle: self.sqlite_statements
-
         // TODO: port the rest of the component loader logic
         struct SimpleComponentLoader;
 
@@ -421,6 +424,7 @@ impl<T: Trigger> TriggerAppBuilder<T> {
         // TODO:
         // builder.hooks(SummariseRuntimeConfigHook::new(&self.runtime_config_file));
         executor.add_hooks(KeyValueDefaultStoreSummaryHook);
+        executor.add_hooks(SqlStatementExecutorHook::new(options.sqlite_statements));
         // builder.hooks(SqlitePersistenceMessageHook);
 
         let configured_app = {

crates/trigger/src/cli/sqlite_statements.rs

@@ -0,0 +1,99 @@
+use anyhow::Context as _;
+use spin_factor_sqlite::SqliteFactor;
+use spin_factors::RuntimeFactors;
+use spin_factors_executor::ExecutorHooks;
+
+/// The default sqlite label
+const DEFAULT_SQLITE_LABEL: &str = "default";
+
+/// ExecutorHook for executing sqlite statements.
+///
+/// This executor assumes that the configured app has access to `SqliteFactor`.
+/// It will silently ignore the hook if the app does not have access to `SqliteFactor`.
+pub struct SqlStatementExecutorHook {
+    sql_statements: Vec<String>,
+}
+
+impl SqlStatementExecutorHook {
+    /// Creates a new SqlStatementExecutorHook
+    ///
+    /// The statements can be either a list of raw SQL statements or a list of `@{file:label}` statements.
+    pub fn new(sql_statements: Vec<String>) -> Self {
+        Self { sql_statements }
+    }
+}
+
+impl<F: RuntimeFactors, U> ExecutorHooks<F, U> for SqlStatementExecutorHook {
+    fn configure_app(
+        &mut self,
+        configured_app: &spin_factors::ConfiguredApp<F>,
+    ) -> anyhow::Result<()> {
+        if self.sql_statements.is_empty() {
+            return Ok(());
+        }
+        let Some(sqlite) = configured_app.app_state::<SqliteFactor>().ok() else {
+            return Ok(());
+        };
+        if let Ok(current) = tokio::runtime::Handle::try_current() {
+            let _ = current.spawn(execute(sqlite.clone(), self.sql_statements.clone()));
+        }
+        Ok(())
+    }
+}
+
+/// Executes the sql statements.
+pub async fn execute(
+    sqlite: spin_factor_sqlite::AppState,
+    sql_statements: Vec<String>,
+) -> anyhow::Result<()> {
+    let get_database = |label| {
+        let sqlite = &sqlite;
+        async move {
+            sqlite
+                .get_connection(label)
+                .await
+                .transpose()
+                .with_context(|| format!("failed connect to database with label '{label}'"))
+        }
+    };
+
+    for statement in &sql_statements {
+        if let Some(config) = statement.strip_prefix('@') {
+            let (file, label) = parse_file_and_label(config)?;
+            let database = get_database(label).await?.with_context(|| {
+                    format!(
+                        "based on the '@{config}' a registered database named '{label}' was expected but not found."
+                    )
+                })?;
+            let sql = std::fs::read_to_string(file).with_context(|| {
+                format!("could not read file '{file}' containing sql statements")
+            })?;
+            database.execute_batch(&sql).await.with_context(|| {
+                format!("failed to execute sql against database '{label}' from file '{file}'")
+            })?;
+        } else {
+            let Some(default) = get_database(DEFAULT_SQLITE_LABEL).await? else {
+                debug_assert!(false, "the '{DEFAULT_SQLITE_LABEL}' sqlite database should always be available but for some reason was not");
+                return Ok(());
+            };
+            default
+                    .query(statement, Vec::new())
+                    .await
+                    .with_context(|| format!("failed to execute following sql statement against default database: '{statement}'"))?;
+        }
+    }
+    Ok(())
+}
+
+/// Parses a @{file:label} sqlite statement
+fn parse_file_and_label(config: &str) -> anyhow::Result<(&str, &str)> {
+    let config = config.trim();
+    let (file, label) = match config.split_once(':') {
+        Some((_, label)) if label.trim().is_empty() => {
+            anyhow::bail!("database label is empty in the '@{config}' sqlite statement")
+        }
+        Some((file, label)) => (file.trim(), label.trim()),
+        None => (config, "default"),
+    };
+    Ok((file, label))
+}