Add bundle-schema option to CLI

bundler/Cargo.toml

@@ -4,14 +4,16 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-anyhow = "1.0.79"
+anyhow = { version = "1.0.79" }
 axum = { version = "0.7.3" }
 axum-extra = { version = "0.9.1", features = ["typed-header"] }
 clap = { version = "4.4.13", features = ["derive", "env"] }
+clio = { version = "0.3.5", features = ["clap-parse"] }
 dotenvy = { version = "0.15.7" }
 flate2 = { version = "1.0.28" }
 headers = { version = "0.4.0" }
 humantime = { version = "2.1.0" }
+schemars = { version = "0.8.16" }
 serde = { version = "1.0.195", features = ["derive"] }
 serde_json = { version = "1.0.111" }
 sqlx = { version = "0.7.3", features = [

bundler/src/bundle.rs

@@ -1,9 +1,10 @@
 use crate::permissionables::{permissions::Permissions, proposals::Proposals, sessions::Sessions};
 use flate2::{write::GzEncoder, Compression};
+use schemars::{schema::RootSchema, schema_for, JsonSchema};
 use serde::Serialize;
 use sqlx::MySqlPool;
 use std::{
-    collections::hash_map::DefaultHasher,
+    collections::{hash_map::DefaultHasher, BTreeMap},
     hash::{Hash, Hasher},
 };
 use tar::Header;
@@ -148,4 +149,13 @@ where
 
         Ok(bundle_builder.into_inner()?.finish()?)
     }
+
+    /// Produces a set of schemas associated with the data in the bundle
+    pub fn schemas() -> BTreeMap<String, RootSchema> {
+        BTreeMap::from([
+            (Proposals::schema_name(), schema_for!(Proposals)),
+            (Sessions::schema_name(), schema_for!(Sessions)),
+            (Permissions::schema_name(), schema_for!(Permissions)),
+        ])
+    }
 }

bundler/src/main.rs

@@ -18,16 +18,19 @@ use axum::{
     http::{HeaderMap, StatusCode},
     response::IntoResponse,
     routing::get,
-    serve, Router,
+    Router,
 };
 use axum_extra::TypedHeader;
 use clap::Parser;
+use clio::ClioPath;
 use headers::{ETag, HeaderMapExt, IfNoneMatch};
 use require_bearer::RequireBearerLayer;
 use serde::Serialize;
 use sqlx::{mysql::MySqlPoolOptions, MySqlPool};
 use std::{
+    fs::File,
     hash::Hash,
+    io::Write,
     net::{Ipv4Addr, SocketAddr, SocketAddrV4},
     ops::Add,
     str::FromStr,
@@ -70,11 +73,20 @@ where
 
 /// A thread safe, mutable, wrapper around the [`BundleFile`]
 type CurrentBundle = Arc<RwLock<BundleFile<NoMetadata>>>;
-
 /// Bundler acts as a Open Policy Agent bundle server, providing permissionable data from the ISPyB database
+
 #[derive(Debug, Parser)]
 #[command(author, version, about, long_about= None)]
-struct Cli {
+enum Cli {
+    /// Run the service providing bundle data
+    Serve(ServeArgs),
+    /// Output the bundle schema
+    BundleSchema(BundleSchemaArgs),
+}
+
+/// Arguments to run the service with
+#[derive(Debug, Parser)]
+struct ServeArgs {
     /// The port to which this application should bind
     #[arg(short, long, env = "BUNDLER_PORT", default_value_t = 80)]
     port: u16,
@@ -92,11 +104,27 @@ struct Cli {
     polling_interval: humantime::Duration,
 }
 
+/// Arguments to output the schema with
+#[derive(Debug, Parser)]
+struct BundleSchemaArgs {
+    /// The path to write the schema to
+    #[arg(short, long, value_parser = clap::value_parser!(ClioPath).exists().is_dir())]
+    path: Option<ClioPath>,
+}
+
 #[tokio::main]
 async fn main() {
     dotenvy::dotenv().ok();
     let args = Cli::parse();
 
+    match args {
+        Cli::Serve(args) => serve(args).await,
+        Cli::BundleSchema(args) => bundle_schema(args),
+    }
+}
+
+/// Runs the service, pulling fresh bundles from ISPyB and serving them via the API
+async fn serve(args: ServeArgs) {
     tracing_subscriber::FmtSubscriber::builder()
         .with_max_level(args.log_level)
         .finish()
@@ -122,15 +150,15 @@ async fn main() {
         ispyb_pool,
         args.polling_interval.into(),
     ));
-    tasks.spawn(serve_app(args.port, app));
+    tasks.spawn(serve_endpoints(args.port, app));
     tasks.join_next().await.unwrap().unwrap()
 }
 
 /// Bind to the provided socket address and serve the application endpoints
-async fn serve_app(port: u16, app: Router) {
+async fn serve_endpoints(port: u16, app: Router) {
     let socket_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port));
     let listener = TcpListener::bind(socket_addr).await.unwrap();
-    serve(listener, app).await.unwrap()
+    axum::serve(listener, app).await.unwrap()
 }
 
 /// Periodically update the bundle with new data from ISPyB
@@ -180,3 +208,25 @@ async fn bundle_endpoint(
 async fn fallback_endpoint() -> impl IntoResponse {
     StatusCode::NOT_FOUND
 }
+
+/// Outputs the bundle schema as a set of files or to standard output
+fn bundle_schema(args: BundleSchemaArgs) {
+    let schemas = Bundle::<NoMetadata>::schemas()
+        .into_iter()
+        .map(|(name, schema)| (name, serde_json::to_string_pretty(&schema).unwrap()));
+    if let Some(path) = args.path {
+        for (name, schema) in schemas {
+            let mut schema_file =
+                File::create(path.clone().join(name).with_extension("json")).unwrap();
+            schema_file.write_all(schema.as_bytes()).unwrap();
+        }
+    } else {
+        println!(
+            "{}",
+            schemas
+                .map(|(_, schema)| schema)
+                .collect::<Vec<_>>()
+                .join("\n\n---\n\n")
+        )
+    }
+}

bundler/src/permissionables/permissions.rs

@@ -1,9 +1,10 @@
+use schemars::JsonSchema;
 use serde::Serialize;
 use sqlx::{query_as, MySqlPool};
 use std::collections::BTreeMap;
 
 /// A mapping of users to their permissions via groups
-#[derive(Debug, Default, PartialEq, Eq, Hash, Serialize)]
+#[derive(Debug, Default, PartialEq, Eq, Hash, Serialize, JsonSchema)]
 pub struct Permissions(BTreeMap<String, Vec<String>>);
 
 impl Permissions {

bundler/src/permissionables/proposals.rs

@@ -1,9 +1,10 @@
+use schemars::JsonSchema;
 use serde::Serialize;
 use sqlx::{query_as, MySqlPool};
 use std::collections::BTreeMap;
 
 /// A mapping of users to their proposals
-#[derive(Debug, Default, PartialEq, Eq, Hash, Serialize)]
+#[derive(Debug, Default, PartialEq, Eq, Hash, Serialize, JsonSchema)]
 pub struct Proposals(BTreeMap<String, Vec<u32>>);
 
 impl Proposals {

bundler/src/permissionables/sessions.rs

@@ -1,9 +1,10 @@
+use schemars::JsonSchema;
 use serde::Serialize;
 use sqlx::{query_as, MySqlPool};
 use std::collections::BTreeMap;
 
 /// A mapping of users to their sessions, possibly via proposals
-#[derive(Debug, Default, PartialEq, Eq, Hash, Serialize)]
+#[derive(Debug, Default, PartialEq, Eq, Hash, Serialize, JsonSchema)]
 pub struct Sessions(BTreeMap<String, Vec<(u32, u32)>>);
 
 impl Sessions {

charts/bundler/templates/deployment.yaml

@@ -30,6 +30,8 @@ spec:
             {{- toYaml .Values.securityContext | nindent 12 }}
           image: "{{ .Values.image.repository}}:{{ .Values.image.tag }}"
           imagePullPolicy: {{ .Values.image.pullPolicy }}
+          args:
+            - serve
           env:
             - name: BUNDLER_DATABASE_PASSWORD
               valueFrom: