add ability to get nex keys for staged accounts

Cargo.lock

@@ -959,7 +959,7 @@ checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084"
 
 [[package]]
 name = "nex-account"
-version = "0.2.3"
+version = "0.2.4"
 dependencies = [
  "base64",
  "bytemuck",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nex-account"
-version = "0.2.3"
+version = "0.2.4"
 edition = "2024"
 publish = ["spbr"]
 

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);
 }

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>,