From d94ca9e1bb79fcd9d4f64c9a6270d2623bf565a3 Mon Sep 17 00:00:00 2001
From: kieran <kieran@harkin.me>
Date: Mon, 10 Mar 2025 12:33:11 +0000
Subject: [PATCH] feat: add data migration for dns entries

---
 src/bin/api.rs            |  4 +++
 src/data_migration/dns.rs | 55 +++++++++++++++++++++++++++++++++++++++
 src/data_migration/mod.rs | 31 ++++++++++++++++++++++
 src/lib.rs                |  2 +-
 4 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 src/data_migration/dns.rs
 create mode 100644 src/data_migration/mod.rs

diff --git a/src/bin/api.rs b/src/bin/api.rs
index 024e527..5338c9b 100644
--- a/src/bin/api.rs
+++ b/src/bin/api.rs
@@ -4,6 +4,7 @@ use clap::Parser;
 use config::{Config, File};
 use lnvps::api;
 use lnvps::cors::CORS;
+use lnvps::data_migration::run_data_migrations;
 use lnvps::exchange::{DefaultRateCache, ExchangeRateService};
 use lnvps::invoice::InvoiceHandler;
 use lnvps::lightning::get_node;
@@ -84,6 +85,9 @@ async fn main() -> Result<(), Error> {
     }
     let db: Arc<dyn LNVpsDb> = Arc::new(db);
 
+    // run data migrations
+    run_data_migrations(db.clone(), &settings).await?;
+
     let nostr_client = if let Some(ref c) = settings.nostr {
         let cx = Client::builder().signer(Keys::parse(&c.nsec)?).build();
         for r in &c.relays {
diff --git a/src/data_migration/dns.rs b/src/data_migration/dns.rs
new file mode 100644
index 0000000..c673cc0
--- /dev/null
+++ b/src/data_migration/dns.rs
@@ -0,0 +1,55 @@
+use crate::data_migration::DataMigration;
+use crate::dns::{BasicRecord, DnsServer};
+use crate::settings::Settings;
+use anyhow::Result;
+use lnvps_db::LNVpsDb;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+pub struct DnsDataMigration {
+    db: Arc<dyn LNVpsDb>,
+    dns: Arc<dyn DnsServer>,
+}
+
+impl DnsDataMigration {
+    pub fn new(db: Arc<dyn LNVpsDb>, settings: &Settings) -> Option<Self> {
+        let dns = settings.get_dns().ok().flatten()?;
+        Some(Self { db, dns })
+    }
+}
+
+impl DataMigration for DnsDataMigration {
+    fn migrate(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
+        let db = self.db.clone();
+        let dns = self.dns.clone();
+        Box::pin(async move {
+            let vms = db.list_vms().await?;
+
+            for vm in vms {
+                let mut ips = db.list_vm_ip_assignments(vm.id).await?;
+                for mut ip in &mut ips {
+                    let mut did_change = false;
+                    if ip.dns_forward.is_none() {
+                        let rec = BasicRecord::forward(&ip)?;
+                        let r = dns.add_record(&rec).await?;
+                        ip.dns_forward = Some(r.name);
+                        ip.dns_forward_ref = r.id;
+                        did_change = true;
+                    }
+                    if ip.dns_reverse.is_none() {
+                        let rec = BasicRecord::reverse(&ip)?;
+                        let r = dns.add_record(&rec).await?;
+                        ip.dns_reverse = Some(r.name);
+                        ip.dns_reverse_ref = r.id;
+                        did_change = true;
+                    }
+                    if did_change {
+                        db.update_vm_ip_assignment(ip).await?;
+                    }
+                }
+            }
+            Ok(())
+        })
+    }
+}
diff --git a/src/data_migration/mod.rs b/src/data_migration/mod.rs
new file mode 100644
index 0000000..3f758fe
--- /dev/null
+++ b/src/data_migration/mod.rs
@@ -0,0 +1,31 @@
+use crate::data_migration::dns::DnsDataMigration;
+use crate::settings::Settings;
+use anyhow::Result;
+use lnvps_db::LNVpsDb;
+use log::{error, info};
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+mod dns;
+
+/// Basic data migration to run at startup
+pub trait DataMigration: Send + Sync {
+    fn migrate(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>;
+}
+
+pub async fn run_data_migrations(db: Arc<dyn LNVpsDb>, settings: &Settings) -> Result<()> {
+    let mut migrations: Vec<Arc<dyn DataMigration>> = vec![];
+    if let Some(d) = DnsDataMigration::new(db.clone(), settings) {
+        migrations.push(Arc::new(d));
+    }
+
+    info!("Running {} data migrations", migrations.len());
+    for migration in migrations {
+        if let Err(e) = migration.migrate().await {
+            error!("Error running data migration: {}", e);
+        }
+    }
+
+    Ok(())
+}
diff --git a/src/lib.rs b/src/lib.rs
index 0a07ea0..a1ed9af 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -14,6 +14,6 @@ pub mod settings;
 pub mod ssh_client;
 pub mod status;
 pub mod worker;
-
+pub mod data_migration;
 #[cfg(test)]
 pub mod mocks;