From 411cd681a92383c8b7fd8c5966813e109099acbd Mon Sep 17 00:00:00 2001
From: Orne Brocaar <info@brocaar.com>
Date: Thu, 22 Feb 2024 12:44:14 +0000
Subject: [PATCH] Add cmd to migrate ds from Redis to Pg.

---
 chirpstack/src/cmd/migrate_ds_to_pg.rs | 56 ++++++++++++++++++++++++++
 chirpstack/src/cmd/mod.rs              |  1 +
 chirpstack/src/main.rs                 |  4 ++
 3 files changed, 61 insertions(+)
 create mode 100644 chirpstack/src/cmd/migrate_ds_to_pg.rs

diff --git a/chirpstack/src/cmd/migrate_ds_to_pg.rs b/chirpstack/src/cmd/migrate_ds_to_pg.rs
new file mode 100644
index 00000000..4b0c26cf
--- /dev/null
+++ b/chirpstack/src/cmd/migrate_ds_to_pg.rs
@@ -0,0 +1,56 @@
+use anyhow::Result;
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use tracing::{debug, info};
+
+use crate::storage::{self, device_session, error::Error, get_async_db_conn, schema::device};
+use lrwn::{DevAddr, EUI64};
+
+pub async fn run() -> Result<()> {
+    storage::setup().await?;
+
+    info!("Migrating device-sessions from Redis to PostgreSQL");
+    info!("Getting DevEUIs from PostgreSQL without device-session");
+
+    let dev_euis: Vec<EUI64> = device::dsl::device
+        .select(device::dsl::dev_eui)
+        .filter(device::dsl::device_session.is_null())
+        .load(&mut get_async_db_conn().await?)
+        .await?;
+
+    info!(
+        "There are {} devices in PostgreSQL without device-session set",
+        dev_euis.len()
+    );
+
+    for dev_eui in &dev_euis {
+        debug!(dev_eui = %dev_eui, "Migrating device-session");
+
+        let ds = match device_session::get(dev_eui).await {
+            Ok(v) => v,
+            Err(e) => match e {
+                Error::NotFound(_) => {
+                    debug!(dev_eui = %dev_eui, "Device does not have a device-session");
+                    continue;
+                }
+                _ => {
+                    return Err(anyhow::Error::new(e));
+                }
+            },
+        };
+
+        storage::device::partial_update(
+            *dev_eui,
+            &storage::device::DeviceChangeset {
+                dev_addr: Some(Some(DevAddr::from_slice(&ds.dev_addr)?)),
+                device_session: Some(Some(ds)),
+                ..Default::default()
+            },
+        )
+        .await?;
+
+        debug!(dev_eui = %dev_eui, "Device-session migrated");
+    }
+
+    Ok(())
+}
diff --git a/chirpstack/src/cmd/mod.rs b/chirpstack/src/cmd/mod.rs
index 08970f81..e85e3ed5 100644
--- a/chirpstack/src/cmd/mod.rs
+++ b/chirpstack/src/cmd/mod.rs
@@ -1,5 +1,6 @@
 pub mod configfile;
 pub mod create_api_key;
 pub mod import_legacy_lorawan_devices_repository;
+pub mod migrate_ds_to_pg;
 pub mod print_ds;
 pub mod root;
diff --git a/chirpstack/src/main.rs b/chirpstack/src/main.rs
index afcac15c..35f90f8d 100644
--- a/chirpstack/src/main.rs
+++ b/chirpstack/src/main.rs
@@ -78,6 +78,9 @@ enum Commands {
         #[arg(short, long, value_name = "NAME")]
         name: String,
     },
+
+    /// Migrate device-sessions from Redis to PostgreSQL.
+    MigrateDeviceSessionsToPostgres {},
 }
 
 #[tokio::main]
@@ -116,6 +119,7 @@ async fn main() -> Result<()> {
                 .unwrap()
         }
         Some(Commands::CreateApiKey { name }) => cmd::create_api_key::run(name).await?,
+        Some(Commands::MigrateDeviceSessionsToPostgres {}) => cmd::migrate_ds_to_pg::run().await?,
         None => cmd::root::run().await?,
     }