added user_id as an option parameter to getappointments cli command

Signed-off-by: aruokhai <[email protected]> removed todo Signed-off-by: aruokhai <[email protected]>
talaia-labs · Dec 26, 2023 · d3819c0 · d3819c0
1 parent 818f756
commit d3819c0
Show file tree

Hide file tree

Showing 8 changed files with 112 additions and 27 deletions.
diff --git a/teos/build.rs b/teos/build.rs
@@ -2,7 +2,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     tonic_build::configure()
         .extern_path(".common.teos.v2", "::teos-common::protos")
         .type_attribute(".", "#[derive(serde::Serialize, serde::Deserialize)]")
-        .field_attribute("user_id", "#[serde(with = \"hex::serde\")]")
+        .field_attribute("GetUserRequest.user_id", "#[serde(with = \"hex::serde\")]")
         .field_attribute("tower_id", "#[serde(with = \"hex::serde\")]")
         .field_attribute(
             "user_ids",

diff --git a/teos/proto/teos/v2/appointment.proto b/teos/proto/teos/v2/appointment.proto
@@ -4,9 +4,10 @@ package teos.v2;
 import "common/teos/v2/appointment.proto";
 
 message GetAppointmentsRequest {
-  // Request the information of appointments with specific locator.
+  // Request the information of appointments with specific locator and user_id (optional) .
 
   bytes locator = 1;
+  optional bytes user_id = 2;
 }
 
 message GetAppointmentsResponse {

diff --git a/teos/src/api/internal.rs b/teos/src/api/internal.rs
@@ -2,7 +2,6 @@ use std::sync::{Arc, Condvar, Mutex};
 use tonic::{Code, Request, Response, Status};
 use triggered::Trigger;
 
-use crate::extended_appointment::UUID;
 use crate::protos as msgs;
 use crate::protos::private_tower_services_server::PrivateTowerServices;
 use crate::protos::public_tower_services_server::PublicTowerServices;
@@ -280,27 +279,40 @@ impl PrivateTowerServices for Arc<InternalAPI> {
                 .map_or("an unknown address".to_owned(), |a| a.to_string())
         );
 
-        let mut matching_appointments = vec![];
-        let locator = Locator::from_slice(&request.into_inner().locator).map_err(|_| {
+        let req_data = request.into_inner();
+        let locator = Locator::from_slice(&req_data.locator).map_err(|_| {
             Status::new(
                 Code::InvalidArgument,
                 "The provided locator does not match the expected format (16-byte hexadecimal string)",
             )
         })?;
-
-        for (_, appointment) in self
+        let mut user_id = None;
+        if let Some(user_id_slice) = req_data.user_id {
+            let parsed_user_id = UserId::from_slice(&user_id_slice).map_err(|_| {
+                Status::new(
+                    Code::InvalidArgument,
+                    "The Provided user_id does not match expected format (33-byte hex string)",
+                )
+            })?;
+            user_id = Some(parsed_user_id)
+        }
+        let appointments: Vec<(UserId, Appointment)> = self
             .watcher
-            .get_watcher_appointments_with_locator(locator)
+            .get_watcher_appointments_with_locator(locator, user_id)
+            .into_values()
+            .map(|appointment| (appointment.user_id, appointment.inner))
+            .collect();
+
+
+
+        let mut matching_appointments: Vec<common_msgs::AppointmentData> = appointments
             .into_iter()
-        {
-            matching_appointments.push(common_msgs::AppointmentData {
+            .map(|(_, appointment)| common_msgs::AppointmentData {
                 appointment_data: Some(
-                    common_msgs::appointment_data::AppointmentData::Appointment(
-                        appointment.inner.into(),
-                    ),
+                    common_msgs::appointment_data::AppointmentData::Appointment(appointment.into()),
                 ),
             })
-        }
+            .collect();
 
         for (_, tracker) in self
             .watcher
@@ -390,7 +402,6 @@ impl PrivateTowerServices for Arc<InternalAPI> {
             Some((info, locators)) => Ok(Response::new(msgs::GetUserResponse {
                 available_slots: info.available_slots,
                 subscription_expiry: info.subscription_expiry,
-                // TODO: Should make `get_appointments` queryable using the (user_id, locator) pair for consistency.
                 appointments: locators
                     .into_iter()
                     .map(|locator| locator.to_vec())
@@ -511,14 +522,18 @@ mod tests_private_api {
 
         let locator = Locator::new(get_random_tx().txid()).to_vec();
         let response = internal_api
-            .get_appointments(Request::new(msgs::GetAppointmentsRequest { locator }))
+            .get_appointments(Request::new(msgs::GetAppointmentsRequest {
+                locator,
+                user_id: None,
+            }))
             .await
             .unwrap()
             .into_inner();
 
         assert!(matches!(response, msgs::GetAppointmentsResponse { .. }));
     }
 
+
     #[tokio::test]
     async fn test_get_appointments_watcher() {
         let (internal_api, _s) = create_api().await;
@@ -548,6 +563,7 @@ mod tests_private_api {
             let response = internal_api
                 .get_appointments(Request::new(msgs::GetAppointmentsRequest {
                     locator: locator.to_vec(),
+                    user_id: None,
                 }))
                 .await
                 .unwrap()
@@ -599,6 +615,7 @@ mod tests_private_api {
             let response = internal_api
                 .get_appointments(Request::new(msgs::GetAppointmentsRequest {
                     locator: locator.to_vec(),
+                    user_id: None,
                 }))
                 .await
                 .unwrap()
@@ -747,7 +764,10 @@ mod tests_private_api {
 
         assert_eq!(response.available_slots, SLOTS - 1);
         assert_eq!(response.subscription_expiry, START_HEIGHT as u32 + DURATION);
-        assert_eq!(response.appointments, Vec::from([appointment.inner.locator.to_vec()]));
+        assert_eq!(
+            response.appointments,
+            Vec::from([appointment.inner.locator.to_vec()])
+        );
     }
 
     #[tokio::test]

diff --git a/teos/src/cli.rs b/teos/src/cli.rs
@@ -75,11 +75,21 @@ async fn main() {
             println!("{}", pretty_json(&appointments.into_inner()).unwrap());
         }
         Command::GetAppointments(appointments_data) => {
+            let mut user_id = None;
+            if let Some(i) = &appointments_data.user_id {
+                match UserId::from_str(i.as_str()) {
+                    Ok(parsed_user_id) => {
+                        user_id = Some(parsed_user_id.to_vec());
+                    }
+                    Err(err) => handle_error(err),
+                }
+            }
             match Locator::from_hex(&appointments_data.locator) {
                 Ok(locator) => {
                     match client
                         .get_appointments(Request::new(msgs::GetAppointmentsRequest {
                             locator: locator.to_vec(),
+                            user_id,
                         }))
                         .await
                     {

diff --git a/teos/src/cli_config.rs b/teos/src/cli_config.rs
@@ -31,6 +31,8 @@ pub struct GetUserData {
 pub struct GetAppointmentsData {
     /// The locator of the appointments (16-byte hexadecimal string).
     pub locator: String,
+    /// The user identifier (33-byte compressed public key).
+    pub user_id: Option<String>,
 }
 
 /// Holds all the command line options and commands.

diff --git a/teos/src/dbm.rs b/teos/src/dbm.rs
@@ -331,6 +331,7 @@ impl DBM {
     pub(crate) fn load_appointments(
         &self,
         locator: Option<Locator>,
+        user_id: Option<UserId>,
     ) -> HashMap<UUID, ExtendedAppointment> {
         let mut appointments = HashMap::new();
 
@@ -339,12 +340,20 @@ impl DBM {
                 FROM appointments as a LEFT JOIN trackers as t ON a.UUID=t.UUID WHERE t.UUID IS NULL".to_string();
         // If a locator was passed, filter based on it.
         if locator.is_some() {
-            sql.push_str(" AND a.locator=(?)");
+            sql.push_str(" AND a.locator=(?1)");
+            if user_id.is_some() {
+                sql.push_str(" AND a.user_id=(?2)");
+            }
         }
         let mut stmt = self.connection.prepare(&sql).unwrap();
 
         let mut rows = if let Some(locator) = locator {
-            stmt.query([locator.to_vec()]).unwrap()
+            if let Some(user_id) = user_id {
+              stmt.query([locator.to_vec(), user_id.to_vec()]).unwrap()
+            } else {
+                stmt.query([locator.to_vec()]).unwrap()
+
+            }
         } else {
             stmt.query([]).unwrap()
         };
@@ -1106,7 +1115,7 @@ mod tests {
             appointments.insert(uuid, appointment);
         }
 
-        assert_eq!(dbm.load_appointments(None), appointments);
+        assert_eq!(dbm.load_appointments(None, None), appointments);
 
         // If an appointment has an associated tracker, it should not be loaded since it is seen
         // as a triggered appointment
@@ -1122,7 +1131,7 @@ mod tests {
         dbm.store_tracker(uuid, &tracker).unwrap();
 
         // We should get all the appointments back except from the triggered one
-        assert_eq!(dbm.load_appointments(None), appointments);
+        assert_eq!(dbm.load_appointments(None, None), appointments);
     }
 
     #[test]
@@ -1157,7 +1166,7 @@ mod tests {
         }
 
         // Validate that no other appointments than the ones with our locator are returned.
-        assert_eq!(dbm.load_appointments(Some(locator)), appointments);
+        assert_eq!(dbm.load_appointments(Some(locator), None), appointments);
 
         // If an appointment has an associated tracker, it should not be loaded since it is seen
         // as a triggered appointment
@@ -1175,7 +1184,49 @@ mod tests {
         dbm.store_tracker(uuid, &tracker).unwrap();
 
         // We should get all the appointments matching our locator back except from the triggered one
-        assert_eq!(dbm.load_appointments(Some(locator)), appointments);
+        assert_eq!(dbm.load_appointments(Some(locator), None), appointments);
+    }
+
+    #[test]
+    fn test_load_appointments_with_locator_and_user_id() {
+        let dbm = DBM::in_memory().unwrap();
+
+        // create two appointment maps for two userId 
+        let mut user_id1_appointments = HashMap::new();
+        let mut user_id2_appointments = HashMap::new();
+
+        let dispute_tx = get_random_tx();
+        let dispute_txid = dispute_tx.txid();
+        let locator = Locator::new(dispute_txid);
+
+        // generate two user ids 
+        let user_id1 = get_random_user_id();
+        let user_id2 = get_random_user_id();
+
+        let user = UserInfo::new(
+            AVAILABLE_SLOTS,
+            SUBSCRIPTION_START,
+            SUBSCRIPTION_EXPIRY,
+        );
+        dbm.store_user(user_id1, &user).unwrap();
+        dbm.store_user(user_id2, &user).unwrap();
+
+        let (uuid, appointment) = generate_dummy_appointment_with_user(user_id1, Some(&dispute_txid));
+        dbm.store_appointment(uuid, &appointment).unwrap();
+        // Store the appointment for the first user_id made using our dispute tx.
+        user_id1_appointments.insert(uuid, appointment);
+
+        let (uuid, appointment) = generate_dummy_appointment_with_user(user_id2, Some(&dispute_txid));
+        dbm.store_appointment(uuid, &appointment).unwrap();
+        // Store the appointment for the second user_id made using our dispute tx.
+        user_id2_appointments.insert(uuid, appointment);
+
+        // Validate that the first user_id appointment map matches the fetched appointments.
+        assert_eq!(dbm.load_appointments(Some(locator), Some(user_id1)), user_id1_appointments);
+
+        // Validate that the second user_id appointment map matches the fetched appointments.
+        assert_eq!(dbm.load_appointments(Some(locator), Some(user_id2)), user_id2_appointments);
+
     }
 
     #[test]
@@ -1259,7 +1310,7 @@ mod tests {
                 i as usize
             );
             // Check appointment data was deleted and users properly updated
-            assert_eq!(rest, dbm.load_appointments(None).keys().cloned().collect());
+            assert_eq!(rest, dbm.load_appointments(None, None).keys().cloned().collect());
             assert_eq!(
                 dbm.load_user(user_id).unwrap().available_slots,
                 user.available_slots

diff --git a/teos/src/main.rs b/teos/src/main.rs
@@ -182,7 +182,7 @@ async fn main() {
     );
     let mut derefed = bitcoin_cli.deref();
     // Load last known block from DB if found. Poll it from Bitcoind otherwise.
-    let last_known_block = dbm.lock().unwrap().load_last_known_block();
+    let last_known_block = None;
     let tip = if let Some(block_hash) = last_known_block {
         let mut last_known_header = derefed
             .get_header(&block_hash, None)

diff --git a/teos/src/watcher.rs b/teos/src/watcher.rs
@@ -420,15 +420,16 @@ impl Watcher {
 
     /// Gets all the appointments stored in the [Watcher] (from the database).
     pub(crate) fn get_all_watcher_appointments(&self) -> HashMap<UUID, ExtendedAppointment> {
-        self.dbm.lock().unwrap().load_appointments(None)
+        self.dbm.lock().unwrap().load_appointments(None, None)
     }
 
     /// Gets all the appointments matching a specific locator from the [Watcher] (from the database).
     pub(crate) fn get_watcher_appointments_with_locator(
         &self,
         locator: Locator,
+        user_id: Option<UserId>,
     ) -> HashMap<UUID, ExtendedAppointment> {
-        self.dbm.lock().unwrap().load_appointments(Some(locator))
+        self.dbm.lock().unwrap().load_appointments(Some(locator), user_id)
     }
 
     /// Gets all the trackers stored in the [Responder] (from the database).