From dcc0d1b2b79c889046acab21e5ba9dba2dd40816 Mon Sep 17 00:00:00 2001
From: Maple Nebel <maple.neb@gmail.com>
Date: Tue, 7 Jul 2026 07:46:55 +0200
Subject: [PATCH] add ability to get nex keys for staged accounts

---
 Cargo.lock  |  2 +-
 Cargo.toml  |  2 +-
 grpc.proto  |  1 +
 src/main.rs | 31 ++++++++++++++++++++++++++++++-
 4 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1c042b6..349b49f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -959,7 +959,7 @@ checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084"
 
 [[package]]
 name = "nex-account"
-version = "0.2.3"
+version = "0.2.4"
 dependencies = [
  "base64",
  "bytemuck",
diff --git a/Cargo.toml b/Cargo.toml
index 7bdbfaf..f314775 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "nex-account"
-version = "0.2.3"
+version = "0.2.4"
 edition = "2024"
 publish = ["spbr"]
 
diff --git a/grpc.proto b/grpc.proto
index ddbddbf..2cb0247 100644
--- a/grpc.proto
+++ b/grpc.proto
@@ -29,5 +29,6 @@ message NexKey {
 service NexAccountService {
   rpc CreateNewSequentialOrUpdateAndGetAccount(ActCreateInfo) returns (NexKey);
   rpc GetNexKeyByPid(PID) returns (NexKey);
+  rpc GetNexKeyByPidMaybeStaged(PID) returns (NexKey);
   rpc StageNewAccount(ActStageInfo) returns (ActStageReturn);
 }
diff --git a/src/main.rs b/src/main.rs
index 1e074d4..2d89256 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -69,7 +69,6 @@ pub fn derive_key(pid: i32, password: &[u8]) -> [u8; 16] {
 
     key
 }
-
 #[tonic::async_trait]
 impl NexAccountService for NexAccountServer {
     async fn create_new_sequential_or_update_and_get_account(
@@ -153,6 +152,36 @@ impl NexAccountService for NexAccountServer {
             ))
         }
     }
+    async fn get_nex_key_by_pid_maybe_staged(
+        &self,
+        request: tonic::Request<Pid>,
+    ) -> Result<Response<NexKey>, tonic::Status> {
+        let pid = request.into_inner().pid;
+
+        let sql_res = query!("select nex_key from nex_accounts where pid = $1", pid)
+            .fetch_optional(&self.pool)
+            .await
+            .map_err(db_neverfail)?;
+
+        if let Some(res) = sql_res {
+            return Ok(Response::new(NexKey { key: res.nex_key }));
+        }
+        let sql_res = query!(
+            "select nex_key from staged_nex_accounts where pid = $1",
+            pid
+        )
+        .fetch_optional(&self.pool)
+        .await
+        .map_err(db_neverfail)?;
+
+        if let Some(res) = sql_res {
+            Ok(Response::new(NexKey { key: res.nex_key }))
+        } else {
+            Err(tonic::Status::not_found(
+                "nex account with provided pid does not exist",
+            ))
+        }
+    }
     async fn stage_new_account(
         &self,
         request: tonic::Request<ActStageInfo>,