From 218b2ebf1e809a8d355fe80d489fc96dd7301bc7 Mon Sep 17 00:00:00 2001
From: Alisina Bahadori <alisina.bm@gmail.com>
Date: Mon, 6 May 2024 01:27:13 -0400
Subject: [PATCH] Add spec generation command

---
 .gitignore  |  1 +
 src/main.rs | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index edb7bf0..c6ac18e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@ db
 
 # Keep everything in tests-data
 !tests-data/**
+openapi-spec.json
diff --git a/src/main.rs b/src/main.rs
index 5a78764..ac3ea7c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,11 @@
-use std::env::{self, args};
+use std::env;
 use std::io::{Error, ErrorKind, Result};
 
+use atlas_rs::api_docs;
+use tokio::io::AsyncWriteExt;
+
+const SPEC_FILENAME: &'static str = "openapi-spec.json";
+
 #[actix_web::main]
 async fn main() -> Result<()> {
     let db_variant = env::var("MAXMIND_DB_VARIANT").unwrap_or("GeoLite2-City".to_string());
@@ -22,18 +27,38 @@ async fn main() -> Result<()> {
         .parse()
         .expect("Invalid SWAGGER_UI_ENABLED value. Expected `false` or `true`");
 
-    let maxmind_db_arc = atlas_rs::init_db(&db_path, &db_variant).await;
-
-    let subcommand = args().skip(1).next();
+    let subcommand = env::args().skip(1).next();
 
     match subcommand.as_deref() {
         Some("server") | None => {
+            // Load or Initialize MaxMind database
+            let maxmind_db_arc = atlas_rs::init_db(&db_path, &db_variant).await;
             // Start Database Updater Daemon
             atlas_rs::start_db_refresher(maxmind_db_arc.clone(), update_interval);
             // Start Server
             atlas_rs::start_server(maxmind_db_arc, &host, port, swagger_ui_enabled).await
         }
-        Some("init") => Ok(()),
+        Some("init") => {
+            let _db = atlas_rs::init_db(&db_path, &db_variant).await;
+            println!("Initiation was successful");
+            Ok(())
+        }
+        Some("spec") => {
+            let api_doc = api_docs::api_doc();
+            let json_api_doc = api_doc.to_json().expect("Failed to generate API spec");
+
+            let mut file = tokio::fs::File::create(SPEC_FILENAME)
+                .await
+                .expect("Could not create spec file");
+
+            file.write_all(json_api_doc.as_bytes())
+                .await
+                .expect("Could not write to file");
+
+            println!("Generated {SPEC_FILENAME}");
+
+            Ok(())
+        }
         Some(command) => Err(Error::new(
             ErrorKind::InvalidInput,
             format!("Invalid command: {}", command),