feat: Add application to schema and API

Cargo.toml

@@ -39,7 +39,7 @@ sentry = { version = "0.41.0", default-features = false, features = [
     "tracing",
     "logs"
 ] }
-sentry_protos = "0.4.10"
+sentry_protos = "0.4.11"
 serde = "1.0.214"
 serde_yaml = "0.9.34"
 sha2 = "0.10.8"

benches/store_bench.rs

@@ -58,7 +58,7 @@ async fn get_pending_activations(num_activations: u32, num_workers: u32) {
             let mut num_activations_processed = 0;
 
             while store
-                .get_pending_activation(Some(&ns))
+                .get_pending_activation(Some("sentry"), Some(&ns))
                 .await
                 .unwrap()
                 .is_some()

migrations/0004_add_application.sql

@@ -0,0 +1 @@
+ALTER TABLE inflight_taskactivations ADD COLUMN application TEXT NOT NULL DEFAULT '';

src/grpc/server.rs

@@ -24,10 +24,11 @@ impl ConsumerService for TaskbrokerServer {
         request: Request<GetTaskRequest>,
     ) -> Result<Response<GetTaskResponse>, Status> {
         let start_time = Instant::now();
+        let application = &request.get_ref().application;
         let namespace = &request.get_ref().namespace;
         let inflight = self
             .store
-            .get_pending_activation(namespace.as_deref())
+            .get_pending_activation(application.as_deref(), namespace.as_deref())
             .await;
 
         match inflight {
@@ -97,14 +98,18 @@ impl ConsumerService for TaskbrokerServer {
         }
         metrics::histogram!("grpc_server.set_status.duration").record(start_time.elapsed());
 
-        let Some(FetchNextTask { ref namespace }) = request.get_ref().fetch_next_task else {
+        let Some(FetchNextTask {
+            ref namespace,
+            ref application,
+        }) = request.get_ref().fetch_next_task
+        else {
             return Ok(Response::new(SetTaskStatusResponse { task: None }));
         };
 
         let start_time = Instant::now();
         let res = match self
             .store
-            .get_pending_activation(namespace.as_deref())
+            .get_pending_activation(application.as_deref(), namespace.as_deref())
             .await
         {
             Err(e) => {

src/grpc/server_tests.rs

@@ -1,5 +1,8 @@
+use prost::Message;
 use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService;
-use sentry_protos::taskbroker::v1::{FetchNextTask, GetTaskRequest, SetTaskStatusRequest};
+use sentry_protos::taskbroker::v1::{
+    FetchNextTask, GetTaskRequest, SetTaskStatusRequest, TaskActivation,
+};
 use tonic::{Code, Request};
 
 use crate::grpc::server::TaskbrokerServer;
@@ -10,7 +13,10 @@ use crate::test_utils::{create_test_store, make_activations};
 async fn test_get_task() {
     let store = create_test_store().await;
     let service = TaskbrokerServer { store };
-    let request = GetTaskRequest { namespace: None };
+    let request = GetTaskRequest {
+        namespace: None,
+        application: None,
+    };
     let response = service.get_task(Request::new(request)).await;
     assert!(response.is_err());
     let e = response.unwrap_err();
@@ -62,7 +68,10 @@ async fn test_get_task_success() {
     store.store(activations).await.unwrap();
 
     let service = TaskbrokerServer { store };
-    let request = GetTaskRequest { namespace: None };
+    let request = GetTaskRequest {
+        namespace: None,
+        application: None,
+    };
     let response = service.get_task(Request::new(request)).await;
     assert!(response.is_ok());
     let resp = response.unwrap();
@@ -71,6 +80,54 @@ async fn test_get_task_success() {
     assert!(task.id == "id_0");
 }
 
+#[tokio::test]
+#[allow(deprecated)]
+async fn test_get_task_with_application_success() {
+    let store = create_test_store().await;
+    let mut activations = make_activations(2);
+
+    let mut payload = TaskActivation::decode(&activations[1].activation as &[u8]).unwrap();
+    payload.application = Some("hammers".into());
+    activations[1].activation = payload.encode_to_vec();
+    activations[1].application = "hammers".into();
+
+    store.store(activations).await.unwrap();
+
+    let service = TaskbrokerServer { store };
+    let request = GetTaskRequest {
+        namespace: None,
+        application: Some("hammers".into()),
+    };
+    let response = service.get_task(Request::new(request)).await;
+    assert!(response.is_ok());
+    let resp = response.unwrap();
+    assert!(resp.get_ref().task.is_some());
+    let task = resp.get_ref().task.as_ref().unwrap();
+    assert_eq!(task.id, "id_1");
+    assert_eq!(task.application, Some("hammers".into()));
+}
+
+#[tokio::test]
+#[allow(deprecated)]
+async fn test_get_task_with_namespace_requires_application() {
+    let store = create_test_store().await;
+    let activations = make_activations(2);
+    let namespace = activations[0].namespace.clone();
+
+    store.store(activations).await.unwrap();
+
+    let service = TaskbrokerServer { store };
+    let request = GetTaskRequest {
+        namespace: Some(namespace),
+        application: None,
+    };
+    let response = service.get_task(Request::new(request)).await;
+
+    assert!(response.is_err());
+    let e = response.unwrap_err();
+    assert_eq!(e.code(), Code::NotFound);
+}
+
 #[tokio::test]
 #[allow(deprecated)]
 async fn test_set_task_status_success() {
@@ -80,7 +137,10 @@ async fn test_set_task_status_success() {
 
     let service = TaskbrokerServer { store };
 
-    let request = GetTaskRequest { namespace: None };
+    let request = GetTaskRequest {
+        namespace: None,
+        application: None,
+    };
     let response = service.get_task(Request::new(request)).await;
     assert!(response.is_ok());
     let resp = response.unwrap();
@@ -91,7 +151,10 @@ async fn test_set_task_status_success() {
     let request = SetTaskStatusRequest {
         id: "id_0".to_string(),
         status: 5, // Complete
-        fetch_next_task: Some(FetchNextTask { namespace: None }),
+        fetch_next_task: Some(FetchNextTask {
+            namespace: None,
+            application: None,
+        }),
     };
     let response = service.set_task_status(Request::new(request)).await;
     assert!(response.is_ok());
@@ -100,3 +163,95 @@ async fn test_set_task_status_success() {
     let task = resp.get_ref().task.as_ref().unwrap();
     assert_eq!(task.id, "id_1");
 }
+
+#[tokio::test]
+#[allow(deprecated)]
+async fn test_set_task_status_with_application() {
+    let store = create_test_store().await;
+    let mut activations = make_activations(2);
+
+    let mut payload = TaskActivation::decode(&activations[1].activation as &[u8]).unwrap();
+    payload.application = Some("hammers".into());
+    activations[1].activation = payload.encode_to_vec();
+    activations[1].application = "hammers".into();
+
+    store.store(activations).await.unwrap();
+
+    let service = TaskbrokerServer { store };
+    let request = SetTaskStatusRequest {
+        id: "id_0".to_string(),
+        status: 5, // Complete
+        fetch_next_task: Some(FetchNextTask {
+            application: Some("hammers".into()),
+            namespace: None,
+        }),
+    };
+    let response = service.set_task_status(Request::new(request)).await;
+    assert!(response.is_ok());
+
+    let resp = response.unwrap();
+    let task_opt = &resp.get_ref().task;
+    assert!(task_opt.is_some());
+
+    let task = task_opt.as_ref().unwrap();
+    assert_eq!(task.id, "id_1");
+    assert_eq!(task.application, Some("hammers".into()));
+}
+
+#[tokio::test]
+#[allow(deprecated)]
+async fn test_set_task_status_with_application_no_match() {
+    let store = create_test_store().await;
+    let mut activations = make_activations(2);
+
+    let mut payload = TaskActivation::decode(&activations[1].activation as &[u8]).unwrap();
+    payload.application = Some("hammers".into());
+    activations[1].activation = payload.encode_to_vec();
+    activations[1].application = "hammers".into();
+
+    store.store(activations).await.unwrap();
+
+    let service = TaskbrokerServer { store };
+    // Request a task from an application without any activations.
+    let request = SetTaskStatusRequest {
+        id: "id_0".to_string(),
+        status: 5, // Complete
+        fetch_next_task: Some(FetchNextTask {
+            application: Some("no-matches".into()),
+            namespace: None,
+        }),
+    };
+    let response = service.set_task_status(Request::new(request)).await;
+    assert!(response.is_err());
+    let e = response.unwrap_err();
+    assert_eq!(e.code(), Code::NotFound);
+}
+
+#[tokio::test]
+#[allow(deprecated)]
+async fn test_set_task_status_with_namespace_requires_application() {
+    let store = create_test_store().await;
+    let activations = make_activations(2);
+    let namespace = activations[0].namespace.clone();
+
+    store.store(activations).await.unwrap();
+
+    let service = TaskbrokerServer { store };
+    let request = SetTaskStatusRequest {
+        id: "id_0".to_string(),
+        status: 5, // Complete
+        fetch_next_task: Some(FetchNextTask {
+            application: None,
+            namespace: Some(namespace),
+        }),
+    };
+    let response = service.set_task_status(Request::new(request)).await;
+    assert!(response.is_err());
+
+    let resp = response.unwrap_err();
+    assert_eq!(
+        resp.code(),
+        Code::NotFound,
+        "No task found as namespace without filter is invalid."
+    );
+}

src/kafka/deserialize_activation.rs

@@ -30,6 +30,7 @@ pub fn new(
         };
 
         let activation = TaskActivation::decode(payload)?;
+        let application = activation.application.unwrap_or("".into());
         let namespace = activation.namespace.clone();
         let taskname = activation.taskname.clone();
 
@@ -91,6 +92,7 @@ pub fn new(
             expires_at,
             delay_until,
             at_most_once,
+            application,
             namespace,
             taskname,
             on_attempts_exceeded,
@@ -123,6 +125,7 @@ mod tests {
         #[allow(deprecated)]
         let activation = TaskActivation {
             id: "id_0".into(),
+            application: Some("sentry".to_string()),
             namespace: generate_unique_namespace(),
             taskname: "taskname".into(),
             parameters: "{}".into(),
@@ -167,6 +170,7 @@ mod tests {
         #[allow(deprecated)]
         let activation = TaskActivation {
             id: "id_0".into(),
+            application: Some("sentry".to_string()),
             namespace: generate_unique_namespace(),
             taskname: "taskname".into(),
             parameters: "{}".into(),
@@ -212,6 +216,7 @@ mod tests {
         #[allow(deprecated)]
         let activation = TaskActivation {
             id: "id_0".into(),
+            application: Some("sentry".to_string()),
             namespace: generate_unique_namespace(),
             taskname: "taskname".into(),
             parameters: "{}".into(),
@@ -257,6 +262,7 @@ mod tests {
         #[allow(deprecated)]
         let activation = TaskActivation {
             id: "id_0".into(),
+            application: Some("sentry".to_string()),
             namespace: generate_unique_namespace(),
             taskname: "taskname".into(),
             parameters: "{}".into(),
@@ -303,6 +309,7 @@ mod tests {
         #[allow(deprecated)]
         let activation = TaskActivation {
             id: "id_0".into(),
+            application: Some("sentry".to_string()),
             namespace: generate_unique_namespace(),
             taskname: "taskname".into(),
             parameters: "{}".into(),