Merge pull request #84 from ontodev/improve-serve

Improve server, rename --connection option
ontodev · Mar 21, 2024 · d58daae · d58daae
2 parents 180c6bb + 5a175e6
commit d58daae
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 19 deletions.
diff --git a/examples/penguins/nanobot.toml b/examples/penguins/nanobot.toml
diff --git a/src/main.rs b/src/main.rs
@@ -59,7 +59,7 @@ async fn main() -> Result<(), NanobotError> {
                 .about("Initialises things")
                 .arg(
                     arg!(
-                        -d --database <FILE> "Specifies a custom database name"
+                        -c --connection <URL> "Specifies a database connection URL or file"
                     )
                     .required(false)
                     .value_parser(value_parser!(String)),
@@ -87,13 +87,21 @@ async fn main() -> Result<(), NanobotError> {
                         .value_parser(value_parser!(String)),
                 ),
         )
-        .subcommand(Command::new("serve").about("Run HTTP server"))
+        .subcommand(
+            Command::new("serve").about("Run HTTP server").arg(
+                arg!(
+                    -c --connection <URL> "Specifies a database connection URL or file"
+                )
+                .required(false)
+                .value_parser(value_parser!(String)),
+            ),
+        )
         .get_matches();
 
     let exit_result = match matches.subcommand() {
         Some(("init", sub_matches)) => {
-            if let Some(d) = sub_matches.get_one::<String>("database") {
-                config.connection(d);
+            if let Some(c) = sub_matches.get_one::<String>("connection") {
+                config.connection(c);
             }
             if sub_matches.get_flag("create_only") {
                 config.create_only(true);
@@ -129,13 +137,29 @@ async fn main() -> Result<(), NanobotError> {
             };
             Ok(result)
         }
-        Some(("serve", _sub_matches)) => {
-            build_valve(&mut config).await?;
+        Some(("serve", sub_matches)) => {
+            if let Some(c) = sub_matches.get_one::<String>("connection") {
+                config.connection(c);
+            }
+            if config.connection == ":memory:" {
+                (config.valve, config.pool) = {
+                    let valve = Valve::build(&config.valve_path, &config.connection).await?;
+                    let pool = valve.pool.clone();
+                    let _ = valve.load_all_tables(true).await;
+                    let table_select = Select::new("\"table\"");
+                    config.table = get_table_from_pool(&pool, &table_select)
+                        .await
+                        .unwrap()
+                        .clone();
+                    (Some(valve), Some(pool))
+                };
+            } else {
+                build_valve(&mut config).await?;
+            }
             serve::app(&config)
         }
         _ => Err(String::from(
-            "Unrecognised or missing subcommand, but CGI environment vars are \
-                             undefined",
+            "Unrecognised or missing subcommand, but CGI environment vars are undefined",
         )),
     };
 

diff --git a/src/serve.rs b/src/serve.rs
@@ -22,6 +22,7 @@ use serde_json::{json, Value as SerdeValue};
 use std::{
     collections::HashMap, collections::HashSet, net::SocketAddr, process::Command, sync::Arc,
 };
+use tokio::signal;
 use tower_http::services::ServeDir;
 use wiring_rs::util::signature;
 
@@ -75,16 +76,41 @@ pub async fn app(config: &Config) -> Result<String, String> {
     // run our app with hyper
     // `axum::Server` is a re-export of `hyper::Server`
     let addr = SocketAddr::from(([0, 0, 0, 0], config.port));
-    tracing::info!("listening on {}", addr);
+    println!("Running Nanobot server at http://{addr}");
+    println!("Press Control-C to quit.");
     if let Err(e) = axum::Server::bind(&addr)
         .serve(app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal())
         .await
     {
         return Err(e.to_string());
     }
 
-    let hello = String::from("Hello, world!");
-    Ok(hello)
+    Ok("Stopping Nanobot server...".into())
+}
+
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
+    }
 }
 
 async fn root() -> impl IntoResponse {