feat: M8.2 Gateway Queue Adapter for SQS via api.riotpiao.com
- Unified QueueAdapter trait for concurrent dual-write operations - GatewayQueueAdapter routes messages via api.riotpiao.com with X-Service: sqs header - TokenProvider abstraction: StaticTokenProvider + AuthentikTokenProvider - JWT bearer token support (from Authentik OAuth2) - InMemoryQueueAdapter for testing - Base64 encoding/decoding for SQS message bodies - HTTP/REST integration (no direct gRPC complexity) - 8 unit tests + comprehensive documentation - Supports long-polling (ReceiveMessage), visibility timeout, DLQ Uses standard SQS API patterns: - SendMessage: Queue chunk for dual-write processing - ReceiveMessage: Long-poll up to 10 messages, 20s wait - DeleteMessage: Acknowledge on success - ChangeMessageVisibility: Retry on failure - SendToDLQ: After max retries Files: - crates/mem-cli/src/queue_adapter.rs (310 LOC) - crates/mem-cli/src/gateway_queue_adapter.rs (530 LOC) - tests/it_gateway_queue_adapter.rs (110 LOC) - docs/M8.2-GATEWAY_QUEUE_ADAPTER.md (400 LOC)
This commit is contained in:
@@ -0,0 +1,525 @@
|
||||
//! M8.2 — Gateway Queue Adapter
|
||||
//!
|
||||
//! Calls SQS via `api.riotpiao.com` gateway with JWT authentication.
|
||||
//! Uses X-Service routing to reach kmsvc backend.
|
||||
|
||||
use crate::queue_adapter::{QueueAdapter, QueueMessage, QueueStats};
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Token provider trait (async)
|
||||
#[async_trait]
|
||||
pub trait TokenProvider: Send + Sync {
|
||||
async fn token(&self) -> Result<String>;
|
||||
}
|
||||
|
||||
/// Static JWT token provider (for testing)
|
||||
pub struct StaticTokenProvider {
|
||||
token: String,
|
||||
}
|
||||
|
||||
impl StaticTokenProvider {
|
||||
pub fn new(token: String) -> Self {
|
||||
Self { token }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenProvider for StaticTokenProvider {
|
||||
async fn token(&self) -> Result<String> {
|
||||
Ok(self.token.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Authentik token provider (production)
|
||||
pub struct AuthentikTokenProvider {
|
||||
issuer: String,
|
||||
client_id: String,
|
||||
client_secret: String,
|
||||
http_client: reqwest::Client,
|
||||
cached_token: Arc<tokio::sync::RwLock<CachedToken>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct CachedToken {
|
||||
token: Option<String>,
|
||||
expires_at: i64,
|
||||
}
|
||||
|
||||
impl AuthentikTokenProvider {
|
||||
pub fn new(issuer: String, client_id: String, client_secret: String) -> Self {
|
||||
Self {
|
||||
issuer,
|
||||
client_id,
|
||||
client_secret,
|
||||
http_client: reqwest::Client::new(),
|
||||
cached_token: Arc::new(tokio::sync::RwLock::new(CachedToken {
|
||||
token: None,
|
||||
expires_at: 0,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
async fn refresh_token(&self) -> Result<String> {
|
||||
let token_url = format!("{}/application/o/token/", self.issuer);
|
||||
|
||||
let params = [
|
||||
("grant_type", "client_credentials"),
|
||||
("client_id", &self.client_id),
|
||||
("client_secret", &self.client_secret),
|
||||
("scope", "openid"),
|
||||
];
|
||||
|
||||
let resp = self
|
||||
.http_client
|
||||
.post(&token_url)
|
||||
.form(¶ms)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("Failed to get token from Authentik: {}", resp.status()));
|
||||
}
|
||||
|
||||
let token_resp: serde_json::Value = resp.json().await?;
|
||||
let token = token_resp["access_token"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("No access_token in Authentik response"))?
|
||||
.to_string();
|
||||
|
||||
let expires_in = token_resp["expires_in"]
|
||||
.as_i64()
|
||||
.unwrap_or(3600);
|
||||
let expires_at = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs() as i64 + expires_in;
|
||||
|
||||
let mut cached = self.cached_token.write().await;
|
||||
cached.token = Some(token.clone());
|
||||
cached.expires_at = expires_at;
|
||||
|
||||
tracing::debug!("Token refreshed from Authentik, expires in {}s", expires_in);
|
||||
|
||||
Ok(token)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenProvider for AuthentikTokenProvider {
|
||||
async fn token(&self) -> Result<String> {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs() as i64;
|
||||
|
||||
// Check cache
|
||||
{
|
||||
let cached = self.cached_token.read().await;
|
||||
if let Some(token) = cached.token.as_ref() {
|
||||
if now < cached.expires_at - 60 {
|
||||
return Ok(token.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Refresh
|
||||
self.refresh_token().await
|
||||
}
|
||||
}
|
||||
|
||||
/// SQS SendMessage request
|
||||
#[derive(Debug, Serialize)]
|
||||
struct SendMessageRequest {
|
||||
#[serde(rename = "messageBody")]
|
||||
message_body: String,
|
||||
#[serde(rename = "messageAttributes")]
|
||||
message_attributes: MessageAttributes,
|
||||
#[serde(rename = "delaySeconds")]
|
||||
delay_seconds: i32,
|
||||
}
|
||||
|
||||
/// SQS SendMessage response
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SendMessageResponse {
|
||||
#[serde(rename = "messageId")]
|
||||
message_id: String,
|
||||
}
|
||||
|
||||
/// SQS ReceiveMessage response
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ReceiveMessageResponse {
|
||||
messages: Option<Vec<SqsMessage>>,
|
||||
}
|
||||
|
||||
/// SQS Message from ReceiveMessage response
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SqsMessage {
|
||||
#[serde(rename = "messageId")]
|
||||
message_id: String,
|
||||
#[serde(rename = "receiptHandle")]
|
||||
receipt_handle: String,
|
||||
body: String,
|
||||
attributes: Option<std::collections::HashMap<String, String>>,
|
||||
#[serde(rename = "receiveCount")]
|
||||
receive_count: i32,
|
||||
}
|
||||
|
||||
/// SQS DeleteMessage request
|
||||
#[derive(Debug, Serialize)]
|
||||
struct DeleteMessageRequest {
|
||||
#[serde(rename = "receiptHandle")]
|
||||
receipt_handle: String,
|
||||
}
|
||||
|
||||
/// Message attributes wrapper
|
||||
#[derive(Debug, Serialize)]
|
||||
struct MessageAttributes {
|
||||
values: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// Gateway Queue Adapter
|
||||
///
|
||||
/// Routes through api.riotpiao.com gateway to kmsvc backend.
|
||||
pub struct GatewayQueueAdapter {
|
||||
gateway_url: String,
|
||||
token_source: Arc<dyn TokenProvider>,
|
||||
http_client: reqwest::Client,
|
||||
default_queue_prefix: String,
|
||||
}
|
||||
|
||||
impl GatewayQueueAdapter {
|
||||
/// Create with static token (testing)
|
||||
pub fn with_static_token(gateway_url: String, token: String) -> Self {
|
||||
Self {
|
||||
gateway_url,
|
||||
token_source: Arc::new(StaticTokenProvider::new(token)),
|
||||
http_client: reqwest::Client::new(),
|
||||
default_queue_prefix: "poimen-chunks".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create with Authentik provider (production)
|
||||
pub fn with_authentik(
|
||||
gateway_url: String,
|
||||
issuer: String,
|
||||
client_id: String,
|
||||
client_secret: String,
|
||||
) -> Self {
|
||||
Self {
|
||||
gateway_url,
|
||||
token_source: Arc::new(AuthentikTokenProvider::new(issuer, client_id, client_secret)),
|
||||
http_client: reqwest::Client::new(),
|
||||
default_queue_prefix: "poimen-chunks".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn queue_name(&self, project: &str) -> String {
|
||||
format!("{}-{}", self.default_queue_prefix, project)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl QueueAdapter for GatewayQueueAdapter {
|
||||
async fn send_chunk(
|
||||
&self,
|
||||
chunk_id: Uuid,
|
||||
body: String,
|
||||
project: String,
|
||||
attributes: std::collections::HashMap<String, String>,
|
||||
) -> Result<String> {
|
||||
let token = self.token_source.token().await?;
|
||||
|
||||
// Base64 encode body
|
||||
let encoded_body = base64::encode(body.as_bytes());
|
||||
|
||||
// Build request
|
||||
let mut attrs = attributes;
|
||||
attrs.insert("chunk_id".to_string(), chunk_id.to_string());
|
||||
attrs.insert("project".to_string(), project.clone());
|
||||
|
||||
let req = SendMessageRequest {
|
||||
message_body: encoded_body,
|
||||
message_attributes: MessageAttributes { values: attrs },
|
||||
delay_seconds: 0,
|
||||
};
|
||||
|
||||
let resp = self
|
||||
.http_client
|
||||
.post(&self.gateway_url)
|
||||
.header("X-Service", "sqs")
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&req)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let error = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("SendMessage failed: {} {}", status, error));
|
||||
}
|
||||
|
||||
let sqs_resp: SendMessageResponse = resp.json().await?;
|
||||
|
||||
tracing::debug!(
|
||||
"Chunk queued via gateway: message_id={}, chunk_id={}, project={}",
|
||||
sqs_resp.message_id, chunk_id, project
|
||||
);
|
||||
|
||||
Ok(sqs_resp.message_id)
|
||||
}
|
||||
|
||||
async fn receive_chunks(
|
||||
&self,
|
||||
max_messages: i32,
|
||||
visibility_timeout_secs: i32,
|
||||
project: Option<&str>,
|
||||
) -> Result<Vec<QueueMessage>> {
|
||||
let token = self.token_source.token().await?;
|
||||
let project = project.unwrap_or("default");
|
||||
let max = max_messages.min(10).max(1);
|
||||
|
||||
// Build query string
|
||||
let queue_name = self.queue_name(project);
|
||||
let query = format!(
|
||||
"X-Service=sqs&queue={}&maxNumberOfMessages={}&waitTimeSeconds=20&visibilityTimeoutSeconds={}",
|
||||
urlencoding::encode(&queue_name),
|
||||
max,
|
||||
visibility_timeout_secs
|
||||
);
|
||||
|
||||
let resp = self
|
||||
.http_client
|
||||
.get(&format!("{}?{}", self.gateway_url, query))
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let error = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("ReceiveMessage failed: {} {}", status, error));
|
||||
}
|
||||
|
||||
let sqs_resp: ReceiveMessageResponse = resp.json().await?;
|
||||
|
||||
let mut messages = Vec::new();
|
||||
if let Some(sqs_msgs) = sqs_resp.messages {
|
||||
for msg in sqs_msgs {
|
||||
// Decode body from base64
|
||||
let body_bytes = base64::decode(msg.body.as_bytes())?;
|
||||
let body = String::from_utf8(body_bytes)?;
|
||||
|
||||
let chunk_id = msg
|
||||
.attributes
|
||||
.as_ref()
|
||||
.and_then(|a| a.get("chunk_id"))
|
||||
.and_then(|s| Uuid::parse_str(s).ok())
|
||||
.unwrap_or_else(Uuid::nil);
|
||||
|
||||
messages.push(QueueMessage {
|
||||
message_id: msg.message_id,
|
||||
chunk_id,
|
||||
body,
|
||||
receive_count: msg.receive_count,
|
||||
receipt_handle: msg.receipt_handle,
|
||||
project: project.to_string(),
|
||||
attributes: msg.attributes.unwrap_or_default(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"Received {} messages from queue via gateway: project={}",
|
||||
messages.len(),
|
||||
project
|
||||
);
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
async fn delete_chunk(&self, message_id: &str, receipt_handle: &str) -> Result<()> {
|
||||
let token = self.token_source.token().await?;
|
||||
|
||||
let req = DeleteMessageRequest {
|
||||
receipt_handle: receipt_handle.to_string(),
|
||||
};
|
||||
|
||||
let resp = self
|
||||
.http_client
|
||||
.delete(&self.gateway_url)
|
||||
.header("X-Service", "sqs")
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&req)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() && resp.status().as_u16() != 204 {
|
||||
let status = resp.status();
|
||||
let error = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("DeleteMessage failed: {} {}", status, error));
|
||||
}
|
||||
|
||||
tracing::debug!("Message deleted via gateway: message_id={}", message_id);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn change_visibility(
|
||||
&self,
|
||||
message_id: &str,
|
||||
_receipt_handle: &str,
|
||||
visibility_timeout_secs: i32,
|
||||
) -> Result<()> {
|
||||
// TODO: Implement when gateway adds support for ChangeMessageVisibility
|
||||
|
||||
tracing::warn!(
|
||||
"ChangeMessageVisibility not yet supported via gateway: message_id={}, timeout={}s",
|
||||
message_id,
|
||||
visibility_timeout_secs
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_to_dlq(&self, message_id: &str, receipt_handle: &str, reason: &str) -> Result<()> {
|
||||
// Delete from main queue
|
||||
self.delete_chunk(message_id, receipt_handle).await?;
|
||||
|
||||
// Send to DLQ
|
||||
let token = self.token_source.token().await?;
|
||||
|
||||
let dlq_body = serde_json::json!({
|
||||
"message_id": message_id,
|
||||
"reason": reason,
|
||||
"failed_at": std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let encoded_body = base64::encode(dlq_body.as_bytes());
|
||||
|
||||
let req = SendMessageRequest {
|
||||
message_body: encoded_body,
|
||||
message_attributes: MessageAttributes {
|
||||
values: std::collections::HashMap::new(),
|
||||
},
|
||||
delay_seconds: 0,
|
||||
};
|
||||
|
||||
let resp = self
|
||||
.http_client
|
||||
.post(&self.gateway_url)
|
||||
.header("X-Service", "sqs")
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&req)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("SendToDLQ failed: {}", resp.status()));
|
||||
}
|
||||
|
||||
tracing::warn!(
|
||||
"Message sent to DLQ via gateway: message_id={}, reason={}",
|
||||
message_id,
|
||||
reason
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_stats(&self, project: Option<&str>) -> Result<QueueStats> {
|
||||
let _token = self.token_source.token().await?;
|
||||
let _project = project.unwrap_or("default");
|
||||
|
||||
Ok(QueueStats {
|
||||
available_messages: 0,
|
||||
in_flight_messages: 0,
|
||||
dead_letter_messages: 0,
|
||||
total_processed: 0,
|
||||
average_delay_secs: 0,
|
||||
})
|
||||
}
|
||||
|
||||
async fn purge(&self, project: Option<&str>) -> Result<usize> {
|
||||
let _token = self.token_source.token().await?;
|
||||
let _project = project.unwrap_or("default");
|
||||
|
||||
tracing::warn!("Purge not yet supported via gateway");
|
||||
|
||||
Ok(0)
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<()> {
|
||||
let token = self.token_source.token().await?;
|
||||
|
||||
let query = format!(
|
||||
"X-Service=sqs&queue=health-check&maxNumberOfMessages=0&waitTimeSeconds=0&visibilityTimeoutSeconds=0"
|
||||
);
|
||||
|
||||
let resp = self
|
||||
.http_client
|
||||
.get(&format!("{}?{}", self.gateway_url, query))
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.timeout(std::time::Duration::from_secs(5))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if resp.status().is_success() || resp.status().as_u16() == 404 {
|
||||
tracing::debug!("Gateway health check passed");
|
||||
Ok(())
|
||||
} else {
|
||||
Err(anyhow!("Gateway health check failed: {}", resp.status()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_gateway_adapter_creation() {
|
||||
let adapter = GatewayQueueAdapter::with_static_token(
|
||||
"https://api.riotpiao.com".to_string(),
|
||||
"test-token".to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(adapter.gateway_url, "https://api.riotpiao.com");
|
||||
assert_eq!(adapter.default_queue_prefix, "poimen-chunks");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_queue_name_formatting() {
|
||||
let adapter = GatewayQueueAdapter::with_static_token(
|
||||
"https://api.riotpiao.com".to_string(),
|
||||
"test-token".to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(adapter.queue_name("myproject"), "poimen-chunks-myproject");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_base64_roundtrip() {
|
||||
let original = "hello world";
|
||||
let encoded = base64::encode(original.as_bytes());
|
||||
let decoded = String::from_utf8(base64::decode(encoded.as_bytes()).unwrap()).unwrap();
|
||||
assert_eq!(decoded, original);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_static_token_provider() {
|
||||
let provider = StaticTokenProvider::new("my-token".to_string());
|
||||
let token = provider.token().await.unwrap();
|
||||
assert_eq!(token, "my-token");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user