mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge remote-tracking branch 'upstream/aether-rust-pioneer' into fix-management-token-oauth-jsonb
# Conflicts: # crates/aether-data/src/lifecycle/bootstrap/postgres.rs # crates/aether-data/src/lifecycle/migrate/tests.rs
This commit is contained in:
@@ -9,6 +9,7 @@ const OAUTH_ACCOUNT_BLOCK_PREFIX: &str = "[ACCOUNT_BLOCK] ";
|
||||
const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] ";
|
||||
const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] ";
|
||||
const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] ";
|
||||
const CODEX_SPARK_LIMIT_NAME: &str = "GPT-5.3-Codex-Spark";
|
||||
|
||||
pub fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
|
||||
config
|
||||
@@ -214,6 +215,29 @@ fn codex_write_window(
|
||||
if let Some(value) = source.get("window_minutes").and_then(coerce_json_u64) {
|
||||
target.insert(format!("{target_prefix}_window_minutes"), json!(value));
|
||||
}
|
||||
if let Some(value) = source
|
||||
.get("limit_window_seconds")
|
||||
.and_then(coerce_json_u64)
|
||||
.map(|seconds| seconds / 60)
|
||||
{
|
||||
target.insert(format!("{target_prefix}_window_minutes"), json!(value));
|
||||
}
|
||||
}
|
||||
|
||||
fn codex_find_spark_rate_limit(
|
||||
root: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Option<&serde_json::Map<String, serde_json::Value>> {
|
||||
root.get("additional_rate_limits")
|
||||
.and_then(serde_json::Value::as_array)?
|
||||
.iter()
|
||||
.filter_map(serde_json::Value::as_object)
|
||||
.find(|item| {
|
||||
item.get("limit_name")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.is_some_and(|name| name.trim() == CODEX_SPARK_LIMIT_NAME)
|
||||
})?
|
||||
.get("rate_limit")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
}
|
||||
|
||||
pub fn parse_codex_wham_usage_response(
|
||||
@@ -256,6 +280,21 @@ pub fn parse_codex_wham_usage_response(
|
||||
codex_write_window(&mut result, &primary_window, "primary");
|
||||
}
|
||||
|
||||
if let Some(spark_rate_limit) = codex_find_spark_rate_limit(root) {
|
||||
if let Some(primary_window) = spark_rate_limit
|
||||
.get("primary_window")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
{
|
||||
codex_write_window(&mut result, primary_window, "spark_primary");
|
||||
}
|
||||
if let Some(secondary_window) = spark_rate_limit
|
||||
.get("secondary_window")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
{
|
||||
codex_write_window(&mut result, secondary_window, "spark_secondary");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(credits) = root.get("credits").and_then(serde_json::Value::as_object) {
|
||||
if let Some(value) = credits.get("has_credits").and_then(coerce_json_bool) {
|
||||
result.insert("has_credits".to_string(), json!(value));
|
||||
@@ -874,8 +913,9 @@ pub fn parse_chatgpt_web_conversation_init_response(
|
||||
mod tests {
|
||||
use super::{
|
||||
codex_build_invalid_state, codex_runtime_invalid_reason,
|
||||
parse_chatgpt_web_conversation_init_response, OAUTH_ACCOUNT_BLOCK_PREFIX,
|
||||
OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
|
||||
parse_chatgpt_web_conversation_init_response, parse_codex_wham_usage_response,
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
|
||||
OAUTH_REQUEST_FAILED_PREFIX,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::json;
|
||||
@@ -970,6 +1010,63 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_codex_spark_quota_from_additional_rate_limits() {
|
||||
let parsed = parse_codex_wham_usage_response(
|
||||
&json!({
|
||||
"plan_type": "plus",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 25.0,
|
||||
"reset_after_seconds": 604800,
|
||||
"reset_at": 1_900_000_000u64
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 10.0,
|
||||
"reset_after_seconds": 18000,
|
||||
"reset_at": 1_800_000_000u64
|
||||
}
|
||||
},
|
||||
"additional_rate_limits": [{
|
||||
"limit_name": "GPT-5.3-Codex-Spark",
|
||||
"metered_feature": "codex_bengalfox",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 40.0,
|
||||
"limit_window_seconds": 18000,
|
||||
"reset_after_seconds": 9000,
|
||||
"reset_at": 1_780_000_000u64
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 5.0,
|
||||
"limit_window_seconds": 604800,
|
||||
"reset_after_seconds": 300000,
|
||||
"reset_at": 1_790_000_000u64
|
||||
}
|
||||
}
|
||||
}]
|
||||
}),
|
||||
1_777_000_000,
|
||||
)
|
||||
.expect("codex wham usage should parse");
|
||||
|
||||
assert_eq!(parsed.get("primary_used_percent"), Some(&json!(10.0)));
|
||||
assert_eq!(parsed.get("secondary_used_percent"), Some(&json!(25.0)));
|
||||
assert_eq!(parsed.get("spark_primary_used_percent"), Some(&json!(40.0)));
|
||||
assert_eq!(
|
||||
parsed.get("spark_primary_window_minutes"),
|
||||
Some(&json!(300u64))
|
||||
);
|
||||
assert_eq!(
|
||||
parsed.get("spark_secondary_used_percent"),
|
||||
Some(&json!(5.0))
|
||||
);
|
||||
assert_eq!(
|
||||
parsed.get("spark_secondary_window_minutes"),
|
||||
Some(&json!(10_080u64))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_chatgpt_web_image_quota_from_conversation_init() {
|
||||
let parsed = parse_chatgpt_web_conversation_init_response(
|
||||
|
||||
@@ -157,7 +157,7 @@ fn request_context(mapped_model: &str, upstream_is_stream: bool) -> FormatContex
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
convert_openai_chat_request_to_claude_request,
|
||||
@@ -215,4 +215,339 @@ mod tests {
|
||||
assert_eq!(converted["messages"][0]["role"], "user");
|
||||
assert_eq!(converted["messages"][0]["content"], "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_normalizer_preserves_multiple_claude_tool_results() {
|
||||
let body = json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_1",
|
||||
"name": "lookup",
|
||||
"input": {"query": "alpha"}
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_2",
|
||||
"name": "lookup",
|
||||
"input": {"query": "beta"}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"content": "alpha result"
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_2",
|
||||
"content": [{"type": "text", "text": "beta result"}]
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"max_tokens": 128,
|
||||
});
|
||||
|
||||
let converted =
|
||||
normalize_claude_request_to_openai_chat_request(&body).expect("openai chat request");
|
||||
let messages = converted["messages"].as_array().expect("messages");
|
||||
|
||||
assert_eq!(messages.len(), 3);
|
||||
assert_eq!(messages[0]["role"], "assistant");
|
||||
assert_eq!(messages[0]["tool_calls"].as_array().unwrap().len(), 2);
|
||||
assert_eq!(messages[0]["tool_calls"][0]["id"], "toolu_1");
|
||||
assert_eq!(messages[0]["tool_calls"][1]["id"], "toolu_2");
|
||||
assert_eq!(messages[1]["role"], "tool");
|
||||
assert_eq!(messages[1]["tool_call_id"], "toolu_1");
|
||||
assert_eq!(messages[1]["content"], "alpha result");
|
||||
assert_eq!(messages[2]["role"], "tool");
|
||||
assert_eq!(messages[2]["tool_call_id"], "toolu_2");
|
||||
assert_eq!(messages[2]["content"], "beta result");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_normalizer_preserves_claude_tool_result_order_around_text() {
|
||||
let body = json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "before"},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"content": "first"
|
||||
},
|
||||
{"type": "text", "text": "between"},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_2",
|
||||
"content": "second"
|
||||
}
|
||||
]
|
||||
}],
|
||||
"max_tokens": 128,
|
||||
});
|
||||
|
||||
let converted =
|
||||
normalize_claude_request_to_openai_chat_request(&body).expect("openai chat request");
|
||||
let messages = converted["messages"].as_array().expect("messages");
|
||||
|
||||
assert_eq!(messages.len(), 4);
|
||||
assert_eq!(messages[0]["role"], "user");
|
||||
assert_eq!(messages[0]["content"], "before");
|
||||
assert_eq!(messages[1]["role"], "tool");
|
||||
assert_eq!(messages[1]["tool_call_id"], "toolu_1");
|
||||
assert_eq!(messages[1]["content"], "first");
|
||||
assert_eq!(messages[2]["role"], "user");
|
||||
assert_eq!(messages[2]["content"], "between");
|
||||
assert_eq!(messages[3]["role"], "tool");
|
||||
assert_eq!(messages[3]["tool_call_id"], "toolu_2");
|
||||
assert_eq!(messages[3]["content"], "second");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_normalizer_marks_claude_error_tool_result_string_and_object_content() {
|
||||
let object_result = json!({"code": "ENOENT", "message": "missing"});
|
||||
let body = json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_error_string",
|
||||
"content": "lookup failed",
|
||||
"is_error": true
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_error_empty",
|
||||
"content": "",
|
||||
"is_error": true
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_error_object",
|
||||
"content": object_result,
|
||||
"is_error": true
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_ok",
|
||||
"content": "still ok"
|
||||
}
|
||||
]
|
||||
}],
|
||||
"max_tokens": 128,
|
||||
});
|
||||
|
||||
let converted =
|
||||
normalize_claude_request_to_openai_chat_request(&body).expect("openai chat request");
|
||||
let messages = converted["messages"].as_array().expect("messages");
|
||||
|
||||
assert_eq!(messages.len(), 4);
|
||||
assert_eq!(messages[0]["role"], "tool");
|
||||
assert_eq!(messages[0]["tool_call_id"], "toolu_error_string");
|
||||
assert_eq!(messages[0]["content"], "[tool error]\nlookup failed");
|
||||
|
||||
assert_eq!(messages[1]["role"], "tool");
|
||||
assert_eq!(messages[1]["tool_call_id"], "toolu_error_empty");
|
||||
assert_eq!(messages[1]["content"], "[tool error]");
|
||||
|
||||
assert_eq!(messages[2]["role"], "tool");
|
||||
assert_eq!(messages[2]["tool_call_id"], "toolu_error_object");
|
||||
let object_content = messages[2]["content"].as_str().expect("object content");
|
||||
let serialized_object = object_content
|
||||
.strip_prefix("[tool error]\n")
|
||||
.expect("error prefix");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(serialized_object).expect("serialized object"),
|
||||
object_result
|
||||
);
|
||||
|
||||
assert_eq!(messages[3]["role"], "tool");
|
||||
assert_eq!(messages[3]["tool_call_id"], "toolu_ok");
|
||||
assert_eq!(messages[3]["content"], "still ok");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_normalizer_marks_claude_error_tool_result_multipart_image_content() {
|
||||
let body = json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_error_image",
|
||||
"content": [
|
||||
{"type": "text", "text": "preview"},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": "aW1hZ2U="
|
||||
}
|
||||
}
|
||||
],
|
||||
"is_error": true
|
||||
}]
|
||||
}],
|
||||
"max_tokens": 128,
|
||||
});
|
||||
|
||||
let converted =
|
||||
normalize_claude_request_to_openai_chat_request(&body).expect("openai chat request");
|
||||
let messages = converted["messages"].as_array().expect("messages");
|
||||
|
||||
assert_eq!(messages.len(), 1);
|
||||
assert_eq!(messages[0]["role"], "tool");
|
||||
assert_eq!(messages[0]["tool_call_id"], "toolu_error_image");
|
||||
let content = messages[0]["content"]
|
||||
.as_array()
|
||||
.expect("multipart error content");
|
||||
assert_eq!(
|
||||
content.as_slice(),
|
||||
&[
|
||||
json!({"type": "text", "text": "[tool error]"}),
|
||||
json!({"type": "text", "text": "preview"}),
|
||||
json!({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64,aW1hZ2U="}
|
||||
}),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_normalizer_preserves_legal_openai_tool_content_for_claude_variants() {
|
||||
let anthropic_blocks = json!([
|
||||
{"type": "text", "text": "preview"},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/jpeg",
|
||||
"data": "aGVsbG8="
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "url",
|
||||
"url": "https://example.com/image.jpg"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "document",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "application/pdf",
|
||||
"data": "JVBERi0x"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "document",
|
||||
"source": {
|
||||
"type": "url",
|
||||
"url": "https://example.com/report.pdf"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "document",
|
||||
"source": {
|
||||
"type": "text",
|
||||
"media_type": "text/plain",
|
||||
"data": "document body"
|
||||
}
|
||||
}
|
||||
]);
|
||||
let object_result = json!({"answer": 42, "ok": true});
|
||||
let body = json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_object",
|
||||
"content": object_result
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_text_blocks",
|
||||
"content": [
|
||||
{"type": "text", "text": "line one"},
|
||||
{"type": "text", "text": "line two"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_anthropic_blocks",
|
||||
"content": anthropic_blocks
|
||||
}
|
||||
]
|
||||
}],
|
||||
"max_tokens": 128,
|
||||
});
|
||||
|
||||
let converted =
|
||||
normalize_claude_request_to_openai_chat_request(&body).expect("openai chat request");
|
||||
let messages = converted["messages"].as_array().expect("messages");
|
||||
|
||||
assert_eq!(messages.len(), 3);
|
||||
assert_eq!(messages[0]["role"], "tool");
|
||||
assert_eq!(messages[0]["tool_call_id"], "toolu_object");
|
||||
let object_content = messages[0]["content"].as_str().expect("object content");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(object_content).expect("serialized object"),
|
||||
object_result
|
||||
);
|
||||
|
||||
assert_eq!(messages[1]["role"], "tool");
|
||||
assert_eq!(messages[1]["tool_call_id"], "toolu_text_blocks");
|
||||
assert_eq!(messages[1]["content"], "line one\n\nline two");
|
||||
|
||||
assert_eq!(messages[2]["role"], "tool");
|
||||
assert_eq!(messages[2]["tool_call_id"], "toolu_anthropic_blocks");
|
||||
let block_content = messages[2]["content"]
|
||||
.as_array()
|
||||
.expect("multipart anthropic block content");
|
||||
assert_eq!(
|
||||
block_content.as_slice(),
|
||||
&[
|
||||
json!({"type": "text", "text": "preview"}),
|
||||
json!({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/jpeg;base64,aGVsbG8="}
|
||||
}),
|
||||
json!({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "https://example.com/image.jpg"}
|
||||
}),
|
||||
json!({
|
||||
"type": "file",
|
||||
"file": {"file_data": "data:application/pdf;base64,JVBERi0x"}
|
||||
}),
|
||||
json!({"type": "text", "text": "[File: https://example.com/report.pdf]"}),
|
||||
json!({
|
||||
"type": "text",
|
||||
"text": "[Claude tool_result document content omitted: text/plain]"
|
||||
}),
|
||||
]
|
||||
);
|
||||
let block_content_json = Value::Array(block_content.clone()).to_string();
|
||||
assert!(!block_content_json.contains("\"source\""));
|
||||
assert!(!block_content_json.contains("document body"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ use serde_json::{json, Value};
|
||||
use crate::{
|
||||
formats::context::FormatContext,
|
||||
protocol::canonical::{
|
||||
canonical_extension_object_mut, canonical_message_to_openai_chat,
|
||||
canonical_extension_object_mut, canonical_message_to_openai_chat_messages,
|
||||
canonical_response_format_to_openai, canonical_tool_choice_to_openai,
|
||||
canonical_tool_to_openai, namespace_extension_object, openai_content_text,
|
||||
openai_extensions, openai_generation_config, openai_message_content_blocks,
|
||||
@@ -148,7 +148,7 @@ pub fn to_raw(canonical: &CanonicalRequest) -> Value {
|
||||
}
|
||||
}
|
||||
for message in &canonical.messages {
|
||||
messages.push(canonical_message_to_openai_chat(message));
|
||||
messages.extend(canonical_message_to_openai_chat_messages(message));
|
||||
}
|
||||
output.insert("messages".to_string(), Value::Array(messages));
|
||||
|
||||
|
||||
@@ -9,6 +9,9 @@ pub use crate::protocol::stream::{CanonicalStreamEvent, CanonicalStreamFrame};
|
||||
|
||||
pub(crate) const OPENAI_RESPONSES_EXTENSION_NAMESPACE: &str = "openai_responses";
|
||||
pub(crate) const OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE: &str = "openai_cli";
|
||||
const AETHER_EXTENSION_NAMESPACE: &str = "aether";
|
||||
const CLAUDE_TOOL_RESULT_SOURCE_MARKER: &str = "claude_tool_result";
|
||||
const OPENAI_CHAT_TOOL_ERROR_PREFIX: &str = "[tool error]";
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
@@ -1191,6 +1194,14 @@ pub(crate) fn claude_block_to_canonical_block(block: &Value) -> Option<Canonical
|
||||
}),
|
||||
"tool_result" => {
|
||||
let content = block_object.get("content").cloned();
|
||||
let mut extensions = claude_extensions(
|
||||
block_object,
|
||||
&["type", "tool_use_id", "content", "is_error"],
|
||||
);
|
||||
extensions.insert(
|
||||
AETHER_EXTENSION_NAMESPACE.to_string(),
|
||||
json!({ "source": CLAUDE_TOOL_RESULT_SOURCE_MARKER }),
|
||||
);
|
||||
Some(CanonicalContentBlock::ToolResult {
|
||||
tool_use_id: block_object
|
||||
.get("tool_use_id")
|
||||
@@ -1204,10 +1215,7 @@ pub(crate) fn claude_block_to_canonical_block(block: &Value) -> Option<Canonical
|
||||
.get("is_error")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
extensions: claude_extensions(
|
||||
block_object,
|
||||
&["type", "tool_use_id", "content", "is_error"],
|
||||
),
|
||||
extensions,
|
||||
})
|
||||
}
|
||||
_ => Some(CanonicalContentBlock::Unknown {
|
||||
@@ -2050,7 +2058,64 @@ pub(crate) fn openai_part_to_canonical_block(part: &Value) -> Option<CanonicalCo
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn canonical_message_to_openai_chat(message: &CanonicalMessage) -> Value {
|
||||
pub(crate) fn canonical_message_to_openai_chat_messages(message: &CanonicalMessage) -> Vec<Value> {
|
||||
let mut messages = Vec::new();
|
||||
let mut pending_start = 0usize;
|
||||
let mut saw_tool_result = false;
|
||||
|
||||
for (index, block) in message.content.iter().enumerate() {
|
||||
if let CanonicalContentBlock::ToolResult { .. } = block {
|
||||
saw_tool_result = true;
|
||||
if pending_start < index {
|
||||
if let Some(message_value) = canonical_message_blocks_to_openai_chat(
|
||||
message,
|
||||
&message.content[pending_start..index],
|
||||
false,
|
||||
) {
|
||||
messages.push(message_value);
|
||||
}
|
||||
}
|
||||
messages.push(canonical_tool_result_to_openai_chat(block));
|
||||
pending_start = index + 1;
|
||||
}
|
||||
}
|
||||
|
||||
if !saw_tool_result {
|
||||
return vec![canonical_message_without_tool_results_to_openai_chat(
|
||||
message,
|
||||
)];
|
||||
}
|
||||
|
||||
if pending_start < message.content.len() {
|
||||
if let Some(message_value) = canonical_message_blocks_to_openai_chat(
|
||||
message,
|
||||
&message.content[pending_start..],
|
||||
false,
|
||||
) {
|
||||
messages.push(message_value);
|
||||
}
|
||||
}
|
||||
|
||||
messages
|
||||
}
|
||||
|
||||
fn canonical_message_without_tool_results_to_openai_chat(message: &CanonicalMessage) -> Value {
|
||||
debug_assert!(
|
||||
!message
|
||||
.content
|
||||
.iter()
|
||||
.any(|block| matches!(block, CanonicalContentBlock::ToolResult { .. })),
|
||||
"single OpenAI Chat message emission requires no ToolResult blocks; use canonical_message_to_openai_chat_messages"
|
||||
);
|
||||
canonical_message_blocks_to_openai_chat(message, &message.content, true)
|
||||
.expect("include_empty=true always emits a chat message")
|
||||
}
|
||||
|
||||
fn canonical_message_blocks_to_openai_chat(
|
||||
message: &CanonicalMessage,
|
||||
content: &[CanonicalContentBlock],
|
||||
include_empty: bool,
|
||||
) -> Option<Value> {
|
||||
let mut output = Map::new();
|
||||
output.insert(
|
||||
"role".to_string(),
|
||||
@@ -2069,7 +2134,7 @@ pub(crate) fn canonical_message_to_openai_chat(message: &CanonicalMessage) -> Va
|
||||
let mut tool_calls = Vec::new();
|
||||
let mut reasoning_segments = Vec::new();
|
||||
let mut reasoning_parts = Vec::new();
|
||||
for block in &message.content {
|
||||
for block in content {
|
||||
match block {
|
||||
CanonicalContentBlock::Thinking {
|
||||
text,
|
||||
@@ -2123,25 +2188,7 @@ pub(crate) fn canonical_message_to_openai_chat(message: &CanonicalMessage) -> Va
|
||||
"arguments": canonicalize_tool_arguments(input),
|
||||
}
|
||||
})),
|
||||
CanonicalContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content_text,
|
||||
output: result_output,
|
||||
..
|
||||
} => {
|
||||
output.insert("role".to_string(), Value::String("tool".to_string()));
|
||||
output.insert(
|
||||
"tool_call_id".to_string(),
|
||||
Value::String(tool_use_id.clone()),
|
||||
);
|
||||
output.insert(
|
||||
"content".to_string(),
|
||||
result_output
|
||||
.clone()
|
||||
.unwrap_or_else(|| Value::String(content_text.clone().unwrap_or_default())),
|
||||
);
|
||||
return Value::Object(output);
|
||||
}
|
||||
CanonicalContentBlock::ToolResult { .. } => {}
|
||||
other => {
|
||||
if let Some(part) = canonical_content_block_to_openai_part(other) {
|
||||
content_parts.push(part);
|
||||
@@ -2149,6 +2196,14 @@ pub(crate) fn canonical_message_to_openai_chat(message: &CanonicalMessage) -> Va
|
||||
}
|
||||
}
|
||||
}
|
||||
if !include_empty
|
||||
&& content_parts.is_empty()
|
||||
&& tool_calls.is_empty()
|
||||
&& reasoning_segments.is_empty()
|
||||
&& reasoning_parts.is_empty()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
output.insert(
|
||||
"content".to_string(),
|
||||
if !tool_calls.is_empty() && content_parts.is_empty() {
|
||||
@@ -2173,9 +2228,254 @@ pub(crate) fn canonical_message_to_openai_chat(message: &CanonicalMessage) -> Va
|
||||
if !reasoning_parts.is_empty() {
|
||||
output.insert("reasoning_parts".to_string(), Value::Array(reasoning_parts));
|
||||
}
|
||||
Some(Value::Object(output))
|
||||
}
|
||||
|
||||
fn canonical_tool_result_to_openai_chat(block: &CanonicalContentBlock) -> Value {
|
||||
let CanonicalContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content_text,
|
||||
output: result_output,
|
||||
is_error,
|
||||
extensions,
|
||||
..
|
||||
} = block
|
||||
else {
|
||||
unreachable!("canonical_tool_result_to_openai_chat requires ToolResult");
|
||||
};
|
||||
|
||||
let mut output = Map::new();
|
||||
output.insert("role".to_string(), Value::String("tool".to_string()));
|
||||
output.insert(
|
||||
"tool_call_id".to_string(),
|
||||
Value::String(tool_use_id.clone()),
|
||||
);
|
||||
let content = if is_claude_tool_result(extensions) {
|
||||
let content =
|
||||
openai_chat_tool_result_content(result_output.as_ref(), content_text.as_deref());
|
||||
if *is_error {
|
||||
openai_chat_tool_error_content(content)
|
||||
} else {
|
||||
content
|
||||
}
|
||||
} else {
|
||||
result_output
|
||||
.clone()
|
||||
.unwrap_or_else(|| Value::String(content_text.clone().unwrap_or_default()))
|
||||
};
|
||||
output.insert("content".to_string(), content);
|
||||
Value::Object(output)
|
||||
}
|
||||
|
||||
fn is_claude_tool_result(extensions: &BTreeMap<String, Value>) -> bool {
|
||||
extensions
|
||||
.get(AETHER_EXTENSION_NAMESPACE)
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(Value::as_str)
|
||||
== Some(CLAUDE_TOOL_RESULT_SOURCE_MARKER)
|
||||
}
|
||||
|
||||
fn openai_chat_tool_result_content(output: Option<&Value>, content_text: Option<&str>) -> Value {
|
||||
match output {
|
||||
Some(Value::String(text)) => Value::String(text.clone()),
|
||||
Some(Value::Array(parts)) => anthropic_tool_result_blocks_to_openai_chat_content(parts),
|
||||
Some(value) => Value::String(value.to_string()),
|
||||
None => Value::String(content_text.unwrap_or_default().to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_chat_tool_error_content(content: Value) -> Value {
|
||||
match content {
|
||||
Value::String(text) if text.is_empty() => {
|
||||
Value::String(OPENAI_CHAT_TOOL_ERROR_PREFIX.to_string())
|
||||
}
|
||||
Value::String(text) => Value::String(format!("{OPENAI_CHAT_TOOL_ERROR_PREFIX}\n{text}")),
|
||||
Value::Array(parts) => {
|
||||
let mut prefixed_parts = Vec::with_capacity(parts.len() + 1);
|
||||
prefixed_parts.push(openai_text_part(OPENAI_CHAT_TOOL_ERROR_PREFIX));
|
||||
prefixed_parts.extend(parts);
|
||||
Value::Array(prefixed_parts)
|
||||
}
|
||||
value => Value::String(format!("{OPENAI_CHAT_TOOL_ERROR_PREFIX}\n{value}")),
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_tool_result_blocks_to_openai_chat_content(parts: &[Value]) -> Value {
|
||||
if let Some(text) = anthropic_text_blocks_to_string(parts) {
|
||||
return Value::String(text);
|
||||
}
|
||||
|
||||
let mut has_media_part = false;
|
||||
let converted_parts = parts
|
||||
.iter()
|
||||
.map(|part| {
|
||||
let openai_part = anthropic_tool_result_block_to_openai_chat_part(part);
|
||||
if !openai_chat_part_is_text(&openai_part) {
|
||||
has_media_part = true;
|
||||
}
|
||||
openai_part
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if has_media_part {
|
||||
Value::Array(converted_parts)
|
||||
} else {
|
||||
Value::String(openai_text_parts_to_string(&converted_parts))
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_text_blocks_to_string(parts: &[Value]) -> Option<String> {
|
||||
let mut texts = Vec::with_capacity(parts.len());
|
||||
for part in parts {
|
||||
let part_object = part.as_object()?;
|
||||
if part_object.get("type").and_then(Value::as_str) != Some("text") {
|
||||
return None;
|
||||
}
|
||||
texts.push(part_object.get("text").and_then(Value::as_str)?);
|
||||
}
|
||||
Some(texts.join("\n\n"))
|
||||
}
|
||||
|
||||
fn anthropic_tool_result_block_to_openai_chat_part(part: &Value) -> Value {
|
||||
let Some(part_object) = part.as_object() else {
|
||||
return openai_text_part("[Claude tool_result non-text content omitted]");
|
||||
};
|
||||
match part_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"text" => openai_text_part(
|
||||
part_object
|
||||
.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default(),
|
||||
),
|
||||
"image" => anthropic_image_block_to_openai_chat_part(part_object).unwrap_or_else(|| {
|
||||
openai_text_part(anthropic_media_block_summary("image", part_object))
|
||||
}),
|
||||
"document" => {
|
||||
anthropic_document_block_to_openai_chat_part(part_object).unwrap_or_else(|| {
|
||||
openai_text_part(anthropic_media_block_summary("document", part_object))
|
||||
})
|
||||
}
|
||||
"file" => anthropic_document_block_to_openai_chat_part(part_object).unwrap_or_else(|| {
|
||||
openai_text_part(anthropic_media_block_summary("file", part_object))
|
||||
}),
|
||||
"" => openai_text_part("[Claude tool_result object content omitted]"),
|
||||
raw_type => openai_text_part(format!("[Claude tool_result {raw_type} content omitted]")),
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_image_block_to_openai_chat_part(block: &Map<String, Value>) -> Option<Value> {
|
||||
let source = block.get("source")?.as_object()?;
|
||||
match source
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"base64" => {
|
||||
let media_type = anthropic_source_media_type(source)?;
|
||||
let data = anthropic_source_str(source, "data")?;
|
||||
Some(json!({
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": format!("data:{media_type};base64,{data}"),
|
||||
},
|
||||
}))
|
||||
}
|
||||
"url" => {
|
||||
let url = anthropic_source_str(source, "url")?;
|
||||
Some(json!({
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": url,
|
||||
},
|
||||
}))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_document_block_to_openai_chat_part(block: &Map<String, Value>) -> Option<Value> {
|
||||
let source = block.get("source")?.as_object()?;
|
||||
match source
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"base64" => {
|
||||
let media_type = anthropic_source_media_type(source)?;
|
||||
let data = anthropic_source_str(source, "data")?;
|
||||
Some(json!({
|
||||
"type": "file",
|
||||
"file": {
|
||||
"file_data": format!("data:{media_type};base64,{data}"),
|
||||
},
|
||||
}))
|
||||
}
|
||||
"url" => {
|
||||
let url = anthropic_source_str(source, "url")?;
|
||||
Some(openai_text_part(format!("[File: {url}]")))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_media_block_summary(kind: &str, block: &Map<String, Value>) -> String {
|
||||
let media_type = block
|
||||
.get("source")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(anthropic_source_media_type);
|
||||
match media_type {
|
||||
Some(media_type) if !media_type.trim().is_empty() => {
|
||||
format!("[Claude tool_result {kind} content omitted: {media_type}]")
|
||||
}
|
||||
_ => format!("[Claude tool_result {kind} content omitted]"),
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_text_part(text: impl Into<String>) -> Value {
|
||||
json!({
|
||||
"type": "text",
|
||||
"text": text.into(),
|
||||
})
|
||||
}
|
||||
|
||||
fn openai_chat_part_is_text(part: &Value) -> bool {
|
||||
part.as_object()
|
||||
.and_then(|object| object.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
== Some("text")
|
||||
}
|
||||
|
||||
fn openai_text_parts_to_string(parts: &[Value]) -> String {
|
||||
parts
|
||||
.iter()
|
||||
.filter_map(|part| {
|
||||
part.as_object()
|
||||
.and_then(|object| object.get("text"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n")
|
||||
}
|
||||
|
||||
fn anthropic_source_media_type(source: &Map<String, Value>) -> Option<&str> {
|
||||
source
|
||||
.get("media_type")
|
||||
.or_else(|| source.get("mime_type"))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
|
||||
fn anthropic_source_str<'a>(source: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
|
||||
source
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
|
||||
pub(crate) fn canonical_content_block_to_openai_part(
|
||||
block: &CanonicalContentBlock,
|
||||
) -> Option<Value> {
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
mod types;
|
||||
|
||||
pub use types::{
|
||||
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskReadRepository,
|
||||
BackgroundTaskRepository, BackgroundTaskStatus, BackgroundTaskSummary,
|
||||
BackgroundTaskWriteRepository, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
|
||||
StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
|
||||
};
|
||||
@@ -0,0 +1,290 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(
|
||||
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
|
||||
)]
|
||||
pub enum BackgroundTaskKind {
|
||||
Scheduled,
|
||||
Daemon,
|
||||
OnDemand,
|
||||
FireAndForget,
|
||||
}
|
||||
|
||||
impl BackgroundTaskKind {
|
||||
pub fn as_database(self) -> &'static str {
|
||||
match self {
|
||||
Self::Scheduled => "scheduled",
|
||||
Self::Daemon => "daemon",
|
||||
Self::OnDemand => "on_demand",
|
||||
Self::FireAndForget => "fire_and_forget",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"scheduled" => Ok(Self::Scheduled),
|
||||
"daemon" => Ok(Self::Daemon),
|
||||
"on_demand" => Ok(Self::OnDemand),
|
||||
"fire_and_forget" => Ok(Self::FireAndForget),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported background_tasks.kind: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(
|
||||
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
|
||||
)]
|
||||
pub enum BackgroundTaskStatus {
|
||||
Queued,
|
||||
Running,
|
||||
Retrying,
|
||||
Succeeded,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Skipped,
|
||||
}
|
||||
|
||||
impl BackgroundTaskStatus {
|
||||
pub fn as_database(self) -> &'static str {
|
||||
match self {
|
||||
Self::Queued => "queued",
|
||||
Self::Running => "running",
|
||||
Self::Retrying => "retrying",
|
||||
Self::Succeeded => "succeeded",
|
||||
Self::Failed => "failed",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::Skipped => "skipped",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"queued" => Ok(Self::Queued),
|
||||
"running" => Ok(Self::Running),
|
||||
"retrying" => Ok(Self::Retrying),
|
||||
"succeeded" => Ok(Self::Succeeded),
|
||||
"failed" => Ok(Self::Failed),
|
||||
"cancelled" => Ok(Self::Cancelled),
|
||||
"skipped" => Ok(Self::Skipped),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported background_tasks.status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredBackgroundTaskRun {
|
||||
pub id: String,
|
||||
pub task_key: String,
|
||||
pub kind: BackgroundTaskKind,
|
||||
pub trigger: String,
|
||||
pub status: BackgroundTaskStatus,
|
||||
pub attempt: u32,
|
||||
pub max_attempts: u32,
|
||||
pub owner_instance: Option<String>,
|
||||
pub progress_percent: u16,
|
||||
pub progress_message: Option<String>,
|
||||
pub payload_json: Option<Value>,
|
||||
pub result_json: Option<Value>,
|
||||
pub error_message: Option<String>,
|
||||
pub cancel_requested: bool,
|
||||
pub created_by: Option<String>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub started_at_unix_secs: Option<u64>,
|
||||
pub finished_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct UpsertBackgroundTaskRun {
|
||||
pub id: String,
|
||||
pub task_key: String,
|
||||
pub kind: BackgroundTaskKind,
|
||||
pub trigger: String,
|
||||
pub status: BackgroundTaskStatus,
|
||||
pub attempt: u32,
|
||||
pub max_attempts: u32,
|
||||
pub owner_instance: Option<String>,
|
||||
pub progress_percent: u16,
|
||||
pub progress_message: Option<String>,
|
||||
pub payload_json: Option<Value>,
|
||||
pub result_json: Option<Value>,
|
||||
pub error_message: Option<String>,
|
||||
pub cancel_requested: bool,
|
||||
pub created_by: Option<String>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub started_at_unix_secs: Option<u64>,
|
||||
pub finished_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl UpsertBackgroundTaskRun {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.id.trim().is_empty()
|
||||
|| self.task_key.trim().is_empty()
|
||||
|| self.trigger.trim().is_empty()
|
||||
{
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"background task run identity is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.progress_percent > 100 {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"background task progress_percent out of range: {}",
|
||||
self.progress_percent
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn into_stored(self) -> StoredBackgroundTaskRun {
|
||||
StoredBackgroundTaskRun {
|
||||
id: self.id,
|
||||
task_key: self.task_key,
|
||||
kind: self.kind,
|
||||
trigger: self.trigger,
|
||||
status: self.status,
|
||||
attempt: self.attempt,
|
||||
max_attempts: self.max_attempts,
|
||||
owner_instance: self.owner_instance,
|
||||
progress_percent: self.progress_percent,
|
||||
progress_message: self.progress_message,
|
||||
payload_json: self.payload_json,
|
||||
result_json: self.result_json,
|
||||
error_message: self.error_message,
|
||||
cancel_requested: self.cancel_requested,
|
||||
created_by: self.created_by,
|
||||
created_at_unix_secs: self.created_at_unix_secs,
|
||||
started_at_unix_secs: self.started_at_unix_secs,
|
||||
finished_at_unix_secs: self.finished_at_unix_secs,
|
||||
updated_at_unix_secs: self.updated_at_unix_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredBackgroundTaskEvent {
|
||||
pub id: String,
|
||||
pub run_id: String,
|
||||
pub event_type: String,
|
||||
pub message: String,
|
||||
pub payload_json: Option<Value>,
|
||||
pub created_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct UpsertBackgroundTaskEvent {
|
||||
pub id: String,
|
||||
pub run_id: String,
|
||||
pub event_type: String,
|
||||
pub message: String,
|
||||
pub payload_json: Option<Value>,
|
||||
pub created_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl UpsertBackgroundTaskEvent {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.id.trim().is_empty()
|
||||
|| self.run_id.trim().is_empty()
|
||||
|| self.event_type.trim().is_empty()
|
||||
|| self.message.trim().is_empty()
|
||||
{
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"background task event identity is empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn into_stored(self) -> StoredBackgroundTaskEvent {
|
||||
StoredBackgroundTaskEvent {
|
||||
id: self.id,
|
||||
run_id: self.run_id,
|
||||
event_type: self.event_type,
|
||||
message: self.message,
|
||||
payload_json: self.payload_json,
|
||||
created_at_unix_secs: self.created_at_unix_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BackgroundTaskListQuery {
|
||||
pub task_key_substring: Option<String>,
|
||||
pub kind: Option<BackgroundTaskKind>,
|
||||
pub status: Option<BackgroundTaskStatus>,
|
||||
pub trigger: Option<String>,
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredBackgroundTaskRunPage {
|
||||
pub items: Vec<StoredBackgroundTaskRun>,
|
||||
pub total: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BackgroundTaskSummary {
|
||||
pub total: u64,
|
||||
pub running_count: u64,
|
||||
pub by_status: BTreeMap<String, u64>,
|
||||
pub by_kind: BTreeMap<String, u64>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait BackgroundTaskReadRepository: Send + Sync {
|
||||
async fn find_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, crate::DataLayerError>;
|
||||
|
||||
async fn list_runs(
|
||||
&self,
|
||||
query: &BackgroundTaskListQuery,
|
||||
) -> Result<StoredBackgroundTaskRunPage, crate::DataLayerError>;
|
||||
|
||||
async fn list_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredBackgroundTaskEvent>, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_runs(&self) -> Result<BackgroundTaskSummary, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait BackgroundTaskWriteRepository: Send + Sync {
|
||||
async fn upsert_run(
|
||||
&self,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Result<StoredBackgroundTaskRun, crate::DataLayerError>;
|
||||
|
||||
async fn request_cancel(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn upsert_event(
|
||||
&self,
|
||||
event: UpsertBackgroundTaskEvent,
|
||||
) -> Result<StoredBackgroundTaskEvent, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait BackgroundTaskRepository:
|
||||
BackgroundTaskReadRepository + BackgroundTaskWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> BackgroundTaskRepository for T where
|
||||
T: BackgroundTaskReadRepository + BackgroundTaskWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod background_tasks;
|
||||
pub mod billing;
|
||||
pub mod candidate_selection;
|
||||
pub mod candidates;
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
CREATE TABLE IF NOT EXISTS background_task_runs (
|
||||
id VARCHAR(64) PRIMARY KEY,
|
||||
task_key VARCHAR(200) NOT NULL,
|
||||
kind VARCHAR(32) NOT NULL,
|
||||
`trigger` VARCHAR(64) NOT NULL,
|
||||
status VARCHAR(32) NOT NULL,
|
||||
attempt INT NOT NULL DEFAULT 0,
|
||||
max_attempts INT NOT NULL DEFAULT 0,
|
||||
owner_instance VARCHAR(200),
|
||||
progress_percent INT NOT NULL DEFAULT 0,
|
||||
progress_message TEXT,
|
||||
payload_json JSON,
|
||||
result_json JSON,
|
||||
error_message TEXT,
|
||||
cancel_requested TINYINT(1) NOT NULL DEFAULT 0,
|
||||
created_by VARCHAR(200),
|
||||
created_at_unix_secs BIGINT NOT NULL,
|
||||
started_at_unix_secs BIGINT NULL,
|
||||
finished_at_unix_secs BIGINT NULL,
|
||||
updated_at_unix_secs BIGINT NOT NULL,
|
||||
INDEX idx_background_task_runs_task_key (task_key),
|
||||
INDEX idx_background_task_runs_status (status),
|
||||
INDEX idx_background_task_runs_kind (kind),
|
||||
INDEX idx_background_task_runs_created_at (created_at_unix_secs)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS background_task_events (
|
||||
id VARCHAR(64) PRIMARY KEY,
|
||||
run_id VARCHAR(64) NOT NULL,
|
||||
event_type VARCHAR(64) NOT NULL,
|
||||
message TEXT NOT NULL,
|
||||
payload_json JSON,
|
||||
created_at_unix_secs BIGINT NOT NULL,
|
||||
INDEX idx_background_task_events_run_id (run_id, created_at_unix_secs),
|
||||
CONSTRAINT fk_background_task_events_run
|
||||
FOREIGN KEY (run_id) REFERENCES background_task_runs(id) ON DELETE CASCADE
|
||||
);
|
||||
@@ -0,0 +1,42 @@
|
||||
CREATE TABLE IF NOT EXISTS background_task_runs (
|
||||
id character varying(64) PRIMARY KEY,
|
||||
task_key character varying(200) NOT NULL,
|
||||
kind character varying(32) NOT NULL,
|
||||
"trigger" character varying(64) NOT NULL,
|
||||
status character varying(32) NOT NULL,
|
||||
attempt integer NOT NULL DEFAULT 0,
|
||||
max_attempts integer NOT NULL DEFAULT 0,
|
||||
owner_instance character varying(200),
|
||||
progress_percent integer NOT NULL DEFAULT 0,
|
||||
progress_message text,
|
||||
payload_json jsonb,
|
||||
result_json jsonb,
|
||||
error_message text,
|
||||
cancel_requested boolean NOT NULL DEFAULT false,
|
||||
created_by character varying(200),
|
||||
created_at_unix_secs bigint NOT NULL,
|
||||
started_at_unix_secs bigint,
|
||||
finished_at_unix_secs bigint,
|
||||
updated_at_unix_secs bigint NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_task_key
|
||||
ON background_task_runs (task_key);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_status
|
||||
ON background_task_runs (status);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind
|
||||
ON background_task_runs (kind);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_created_at
|
||||
ON background_task_runs (created_at_unix_secs DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS background_task_events (
|
||||
id character varying(64) PRIMARY KEY,
|
||||
run_id character varying(64) NOT NULL REFERENCES background_task_runs(id) ON DELETE CASCADE,
|
||||
event_type character varying(64) NOT NULL,
|
||||
message text NOT NULL,
|
||||
payload_json jsonb,
|
||||
created_at_unix_secs bigint NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_events_run_id
|
||||
ON background_task_events (run_id, created_at_unix_secs ASC);
|
||||
@@ -0,0 +1,43 @@
|
||||
CREATE TABLE IF NOT EXISTS background_task_runs (
|
||||
id TEXT PRIMARY KEY,
|
||||
task_key TEXT NOT NULL,
|
||||
kind TEXT NOT NULL,
|
||||
"trigger" TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
attempt INTEGER NOT NULL DEFAULT 0,
|
||||
max_attempts INTEGER NOT NULL DEFAULT 0,
|
||||
owner_instance TEXT,
|
||||
progress_percent INTEGER NOT NULL DEFAULT 0,
|
||||
progress_message TEXT,
|
||||
payload_json TEXT,
|
||||
result_json TEXT,
|
||||
error_message TEXT,
|
||||
cancel_requested INTEGER NOT NULL DEFAULT 0,
|
||||
created_by TEXT,
|
||||
created_at_unix_secs INTEGER NOT NULL,
|
||||
started_at_unix_secs INTEGER,
|
||||
finished_at_unix_secs INTEGER,
|
||||
updated_at_unix_secs INTEGER NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_task_key
|
||||
ON background_task_runs (task_key);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_status
|
||||
ON background_task_runs (status);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind
|
||||
ON background_task_runs (kind);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_created_at
|
||||
ON background_task_runs (created_at_unix_secs DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS background_task_events (
|
||||
id TEXT PRIMARY KEY,
|
||||
run_id TEXT NOT NULL,
|
||||
event_type TEXT NOT NULL,
|
||||
message TEXT NOT NULL,
|
||||
payload_json TEXT,
|
||||
created_at_unix_secs INTEGER NOT NULL,
|
||||
FOREIGN KEY (run_id) REFERENCES background_task_runs(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_events_run_id
|
||||
ON background_task_events (run_id, created_at_unix_secs ASC);
|
||||
@@ -0,0 +1,42 @@
|
||||
CREATE TABLE IF NOT EXISTS public.background_task_runs (
|
||||
id character varying(64) PRIMARY KEY,
|
||||
task_key character varying(200) NOT NULL,
|
||||
kind character varying(32) NOT NULL,
|
||||
"trigger" character varying(64) NOT NULL,
|
||||
status character varying(32) NOT NULL,
|
||||
attempt integer NOT NULL DEFAULT 0,
|
||||
max_attempts integer NOT NULL DEFAULT 0,
|
||||
owner_instance character varying(200),
|
||||
progress_percent integer NOT NULL DEFAULT 0,
|
||||
progress_message text,
|
||||
payload_json jsonb,
|
||||
result_json jsonb,
|
||||
error_message text,
|
||||
cancel_requested boolean NOT NULL DEFAULT false,
|
||||
created_by character varying(200),
|
||||
created_at_unix_secs bigint NOT NULL,
|
||||
started_at_unix_secs bigint,
|
||||
finished_at_unix_secs bigint,
|
||||
updated_at_unix_secs bigint NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_task_key
|
||||
ON public.background_task_runs (task_key);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_status
|
||||
ON public.background_task_runs (status);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind
|
||||
ON public.background_task_runs (kind);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_created_at
|
||||
ON public.background_task_runs (created_at_unix_secs DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.background_task_events (
|
||||
id character varying(64) PRIMARY KEY,
|
||||
run_id character varying(64) NOT NULL REFERENCES public.background_task_runs(id) ON DELETE CASCADE,
|
||||
event_type character varying(64) NOT NULL,
|
||||
message text NOT NULL,
|
||||
payload_json jsonb,
|
||||
created_at_unix_secs bigint NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_events_run_id
|
||||
ON public.background_task_events (run_id, created_at_unix_secs ASC);
|
||||
@@ -9,3 +9,4 @@
|
||||
120_stats_rollups.sql
|
||||
130_stats_cost_savings.sql
|
||||
140_proxy_node_metrics.sql
|
||||
150_background_tasks.sql
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
-- Generated by aether-data-schema from schema/logical/*.toml.
|
||||
-- Do not edit generated files directly; edit logical schema or explicit overrides instead.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS background_task_runs (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`task_key` VARCHAR(200) NOT NULL,
|
||||
`kind` VARCHAR(32) NOT NULL,
|
||||
`trigger` VARCHAR(64) NOT NULL,
|
||||
`status` VARCHAR(32) NOT NULL,
|
||||
`attempt` INT NOT NULL DEFAULT 0,
|
||||
`max_attempts` INT NOT NULL DEFAULT 0,
|
||||
`owner_instance` VARCHAR(200),
|
||||
`progress_percent` INT NOT NULL DEFAULT 0,
|
||||
`progress_message` TEXT,
|
||||
`payload_json` JSON,
|
||||
`result_json` JSON,
|
||||
`error_message` TEXT,
|
||||
`cancel_requested` TINYINT(1) NOT NULL DEFAULT 0,
|
||||
`created_by` VARCHAR(200),
|
||||
`created_at_unix_secs` BIGINT NOT NULL,
|
||||
`started_at_unix_secs` BIGINT,
|
||||
`finished_at_unix_secs` BIGINT,
|
||||
`updated_at_unix_secs` BIGINT NOT NULL,
|
||||
PRIMARY KEY (`id`),
|
||||
KEY idx_background_task_runs_task_key (`task_key`),
|
||||
KEY idx_background_task_runs_status (`status`),
|
||||
KEY idx_background_task_runs_kind (`kind`),
|
||||
KEY idx_background_task_runs_created_at (`created_at_unix_secs`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS background_task_events (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`run_id` VARCHAR(64) NOT NULL,
|
||||
`event_type` VARCHAR(64) NOT NULL,
|
||||
`message` TEXT NOT NULL,
|
||||
`payload_json` JSON,
|
||||
`created_at_unix_secs` BIGINT NOT NULL,
|
||||
PRIMARY KEY (`id`),
|
||||
KEY idx_background_task_events_run_id (`run_id`, `created_at_unix_secs`),
|
||||
CONSTRAINT fk_background_task_events_run FOREIGN KEY (`run_id`) REFERENCES background_task_runs (`id`) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
@@ -8,3 +8,4 @@
|
||||
005_wallet_billing.sql
|
||||
006_usage.sql
|
||||
007_stats.sql
|
||||
008_background_tasks.sql
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
-- Generated by aether-data-schema from schema/logical/*.toml.
|
||||
-- Do not edit generated files directly; edit logical schema or explicit overrides instead.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.background_task_runs (
|
||||
id character varying(64) NOT NULL,
|
||||
task_key character varying(200) NOT NULL,
|
||||
kind character varying(32) NOT NULL,
|
||||
trigger character varying(64) NOT NULL,
|
||||
status character varying(32) NOT NULL,
|
||||
attempt integer DEFAULT 0 NOT NULL,
|
||||
max_attempts integer DEFAULT 0 NOT NULL,
|
||||
owner_instance character varying(200),
|
||||
progress_percent integer DEFAULT 0 NOT NULL,
|
||||
progress_message text,
|
||||
payload_json jsonb,
|
||||
result_json jsonb,
|
||||
error_message text,
|
||||
cancel_requested boolean DEFAULT false NOT NULL,
|
||||
created_by character varying(200),
|
||||
created_at_unix_secs bigint NOT NULL,
|
||||
started_at_unix_secs bigint,
|
||||
finished_at_unix_secs bigint,
|
||||
updated_at_unix_secs bigint NOT NULL
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.background_task_runs ADD CONSTRAINT background_task_runs_pkey PRIMARY KEY (id);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_task_key ON public.background_task_runs USING btree (task_key);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_status ON public.background_task_runs USING btree (status);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind ON public.background_task_runs USING btree (kind);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_created_at ON public.background_task_runs USING btree (created_at_unix_secs);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.background_task_events (
|
||||
id character varying(64) NOT NULL,
|
||||
run_id character varying(64) NOT NULL,
|
||||
event_type character varying(64) NOT NULL,
|
||||
message text NOT NULL,
|
||||
payload_json jsonb,
|
||||
created_at_unix_secs bigint NOT NULL
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.background_task_events ADD CONSTRAINT background_task_events_pkey PRIMARY KEY (id);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_events_run_id ON public.background_task_events USING btree (run_id, created_at_unix_secs);
|
||||
ALTER TABLE ONLY public.background_task_events ADD CONSTRAINT fk_background_task_events_run FOREIGN KEY (run_id) REFERENCES public.background_task_runs(id) ON DELETE CASCADE;
|
||||
|
||||
@@ -8,3 +8,4 @@
|
||||
005_wallet_billing.sql
|
||||
006_usage.sql
|
||||
007_stats.sql
|
||||
008_background_tasks.sql
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
-- Generated by aether-data-schema from schema/logical/*.toml.
|
||||
-- Do not edit generated files directly; edit logical schema or explicit overrides instead.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS background_task_runs (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
task_key TEXT NOT NULL,
|
||||
kind TEXT NOT NULL,
|
||||
trigger TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
attempt INTEGER NOT NULL DEFAULT 0,
|
||||
max_attempts INTEGER NOT NULL DEFAULT 0,
|
||||
owner_instance TEXT,
|
||||
progress_percent INTEGER NOT NULL DEFAULT 0,
|
||||
progress_message TEXT,
|
||||
payload_json TEXT,
|
||||
result_json TEXT,
|
||||
error_message TEXT,
|
||||
cancel_requested INTEGER NOT NULL DEFAULT 0,
|
||||
created_by TEXT,
|
||||
created_at_unix_secs INTEGER NOT NULL,
|
||||
started_at_unix_secs INTEGER,
|
||||
finished_at_unix_secs INTEGER,
|
||||
updated_at_unix_secs INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_task_key ON background_task_runs (task_key);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_status ON background_task_runs (status);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind ON background_task_runs (kind);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_runs_created_at ON background_task_runs (created_at_unix_secs);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS background_task_events (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
run_id TEXT NOT NULL,
|
||||
event_type TEXT NOT NULL,
|
||||
message TEXT NOT NULL,
|
||||
payload_json TEXT,
|
||||
created_at_unix_secs INTEGER NOT NULL,
|
||||
CONSTRAINT fk_background_task_events_run FOREIGN KEY (run_id) REFERENCES background_task_runs (id) ON DELETE CASCADE
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_background_task_events_run_id ON background_task_events (run_id, created_at_unix_secs);
|
||||
|
||||
@@ -8,3 +8,4 @@
|
||||
005_wallet_billing.sql
|
||||
006_usage.sql
|
||||
007_stats.sql
|
||||
008_background_tasks.sql
|
||||
|
||||
159
crates/aether-data/schema/logical/008_background_tasks.toml
Normal file
159
crates/aether-data/schema/logical/008_background_tasks.toml
Normal file
@@ -0,0 +1,159 @@
|
||||
[table.background_task_runs]
|
||||
domain = "background_tasks"
|
||||
order = 10
|
||||
primary_key = ["id"]
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "task_key"
|
||||
type = "text"
|
||||
length = 200
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "kind"
|
||||
type = "text"
|
||||
length = 32
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "trigger"
|
||||
type = "text"
|
||||
length = 64
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "status"
|
||||
type = "text"
|
||||
length = 32
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "attempt"
|
||||
type = "int32"
|
||||
default = 0
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "max_attempts"
|
||||
type = "int32"
|
||||
default = 0
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "owner_instance"
|
||||
type = "text"
|
||||
length = 200
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "progress_percent"
|
||||
type = "int32"
|
||||
default = 0
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "progress_message"
|
||||
type = "text"
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "payload_json"
|
||||
type = "json"
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "result_json"
|
||||
type = "json"
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "error_message"
|
||||
type = "text"
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "cancel_requested"
|
||||
type = "bool"
|
||||
default = false
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "created_by"
|
||||
type = "text"
|
||||
length = 200
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "created_at_unix_secs"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "started_at_unix_secs"
|
||||
type = "unix_seconds"
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "finished_at_unix_secs"
|
||||
type = "unix_seconds"
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_runs.columns]]
|
||||
name = "updated_at_unix_secs"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.background_task_runs.indexes]]
|
||||
name = "idx_background_task_runs_task_key"
|
||||
columns = ["task_key"]
|
||||
|
||||
[[table.background_task_runs.indexes]]
|
||||
name = "idx_background_task_runs_status"
|
||||
columns = ["status"]
|
||||
|
||||
[[table.background_task_runs.indexes]]
|
||||
name = "idx_background_task_runs_kind"
|
||||
columns = ["kind"]
|
||||
|
||||
[[table.background_task_runs.indexes]]
|
||||
name = "idx_background_task_runs_created_at"
|
||||
columns = ["created_at_unix_secs"]
|
||||
|
||||
[table.background_task_events]
|
||||
domain = "background_tasks"
|
||||
order = 20
|
||||
primary_key = ["id"]
|
||||
|
||||
[[table.background_task_events.columns]]
|
||||
name = "id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.background_task_events.columns]]
|
||||
name = "run_id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.background_task_events.columns]]
|
||||
name = "event_type"
|
||||
type = "text"
|
||||
length = 64
|
||||
|
||||
[[table.background_task_events.columns]]
|
||||
name = "message"
|
||||
type = "text"
|
||||
|
||||
[[table.background_task_events.columns]]
|
||||
name = "payload_json"
|
||||
type = "json"
|
||||
nullable = true
|
||||
|
||||
[[table.background_task_events.columns]]
|
||||
name = "created_at_unix_secs"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.background_task_events.indexes]]
|
||||
name = "idx_background_task_events_run_id"
|
||||
columns = ["run_id", "created_at_unix_secs"]
|
||||
|
||||
[[table.background_task_events.foreign_keys]]
|
||||
name = "fk_background_task_events_run"
|
||||
columns = ["run_id"]
|
||||
references_table = "background_task_runs"
|
||||
references_columns = ["id"]
|
||||
on_delete = "cascade"
|
||||
@@ -194,6 +194,16 @@ impl DataBackends {
|
||||
None => Ok(AdminSystemPurgeSummary::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn purge_admin_request_bodies_batch(
|
||||
&self,
|
||||
batch_size: usize,
|
||||
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.purge_admin_request_bodies_batch(batch_size).await,
|
||||
None => Ok(AdminSystemPurgeSummary::default()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PostgresBackend {
|
||||
@@ -494,4 +504,15 @@ impl<'a> SqlBackendRef<'a> {
|
||||
Self::Sqlite(sqlite) => sqlite.purge_admin_system_data(target).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn purge_admin_request_bodies_batch(
|
||||
self,
|
||||
batch_size: usize,
|
||||
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => postgres.purge_admin_request_bodies_batch(batch_size).await,
|
||||
Self::Mysql(mysql) => mysql.purge_admin_request_bodies_batch(batch_size).await,
|
||||
Self::Sqlite(sqlite) => sqlite.purge_admin_request_bodies_batch(batch_size).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,9 @@ use crate::repository::auth_modules::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, MysqlAuthModuleReadRepository,
|
||||
MysqlAuthModuleRepository,
|
||||
};
|
||||
use crate::repository::background_tasks::{
|
||||
BackgroundTaskReadRepository, BackgroundTaskWriteRepository, MysqlBackgroundTaskRepository,
|
||||
};
|
||||
use crate::repository::billing::{BillingReadRepository, MysqlBillingReadRepository};
|
||||
use crate::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, MysqlMinimalCandidateSelectionReadRepository,
|
||||
@@ -123,6 +126,14 @@ impl MysqlBackend {
|
||||
Arc::new(MysqlBillingReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn background_task_read_repository(&self) -> Arc<dyn BackgroundTaskReadRepository> {
|
||||
Arc::new(MysqlBackgroundTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn background_task_write_repository(&self) -> Arc<dyn BackgroundTaskWriteRepository> {
|
||||
Arc::new(MysqlBackgroundTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
|
||||
Arc::new(MysqlRequestCandidateRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
@@ -15,6 +15,9 @@ use crate::repository::auth_modules::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, SqlxAuthModuleReadRepository,
|
||||
SqlxAuthModuleRepository,
|
||||
};
|
||||
use crate::repository::background_tasks::{
|
||||
BackgroundTaskReadRepository, BackgroundTaskWriteRepository, SqlxBackgroundTaskRepository,
|
||||
};
|
||||
use crate::repository::billing::{BillingReadRepository, SqlxBillingReadRepository};
|
||||
use crate::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, SqlxMinimalCandidateSelectionReadRepository,
|
||||
@@ -118,6 +121,14 @@ impl PostgresBackend {
|
||||
Arc::new(SqlxBillingReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn background_task_read_repository(&self) -> Arc<dyn BackgroundTaskReadRepository> {
|
||||
Arc::new(SqlxBackgroundTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn background_task_write_repository(&self) -> Arc<dyn BackgroundTaskWriteRepository> {
|
||||
Arc::new(SqlxBackgroundTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn minimal_candidate_selection_read_repository(
|
||||
&self,
|
||||
) -> Arc<dyn MinimalCandidateSelectionReadRepository> {
|
||||
|
||||
@@ -6,6 +6,7 @@ use crate::repository::announcements::AnnouncementReadRepository;
|
||||
use crate::repository::audit::AuditLogReadRepository;
|
||||
use crate::repository::auth::AuthApiKeyReadRepository;
|
||||
use crate::repository::auth_modules::AuthModuleReadRepository;
|
||||
use crate::repository::background_tasks::BackgroundTaskReadRepository;
|
||||
use crate::repository::billing::BillingReadRepository;
|
||||
use crate::repository::candidate_selection::MinimalCandidateSelectionReadRepository;
|
||||
use crate::repository::candidates::RequestCandidateReadRepository;
|
||||
@@ -27,6 +28,7 @@ pub struct DataReadRepositories {
|
||||
audit_logs: Option<Arc<dyn AuditLogReadRepository>>,
|
||||
auth_api_keys: Option<Arc<dyn AuthApiKeyReadRepository>>,
|
||||
auth_modules: Option<Arc<dyn AuthModuleReadRepository>>,
|
||||
background_tasks: Option<Arc<dyn BackgroundTaskReadRepository>>,
|
||||
billing: Option<Arc<dyn BillingReadRepository>>,
|
||||
gemini_file_mappings: Option<Arc<dyn GeminiFileMappingReadRepository>>,
|
||||
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
|
||||
@@ -50,6 +52,7 @@ impl fmt::Debug for DataReadRepositories {
|
||||
.field("has_announcements", &self.announcements.is_some())
|
||||
.field("has_audit_logs", &self.audit_logs.is_some())
|
||||
.field("has_auth_modules", &self.auth_modules.is_some())
|
||||
.field("has_background_tasks", &self.background_tasks.is_some())
|
||||
.field("has_billing", &self.billing.is_some())
|
||||
.field(
|
||||
"has_gemini_file_mappings",
|
||||
@@ -97,6 +100,10 @@ impl DataReadRepositories {
|
||||
.map(PostgresBackend::auth_module_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::auth_module_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::auth_module_read_repository)),
|
||||
background_tasks: postgres
|
||||
.map(PostgresBackend::background_task_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::background_task_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::background_task_read_repository)),
|
||||
billing: postgres
|
||||
.map(PostgresBackend::billing_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::billing_read_repository))
|
||||
@@ -177,6 +184,10 @@ impl DataReadRepositories {
|
||||
self.auth_modules.clone()
|
||||
}
|
||||
|
||||
pub fn background_tasks(&self) -> Option<Arc<dyn BackgroundTaskReadRepository>> {
|
||||
self.background_tasks.clone()
|
||||
}
|
||||
|
||||
pub fn billing(&self) -> Option<Arc<dyn BillingReadRepository>> {
|
||||
self.billing.clone()
|
||||
}
|
||||
@@ -240,6 +251,7 @@ impl DataReadRepositories {
|
||||
|| self.announcements.is_some()
|
||||
|| self.audit_logs.is_some()
|
||||
|| self.auth_modules.is_some()
|
||||
|| self.background_tasks.is_some()
|
||||
|| self.billing.is_some()
|
||||
|| self.gemini_file_mappings.is_some()
|
||||
|| self.global_models.is_some()
|
||||
|
||||
@@ -13,6 +13,9 @@ use crate::repository::auth_modules::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, SqliteAuthModuleReadRepository,
|
||||
SqliteAuthModuleRepository,
|
||||
};
|
||||
use crate::repository::background_tasks::{
|
||||
BackgroundTaskReadRepository, BackgroundTaskWriteRepository, SqliteBackgroundTaskRepository,
|
||||
};
|
||||
use crate::repository::billing::{BillingReadRepository, SqliteBillingReadRepository};
|
||||
use crate::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, SqliteMinimalCandidateSelectionReadRepository,
|
||||
@@ -124,6 +127,14 @@ impl SqliteBackend {
|
||||
Arc::new(SqliteBillingReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn background_task_read_repository(&self) -> Arc<dyn BackgroundTaskReadRepository> {
|
||||
Arc::new(SqliteBackgroundTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn background_task_write_repository(&self) -> Arc<dyn BackgroundTaskWriteRepository> {
|
||||
Arc::new(SqliteBackgroundTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
|
||||
Arc::new(SqliteRequestCandidateRepository::new(self.pool_clone()))
|
||||
}
|
||||
@@ -423,6 +434,99 @@ VALUES
|
||||
assert_eq!(admin_exists, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_system_request_bodies_purge_clears_inline_usage_body_fields() {
|
||||
let config = SqlDatabaseConfig {
|
||||
driver: DatabaseDriver::Sqlite,
|
||||
url: "sqlite::memory:".to_string(),
|
||||
pool: SqlPoolConfig {
|
||||
max_connections: 1,
|
||||
..SqlPoolConfig::default()
|
||||
},
|
||||
};
|
||||
let backend = SqliteBackend::from_config(config).expect("backend should build");
|
||||
run_sqlite_migrations(backend.pool())
|
||||
.await
|
||||
.expect("sqlite migrations should run");
|
||||
|
||||
for (column, ty) in [
|
||||
("request_body", "TEXT"),
|
||||
("response_body", "TEXT"),
|
||||
("provider_request_body", "TEXT"),
|
||||
("client_response_body", "TEXT"),
|
||||
("request_body_compressed", "BLOB"),
|
||||
("response_body_compressed", "BLOB"),
|
||||
("provider_request_body_compressed", "BLOB"),
|
||||
("client_response_body_compressed", "BLOB"),
|
||||
] {
|
||||
sqlx::query(&format!(r#"ALTER TABLE "usage" ADD COLUMN {column} {ty}"#))
|
||||
.execute(backend.pool())
|
||||
.await
|
||||
.expect("legacy body column should be added");
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO "usage" (
|
||||
request_id,
|
||||
provider_name,
|
||||
model,
|
||||
request_body,
|
||||
response_body,
|
||||
provider_request_body,
|
||||
client_response_body,
|
||||
request_body_compressed,
|
||||
response_body_compressed,
|
||||
provider_request_body_compressed,
|
||||
client_response_body_compressed,
|
||||
created_at_unix_ms
|
||||
)
|
||||
VALUES (
|
||||
'request-1',
|
||||
'openai',
|
||||
'gpt-4.1',
|
||||
'client request',
|
||||
'provider response',
|
||||
'provider request',
|
||||
'client response',
|
||||
X'01',
|
||||
X'02',
|
||||
X'03',
|
||||
X'04',
|
||||
1
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(backend.pool())
|
||||
.await
|
||||
.expect("usage row should insert");
|
||||
|
||||
let summary = backend
|
||||
.purge_admin_system_data(AdminSystemPurgeTarget::RequestBodies)
|
||||
.await
|
||||
.expect("request body purge should run");
|
||||
|
||||
assert_eq!(summary.affected.get("usage_body_fields_cleaned"), Some(&1));
|
||||
let remaining: i64 = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT COUNT(*)
|
||||
FROM "usage"
|
||||
WHERE request_body IS NOT NULL
|
||||
OR response_body IS NOT NULL
|
||||
OR provider_request_body IS NOT NULL
|
||||
OR client_response_body IS NOT NULL
|
||||
OR request_body_compressed IS NOT NULL
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
"#,
|
||||
)
|
||||
.fetch_one(backend.pool())
|
||||
.await
|
||||
.expect("remaining body count should load");
|
||||
assert_eq!(remaining, 0);
|
||||
}
|
||||
|
||||
async fn sqlite_count(pool: &sqlx::SqlitePool, table: &str) -> i64 {
|
||||
let sql = format!("SELECT COUNT(*) FROM \"{table}\"");
|
||||
sqlx::query_scalar::<_, i64>(&sql)
|
||||
|
||||
@@ -90,6 +90,20 @@ impl PostgresBackend {
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn purge_admin_request_bodies_batch(
|
||||
&self,
|
||||
batch_size: usize,
|
||||
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
|
||||
if batch_size == 0 {
|
||||
return Ok(AdminSystemPurgeSummary::default());
|
||||
}
|
||||
let mut tx = self.pool().begin().await.map_postgres_err()?;
|
||||
let mut summary = AdminSystemPurgeSummary::default();
|
||||
purge_postgres_request_bodies_batch(&mut tx, batch_size, &mut summary).await?;
|
||||
tx.commit().await.map_postgres_err()?;
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
@@ -195,6 +209,20 @@ impl MysqlBackend {
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn purge_admin_request_bodies_batch(
|
||||
&self,
|
||||
batch_size: usize,
|
||||
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
|
||||
if batch_size == 0 {
|
||||
return Ok(AdminSystemPurgeSummary::default());
|
||||
}
|
||||
let mut tx = self.pool().begin().await.map_sql_err()?;
|
||||
let mut summary = AdminSystemPurgeSummary::default();
|
||||
purge_mysql_request_bodies_batch(&mut tx, batch_size, &mut summary).await?;
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
@@ -335,6 +363,20 @@ impl SqliteBackend {
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn purge_admin_request_bodies_batch(
|
||||
&self,
|
||||
batch_size: usize,
|
||||
) -> Result<AdminSystemPurgeSummary, DataLayerError> {
|
||||
if batch_size == 0 {
|
||||
return Ok(AdminSystemPurgeSummary::default());
|
||||
}
|
||||
let mut tx = self.pool().begin().await.map_sql_err()?;
|
||||
let mut summary = AdminSystemPurgeSummary::default();
|
||||
purge_sqlite_request_bodies_batch(&mut tx, batch_size, &mut summary).await?;
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
@@ -519,6 +561,17 @@ const ADMIN_USAGE_CHILD_TABLES: &[&str] = &[
|
||||
"usage_settlement_snapshots",
|
||||
];
|
||||
|
||||
const USAGE_BODY_FIELD_COLUMNS: &[&str] = &[
|
||||
"request_body",
|
||||
"response_body",
|
||||
"provider_request_body",
|
||||
"client_response_body",
|
||||
"request_body_compressed",
|
||||
"response_body_compressed",
|
||||
"provider_request_body_compressed",
|
||||
"client_response_body_compressed",
|
||||
];
|
||||
|
||||
const ADMIN_USER_SCOPED_TABLES: &[&str] = &[
|
||||
"stats_user_daily_cost_savings_model_provider",
|
||||
"stats_user_daily_cost_savings_model",
|
||||
@@ -678,9 +731,10 @@ WHERE request_count <> 0
|
||||
}
|
||||
AdminSystemPurgeTarget::RequestBodies => {
|
||||
pg_delete_table(tx, "usage_body_blobs", summary).await?;
|
||||
pg_execute_if_table(
|
||||
pg_execute_if_table_has_columns(
|
||||
tx,
|
||||
"usage",
|
||||
USAGE_BODY_FIELD_COLUMNS,
|
||||
"usage_body_fields_cleaned",
|
||||
r#"
|
||||
UPDATE public.usage
|
||||
@@ -1038,6 +1092,33 @@ WHERE request_count <> 0
|
||||
}
|
||||
AdminSystemPurgeTarget::RequestBodies => {
|
||||
mysql_delete_table(tx, "usage_body_blobs", summary).await?;
|
||||
mysql_execute_if_table_has_columns(
|
||||
tx,
|
||||
"usage",
|
||||
USAGE_BODY_FIELD_COLUMNS,
|
||||
"usage_body_fields_cleaned",
|
||||
r#"
|
||||
UPDATE `usage`
|
||||
SET request_body = NULL,
|
||||
response_body = NULL,
|
||||
provider_request_body = NULL,
|
||||
client_response_body = NULL,
|
||||
request_body_compressed = NULL,
|
||||
response_body_compressed = NULL,
|
||||
provider_request_body_compressed = NULL,
|
||||
client_response_body_compressed = NULL
|
||||
WHERE request_body IS NOT NULL
|
||||
OR response_body IS NOT NULL
|
||||
OR provider_request_body IS NOT NULL
|
||||
OR client_response_body IS NOT NULL
|
||||
OR request_body_compressed IS NOT NULL
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
"#,
|
||||
summary,
|
||||
)
|
||||
.await?;
|
||||
mysql_execute_if_table(
|
||||
tx,
|
||||
"usage_http_audits",
|
||||
@@ -1333,6 +1414,33 @@ WHERE request_count <> 0
|
||||
}
|
||||
AdminSystemPurgeTarget::RequestBodies => {
|
||||
sqlite_delete_table(tx, "usage_body_blobs", summary).await?;
|
||||
sqlite_execute_if_table_has_columns(
|
||||
tx,
|
||||
"usage",
|
||||
USAGE_BODY_FIELD_COLUMNS,
|
||||
"usage_body_fields_cleaned",
|
||||
r#"
|
||||
UPDATE "usage"
|
||||
SET request_body = NULL,
|
||||
response_body = NULL,
|
||||
provider_request_body = NULL,
|
||||
client_response_body = NULL,
|
||||
request_body_compressed = NULL,
|
||||
response_body_compressed = NULL,
|
||||
provider_request_body_compressed = NULL,
|
||||
client_response_body_compressed = NULL
|
||||
WHERE request_body IS NOT NULL
|
||||
OR response_body IS NOT NULL
|
||||
OR provider_request_body IS NOT NULL
|
||||
OR client_response_body IS NOT NULL
|
||||
OR request_body_compressed IS NOT NULL
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
"#,
|
||||
summary,
|
||||
)
|
||||
.await?;
|
||||
sqlite_execute_if_table(
|
||||
tx,
|
||||
"usage_http_audits",
|
||||
@@ -1506,6 +1614,304 @@ WHERE user_id IN ({non_admin_users})
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn purge_postgres_request_bodies_batch(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
batch_size: usize,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
|
||||
pg_execute_batch_if_table(
|
||||
tx,
|
||||
"usage_body_blobs",
|
||||
"usage_body_blobs",
|
||||
r#"
|
||||
WITH doomed AS (
|
||||
SELECT body_ref
|
||||
FROM public.usage_body_blobs
|
||||
ORDER BY body_ref ASC
|
||||
LIMIT $1
|
||||
)
|
||||
DELETE FROM public.usage_body_blobs AS blobs
|
||||
USING doomed
|
||||
WHERE blobs.body_ref = doomed.body_ref
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
pg_execute_batch_if_table_has_columns(
|
||||
tx,
|
||||
"usage",
|
||||
USAGE_BODY_FIELD_COLUMNS,
|
||||
"usage_body_fields_cleaned",
|
||||
r#"
|
||||
WITH batch AS (
|
||||
SELECT request_id
|
||||
FROM public.usage
|
||||
WHERE request_body IS NOT NULL
|
||||
OR response_body IS NOT NULL
|
||||
OR provider_request_body IS NOT NULL
|
||||
OR client_response_body IS NOT NULL
|
||||
OR request_body_compressed IS NOT NULL
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
ORDER BY created_at_unix_ms ASC, request_id ASC
|
||||
LIMIT $1
|
||||
)
|
||||
UPDATE public.usage AS usage_rows
|
||||
SET request_body = NULL,
|
||||
response_body = NULL,
|
||||
provider_request_body = NULL,
|
||||
client_response_body = NULL,
|
||||
request_body_compressed = NULL,
|
||||
response_body_compressed = NULL,
|
||||
provider_request_body_compressed = NULL,
|
||||
client_response_body_compressed = NULL
|
||||
FROM batch
|
||||
WHERE usage_rows.request_id = batch.request_id
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
pg_execute_batch_if_table(
|
||||
tx,
|
||||
"usage_http_audits",
|
||||
"usage_http_audit_body_refs_cleaned",
|
||||
r#"
|
||||
WITH batch AS (
|
||||
SELECT request_id
|
||||
FROM public.usage_http_audits
|
||||
WHERE request_body_ref IS NOT NULL
|
||||
OR provider_request_body_ref IS NOT NULL
|
||||
OR response_body_ref IS NOT NULL
|
||||
OR client_response_body_ref IS NOT NULL
|
||||
OR request_body_state IS NOT NULL
|
||||
OR provider_request_body_state IS NOT NULL
|
||||
OR response_body_state IS NOT NULL
|
||||
OR client_response_body_state IS NOT NULL
|
||||
OR body_capture_mode <> 'none'
|
||||
ORDER BY request_id ASC
|
||||
LIMIT $1
|
||||
)
|
||||
UPDATE public.usage_http_audits AS audits
|
||||
SET request_body_ref = NULL,
|
||||
provider_request_body_ref = NULL,
|
||||
response_body_ref = NULL,
|
||||
client_response_body_ref = NULL,
|
||||
request_body_state = NULL,
|
||||
provider_request_body_state = NULL,
|
||||
response_body_state = NULL,
|
||||
client_response_body_state = NULL,
|
||||
body_capture_mode = 'none',
|
||||
updated_at = NOW()
|
||||
FROM batch
|
||||
WHERE audits.request_id = batch.request_id
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn purge_mysql_request_bodies_batch(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
batch_size: usize,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
|
||||
mysql_execute_batch_if_table(
|
||||
tx,
|
||||
"usage_body_blobs",
|
||||
"usage_body_blobs",
|
||||
r#"
|
||||
DELETE FROM usage_body_blobs
|
||||
WHERE body_ref IN (
|
||||
SELECT body_ref FROM (
|
||||
SELECT body_ref
|
||||
FROM usage_body_blobs
|
||||
ORDER BY body_ref ASC
|
||||
LIMIT ?
|
||||
) AS doomed
|
||||
)
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
mysql_execute_batch_if_table_has_columns(
|
||||
tx,
|
||||
"usage",
|
||||
USAGE_BODY_FIELD_COLUMNS,
|
||||
"usage_body_fields_cleaned",
|
||||
r#"
|
||||
UPDATE `usage`
|
||||
SET request_body = NULL,
|
||||
response_body = NULL,
|
||||
provider_request_body = NULL,
|
||||
client_response_body = NULL,
|
||||
request_body_compressed = NULL,
|
||||
response_body_compressed = NULL,
|
||||
provider_request_body_compressed = NULL,
|
||||
client_response_body_compressed = NULL
|
||||
WHERE request_id IN (
|
||||
SELECT request_id FROM (
|
||||
SELECT request_id
|
||||
FROM `usage`
|
||||
WHERE request_body IS NOT NULL
|
||||
OR response_body IS NOT NULL
|
||||
OR provider_request_body IS NOT NULL
|
||||
OR client_response_body IS NOT NULL
|
||||
OR request_body_compressed IS NOT NULL
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
ORDER BY created_at_unix_ms ASC, request_id ASC
|
||||
LIMIT ?
|
||||
) AS batch
|
||||
)
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
mysql_execute_batch_if_table(
|
||||
tx,
|
||||
"usage_http_audits",
|
||||
"usage_http_audit_body_refs_cleaned",
|
||||
r#"
|
||||
UPDATE usage_http_audits
|
||||
SET request_body_ref = NULL,
|
||||
provider_request_body_ref = NULL,
|
||||
response_body_ref = NULL,
|
||||
client_response_body_ref = NULL,
|
||||
request_body_state = NULL,
|
||||
provider_request_body_state = NULL,
|
||||
response_body_state = NULL,
|
||||
client_response_body_state = NULL,
|
||||
body_capture_mode = 'none'
|
||||
WHERE request_id IN (
|
||||
SELECT request_id FROM (
|
||||
SELECT request_id
|
||||
FROM usage_http_audits
|
||||
WHERE request_body_ref IS NOT NULL
|
||||
OR provider_request_body_ref IS NOT NULL
|
||||
OR response_body_ref IS NOT NULL
|
||||
OR client_response_body_ref IS NOT NULL
|
||||
OR request_body_state IS NOT NULL
|
||||
OR provider_request_body_state IS NOT NULL
|
||||
OR response_body_state IS NOT NULL
|
||||
OR client_response_body_state IS NOT NULL
|
||||
OR body_capture_mode <> 'none'
|
||||
ORDER BY request_id ASC
|
||||
LIMIT ?
|
||||
) AS batch
|
||||
)
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn purge_sqlite_request_bodies_batch(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
batch_size: usize,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
|
||||
sqlite_execute_batch_if_table(
|
||||
tx,
|
||||
"usage_body_blobs",
|
||||
"usage_body_blobs",
|
||||
r#"
|
||||
DELETE FROM usage_body_blobs
|
||||
WHERE body_ref IN (
|
||||
SELECT body_ref
|
||||
FROM usage_body_blobs
|
||||
ORDER BY body_ref ASC
|
||||
LIMIT ?
|
||||
)
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
sqlite_execute_batch_if_table_has_columns(
|
||||
tx,
|
||||
"usage",
|
||||
USAGE_BODY_FIELD_COLUMNS,
|
||||
"usage_body_fields_cleaned",
|
||||
r#"
|
||||
UPDATE "usage"
|
||||
SET request_body = NULL,
|
||||
response_body = NULL,
|
||||
provider_request_body = NULL,
|
||||
client_response_body = NULL,
|
||||
request_body_compressed = NULL,
|
||||
response_body_compressed = NULL,
|
||||
provider_request_body_compressed = NULL,
|
||||
client_response_body_compressed = NULL
|
||||
WHERE request_id IN (
|
||||
SELECT request_id
|
||||
FROM "usage"
|
||||
WHERE request_body IS NOT NULL
|
||||
OR response_body IS NOT NULL
|
||||
OR provider_request_body IS NOT NULL
|
||||
OR client_response_body IS NOT NULL
|
||||
OR request_body_compressed IS NOT NULL
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
ORDER BY created_at_unix_ms ASC, request_id ASC
|
||||
LIMIT ?
|
||||
)
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
sqlite_execute_batch_if_table(
|
||||
tx,
|
||||
"usage_http_audits",
|
||||
"usage_http_audit_body_refs_cleaned",
|
||||
r#"
|
||||
UPDATE usage_http_audits
|
||||
SET request_body_ref = NULL,
|
||||
provider_request_body_ref = NULL,
|
||||
response_body_ref = NULL,
|
||||
client_response_body_ref = NULL,
|
||||
request_body_state = NULL,
|
||||
provider_request_body_state = NULL,
|
||||
response_body_state = NULL,
|
||||
client_response_body_state = NULL,
|
||||
body_capture_mode = 'none'
|
||||
WHERE request_id IN (
|
||||
SELECT request_id
|
||||
FROM usage_http_audits
|
||||
WHERE request_body_ref IS NOT NULL
|
||||
OR provider_request_body_ref IS NOT NULL
|
||||
OR response_body_ref IS NOT NULL
|
||||
OR client_response_body_ref IS NOT NULL
|
||||
OR request_body_state IS NOT NULL
|
||||
OR provider_request_body_state IS NOT NULL
|
||||
OR response_body_state IS NOT NULL
|
||||
OR client_response_body_state IS NOT NULL
|
||||
OR body_capture_mode <> 'none'
|
||||
ORDER BY request_id ASC
|
||||
LIMIT ?
|
||||
)
|
||||
"#,
|
||||
summary,
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn pg_delete_table(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
table: &str,
|
||||
@@ -1586,6 +1992,69 @@ async fn pg_execute_if_table(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn pg_execute_if_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !pg_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn pg_execute_batch_if_table(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
table: &str,
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
limit: i64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !pg_table_exists(tx, checked_sql_identifier(table)?).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.bind(limit)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn pg_execute_batch_if_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
limit: i64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !pg_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.bind(limit)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn pg_table_exists(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
table: &str,
|
||||
@@ -1598,6 +2067,38 @@ async fn pg_table_exists(
|
||||
.map_postgres_err()
|
||||
}
|
||||
|
||||
async fn pg_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let table = checked_sql_identifier(table)?;
|
||||
if !pg_table_exists(tx, table).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
for column in columns {
|
||||
let column = checked_sql_identifier(column)?;
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS (
|
||||
SELECT 1
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = $1
|
||||
AND column_name = $2
|
||||
)",
|
||||
)
|
||||
.bind(table)
|
||||
.bind(column)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
if !exists {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn mysql_delete_table(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
table: &str,
|
||||
@@ -1678,6 +2179,69 @@ async fn mysql_execute_if_table(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mysql_execute_if_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !mysql_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mysql_execute_batch_if_table(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
table: &str,
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
limit: i64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !mysql_table_exists(tx, checked_sql_identifier(table)?).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.bind(limit)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mysql_execute_batch_if_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
limit: i64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !mysql_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.bind(limit)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mysql_table_exists(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
table: &str,
|
||||
@@ -1693,6 +2257,36 @@ async fn mysql_table_exists(
|
||||
Ok(total > 0)
|
||||
}
|
||||
|
||||
async fn mysql_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let table = checked_sql_identifier(table)?;
|
||||
if !mysql_table_exists(tx, table).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
for column in columns {
|
||||
let column = checked_sql_identifier(column)?;
|
||||
let total: i64 = sqlx::query_scalar(
|
||||
"SELECT COUNT(*)
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = DATABASE()
|
||||
AND table_name = ?
|
||||
AND column_name = ?",
|
||||
)
|
||||
.bind(table)
|
||||
.bind(column)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
if total == 0 {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn sqlite_delete_table(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
table: &str,
|
||||
@@ -1773,6 +2367,69 @@ async fn sqlite_execute_if_table(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn sqlite_execute_if_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !sqlite_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn sqlite_execute_batch_if_table(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
table: &str,
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
limit: i64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !sqlite_table_exists(tx, checked_sql_identifier(table)?).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.bind(limit)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn sqlite_execute_batch_if_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
key: &str,
|
||||
sql: &str,
|
||||
summary: &mut AdminSystemPurgeSummary,
|
||||
limit: i64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !sqlite_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = sqlx::query(sql)
|
||||
.bind(limit)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
summary.add(key, rows);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn sqlite_table_exists(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
table: &str,
|
||||
@@ -1787,6 +2444,31 @@ async fn sqlite_table_exists(
|
||||
Ok(total > 0)
|
||||
}
|
||||
|
||||
async fn sqlite_table_has_columns(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
table: &str,
|
||||
columns: &[&str],
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let table = checked_sql_identifier(table)?;
|
||||
if !sqlite_table_exists(tx, table).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
for column in columns {
|
||||
let column = checked_sql_identifier(column)?;
|
||||
let total: i64 =
|
||||
sqlx::query_scalar("SELECT COUNT(*) FROM pragma_table_info(?) WHERE name = ?")
|
||||
.bind(table)
|
||||
.bind(column)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
if total == 0 {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
chrono::Utc::now().timestamp().max(0) as u64
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ use super::{MysqlBackend, PostgresBackend, SqliteBackend};
|
||||
use crate::repository::announcements::AnnouncementWriteRepository;
|
||||
use crate::repository::auth::AuthApiKeyWriteRepository;
|
||||
use crate::repository::auth_modules::AuthModuleWriteRepository;
|
||||
use crate::repository::background_tasks::BackgroundTaskWriteRepository;
|
||||
use crate::repository::candidates::RequestCandidateWriteRepository;
|
||||
use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
|
||||
use crate::repository::global_models::GlobalModelWriteRepository;
|
||||
@@ -23,6 +24,7 @@ pub struct DataWriteRepositories {
|
||||
announcements: Option<Arc<dyn AnnouncementWriteRepository>>,
|
||||
auth_api_keys: Option<Arc<dyn AuthApiKeyWriteRepository>>,
|
||||
auth_modules: Option<Arc<dyn AuthModuleWriteRepository>>,
|
||||
background_tasks: Option<Arc<dyn BackgroundTaskWriteRepository>>,
|
||||
request_candidates: Option<Arc<dyn RequestCandidateWriteRepository>>,
|
||||
gemini_file_mappings: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
|
||||
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
|
||||
@@ -43,6 +45,7 @@ impl fmt::Debug for DataWriteRepositories {
|
||||
.field("has_announcements", &self.announcements.is_some())
|
||||
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
|
||||
.field("has_auth_modules", &self.auth_modules.is_some())
|
||||
.field("has_background_tasks", &self.background_tasks.is_some())
|
||||
.field("has_request_candidates", &self.request_candidates.is_some())
|
||||
.field(
|
||||
"has_gemini_file_mappings",
|
||||
@@ -81,6 +84,10 @@ impl DataWriteRepositories {
|
||||
.map(PostgresBackend::auth_module_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::auth_module_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::auth_module_write_repository)),
|
||||
background_tasks: postgres
|
||||
.map(PostgresBackend::background_task_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::background_task_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::background_task_write_repository)),
|
||||
request_candidates: postgres
|
||||
.map(PostgresBackend::request_candidate_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::request_candidate_write_repository))
|
||||
@@ -149,6 +156,10 @@ impl DataWriteRepositories {
|
||||
self.auth_modules.clone()
|
||||
}
|
||||
|
||||
pub fn background_tasks(&self) -> Option<Arc<dyn BackgroundTaskWriteRepository>> {
|
||||
self.background_tasks.clone()
|
||||
}
|
||||
|
||||
pub fn usage(&self) -> Option<Arc<dyn UsageWriteRepository>> {
|
||||
self.usage.clone()
|
||||
}
|
||||
@@ -201,6 +212,7 @@ impl DataWriteRepositories {
|
||||
self.announcements.is_some()
|
||||
|| self.auth_api_keys.is_some()
|
||||
|| self.auth_modules.is_some()
|
||||
|| self.background_tasks.is_some()
|
||||
|| self.request_candidates.is_some()
|
||||
|| self.gemini_file_mappings.is_some()
|
||||
|| self.global_models.is_some()
|
||||
|
||||
@@ -294,6 +294,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
||||
20260507000000,
|
||||
20260507120000,
|
||||
20260508000000,
|
||||
20260509000000,
|
||||
20260510000000,
|
||||
]
|
||||
);
|
||||
@@ -552,11 +553,21 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
|
||||
assert_eq!(
|
||||
mysql_versions,
|
||||
vec![20260403000000, 20260507120000, 20260508000000]
|
||||
vec![
|
||||
20260403000000,
|
||||
20260507120000,
|
||||
20260508000000,
|
||||
20260509000000
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
sqlite_versions,
|
||||
vec![20260403000000, 20260507120000, 20260508000000]
|
||||
vec![
|
||||
20260403000000,
|
||||
20260507120000,
|
||||
20260508000000,
|
||||
20260509000000
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1061,6 +1072,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
||||
20260507000000,
|
||||
20260507120000,
|
||||
20260508000000,
|
||||
20260509000000,
|
||||
20260510000000,
|
||||
]
|
||||
);
|
||||
|
||||
218
crates/aether-data/src/repository/background_tasks/memory.rs
Normal file
218
crates/aether-data/src/repository/background_tasks/memory.rs
Normal file
@@ -0,0 +1,218 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::{
|
||||
BackgroundTaskListQuery, BackgroundTaskReadRepository, BackgroundTaskStatus,
|
||||
BackgroundTaskSummary, BackgroundTaskWriteRepository, StoredBackgroundTaskEvent,
|
||||
StoredBackgroundTaskRun, StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent,
|
||||
UpsertBackgroundTaskRun,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct InMemoryBackgroundTaskIndex {
|
||||
runs: BTreeMap<String, StoredBackgroundTaskRun>,
|
||||
events_by_run: BTreeMap<String, Vec<StoredBackgroundTaskEvent>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryBackgroundTaskRepository {
|
||||
index: RwLock<InMemoryBackgroundTaskIndex>,
|
||||
}
|
||||
|
||||
impl InMemoryBackgroundTaskRepository {
|
||||
fn matches_filter(run: &StoredBackgroundTaskRun, query: &BackgroundTaskListQuery) -> bool {
|
||||
if let Some(kind) = query.kind {
|
||||
if run.kind != kind {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(status) = query.status {
|
||||
if run.status != status {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(trigger) = query.trigger.as_deref() {
|
||||
if run.trigger != trigger {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(task_key_substring) = query.task_key_substring.as_deref() {
|
||||
let needle = task_key_substring.to_ascii_lowercase();
|
||||
if !run.task_key.to_ascii_lowercase().contains(&needle) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
pub fn seed_runs<I>(runs: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredBackgroundTaskRun>,
|
||||
{
|
||||
let mut index = InMemoryBackgroundTaskIndex::default();
|
||||
for run in runs {
|
||||
index.runs.insert(run.id.clone(), run);
|
||||
}
|
||||
Self {
|
||||
index: RwLock::new(index),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskReadRepository for InMemoryBackgroundTaskRepository {
|
||||
async fn find_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
||||
Ok(self
|
||||
.index
|
||||
.read()
|
||||
.expect("background task repository lock")
|
||||
.runs
|
||||
.get(run_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_runs(
|
||||
&self,
|
||||
query: &BackgroundTaskListQuery,
|
||||
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
|
||||
let mut items = self
|
||||
.index
|
||||
.read()
|
||||
.expect("background task repository lock")
|
||||
.runs
|
||||
.values()
|
||||
.filter(|run| Self::matches_filter(run, query))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.then_with(|| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs))
|
||||
});
|
||||
|
||||
let total = items.len();
|
||||
let limit = query.limit.max(1);
|
||||
let items = items
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(limit)
|
||||
.collect::<Vec<_>>();
|
||||
Ok(StoredBackgroundTaskRunPage { items, total })
|
||||
}
|
||||
|
||||
async fn list_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
|
||||
let Some(events) = self
|
||||
.index
|
||||
.read()
|
||||
.expect("background task repository lock")
|
||||
.events_by_run
|
||||
.get(run_id)
|
||||
.cloned()
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let limit = limit.max(1);
|
||||
Ok(events.into_iter().skip(offset).take(limit).collect())
|
||||
}
|
||||
|
||||
async fn summarize_runs(&self) -> Result<BackgroundTaskSummary, DataLayerError> {
|
||||
let runs = self
|
||||
.index
|
||||
.read()
|
||||
.expect("background task repository lock")
|
||||
.runs
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let mut by_status = BTreeMap::new();
|
||||
let mut by_kind = BTreeMap::new();
|
||||
let mut running_count = 0_u64;
|
||||
for run in runs {
|
||||
*by_status
|
||||
.entry(run.status.as_database().to_string())
|
||||
.or_insert(0) += 1;
|
||||
*by_kind
|
||||
.entry(run.kind.as_database().to_string())
|
||||
.or_insert(0) += 1;
|
||||
if run.status == BackgroundTaskStatus::Running {
|
||||
running_count += 1;
|
||||
}
|
||||
}
|
||||
let total = by_status.values().copied().sum();
|
||||
Ok(BackgroundTaskSummary {
|
||||
total,
|
||||
running_count,
|
||||
by_status,
|
||||
by_kind,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskWriteRepository for InMemoryBackgroundTaskRepository {
|
||||
async fn upsert_run(
|
||||
&self,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Result<StoredBackgroundTaskRun, DataLayerError> {
|
||||
run.validate()?;
|
||||
let stored = run.into_stored();
|
||||
self.index
|
||||
.write()
|
||||
.expect("background task repository lock")
|
||||
.runs
|
||||
.insert(stored.id.clone(), stored.clone());
|
||||
Ok(stored)
|
||||
}
|
||||
|
||||
async fn request_cancel(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut guard = self.index.write().expect("background task repository lock");
|
||||
let Some(run) = guard.runs.get_mut(run_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
run.cancel_requested = true;
|
||||
run.updated_at_unix_secs = updated_at_unix_secs;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn upsert_event(
|
||||
&self,
|
||||
event: UpsertBackgroundTaskEvent,
|
||||
) -> Result<StoredBackgroundTaskEvent, DataLayerError> {
|
||||
event.validate()?;
|
||||
let stored = event.into_stored();
|
||||
let mut guard = self.index.write().expect("background task repository lock");
|
||||
let entries = guard
|
||||
.events_by_run
|
||||
.entry(stored.run_id.clone())
|
||||
.or_default();
|
||||
if let Some(position) = entries.iter().position(|value| value.id == stored.id) {
|
||||
entries[position] = stored.clone();
|
||||
} else {
|
||||
entries.push(stored.clone());
|
||||
}
|
||||
let mut seen = BTreeSet::new();
|
||||
entries.retain(|entry| seen.insert(entry.id.clone()));
|
||||
entries.sort_by(|left, right| {
|
||||
left.created_at_unix_secs
|
||||
.cmp(&right.created_at_unix_secs)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(stored)
|
||||
}
|
||||
}
|
||||
17
crates/aether-data/src/repository/background_tasks/mod.rs
Normal file
17
crates/aether-data/src/repository/background_tasks/mod.rs
Normal file
@@ -0,0 +1,17 @@
|
||||
mod memory;
|
||||
mod mysql;
|
||||
mod postgres;
|
||||
mod sqlite;
|
||||
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use aether_data_contracts::repository::background_tasks::{
|
||||
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskReadRepository,
|
||||
BackgroundTaskRepository, BackgroundTaskStatus, BackgroundTaskSummary,
|
||||
BackgroundTaskWriteRepository, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
|
||||
StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
|
||||
};
|
||||
|
||||
pub use memory::InMemoryBackgroundTaskRepository;
|
||||
pub use mysql::MysqlBackgroundTaskRepository;
|
||||
pub use postgres::SqlxBackgroundTaskRepository;
|
||||
pub use sqlite::SqliteBackgroundTaskRepository;
|
||||
449
crates/aether-data/src/repository/background_tasks/mysql.rs
Normal file
449
crates/aether-data/src/repository/background_tasks/mysql.rs
Normal file
@@ -0,0 +1,449 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
||||
|
||||
use super::{
|
||||
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskReadRepository,
|
||||
BackgroundTaskStatus, BackgroundTaskSummary, BackgroundTaskWriteRepository,
|
||||
StoredBackgroundTaskEvent, StoredBackgroundTaskRun, StoredBackgroundTaskRunPage,
|
||||
UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
|
||||
};
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const RUN_COLUMNS: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
task_key,
|
||||
kind,
|
||||
`trigger`,
|
||||
status,
|
||||
attempt,
|
||||
max_attempts,
|
||||
owner_instance,
|
||||
progress_percent,
|
||||
progress_message,
|
||||
payload_json,
|
||||
result_json,
|
||||
error_message,
|
||||
cancel_requested,
|
||||
created_by,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
updated_at_unix_secs
|
||||
FROM background_task_runs
|
||||
"#;
|
||||
|
||||
const EVENT_COLUMNS: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
run_id,
|
||||
event_type,
|
||||
message,
|
||||
payload_json,
|
||||
created_at_unix_secs
|
||||
FROM background_task_events
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MysqlBackgroundTaskRepository {
|
||||
pool: MysqlPool,
|
||||
}
|
||||
|
||||
impl MysqlBackgroundTaskRepository {
|
||||
pub fn new(pool: MysqlPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
fn apply_run_filter(builder: &mut QueryBuilder<'_, MySql>, query: &BackgroundTaskListQuery) {
|
||||
let mut has_where = false;
|
||||
if let Some(kind) = query.kind {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
has_where = true;
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder.push("kind = ").push_bind(kind.as_database());
|
||||
}
|
||||
if let Some(status) = query.status {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
has_where = true;
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder.push("status = ").push_bind(status.as_database());
|
||||
}
|
||||
if let Some(trigger) = query.trigger.as_deref() {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
has_where = true;
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder.push("`trigger` = ").push_bind(trigger.to_string());
|
||||
}
|
||||
if let Some(task_key_substring) = query.task_key_substring.as_deref() {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder.push("LOWER(task_key) LIKE ").push_bind(format!(
|
||||
"%{}%",
|
||||
task_key_substring.trim().to_ascii_lowercase()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskReadRepository for MysqlBackgroundTaskRepository {
|
||||
async fn find_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
||||
let row = sqlx::query(&format!("{RUN_COLUMNS} WHERE id = ? LIMIT 1"))
|
||||
.bind(run_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row.as_ref().map(map_run_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_runs(
|
||||
&self,
|
||||
query: &BackgroundTaskListQuery,
|
||||
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
|
||||
let limit = query.limit.max(1);
|
||||
let mut count_builder =
|
||||
QueryBuilder::<MySql>::new("SELECT COUNT(id) AS total FROM background_task_runs");
|
||||
Self::apply_run_filter(&mut count_builder, query);
|
||||
let total = count_builder
|
||||
.build_query_scalar::<i64>()
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
let mut builder = QueryBuilder::<MySql>::new(RUN_COLUMNS);
|
||||
Self::apply_run_filter(&mut builder, query);
|
||||
builder
|
||||
.push(" ORDER BY created_at_unix_secs DESC, updated_at_unix_secs DESC")
|
||||
.push(" LIMIT ")
|
||||
.push_bind(i64_from_usize(limit, "run limit")?)
|
||||
.push(" OFFSET ")
|
||||
.push_bind(i64_from_usize(query.offset, "run offset")?);
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(map_run_row)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(StoredBackgroundTaskRunPage {
|
||||
items,
|
||||
total: usize::try_from(total).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
|
||||
let limit = limit.max(1);
|
||||
let rows = sqlx::query(&format!(
|
||||
"{EVENT_COLUMNS} WHERE run_id = ? ORDER BY created_at_unix_secs ASC, id ASC LIMIT ? OFFSET ?"
|
||||
))
|
||||
.bind(run_id)
|
||||
.bind(i64_from_usize(limit, "event limit")?)
|
||||
.bind(i64_from_usize(offset, "event offset")?)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_event_row).collect()
|
||||
}
|
||||
|
||||
async fn summarize_runs(&self) -> Result<BackgroundTaskSummary, DataLayerError> {
|
||||
let total = sqlx::query_scalar::<_, i64>("SELECT COUNT(id) FROM background_task_runs")
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let running_count = sqlx::query_scalar::<_, i64>(
|
||||
"SELECT COUNT(id) FROM background_task_runs WHERE status = 'running'",
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let status_rows = sqlx::query(
|
||||
"SELECT status, COUNT(id) AS total FROM background_task_runs GROUP BY status",
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let kind_rows =
|
||||
sqlx::query("SELECT kind, COUNT(id) AS total FROM background_task_runs GROUP BY kind")
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
let mut by_status = std::collections::BTreeMap::new();
|
||||
for row in status_rows {
|
||||
let key: String = row.try_get("status").map_sql_err()?;
|
||||
let count: i64 = row.try_get("total").map_sql_err()?;
|
||||
by_status.insert(key, u64::try_from(count).unwrap_or_default());
|
||||
}
|
||||
let mut by_kind = std::collections::BTreeMap::new();
|
||||
for row in kind_rows {
|
||||
let key: String = row.try_get("kind").map_sql_err()?;
|
||||
let count: i64 = row.try_get("total").map_sql_err()?;
|
||||
by_kind.insert(key, u64::try_from(count).unwrap_or_default());
|
||||
}
|
||||
|
||||
Ok(BackgroundTaskSummary {
|
||||
total: u64::try_from(total).unwrap_or_default(),
|
||||
running_count: u64::try_from(running_count).unwrap_or_default(),
|
||||
by_status,
|
||||
by_kind,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskWriteRepository for MysqlBackgroundTaskRepository {
|
||||
async fn upsert_run(
|
||||
&self,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Result<StoredBackgroundTaskRun, DataLayerError> {
|
||||
run.validate()?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO background_task_runs (
|
||||
id,
|
||||
task_key,
|
||||
kind,
|
||||
`trigger`,
|
||||
status,
|
||||
attempt,
|
||||
max_attempts,
|
||||
owner_instance,
|
||||
progress_percent,
|
||||
progress_message,
|
||||
payload_json,
|
||||
result_json,
|
||||
error_message,
|
||||
cancel_requested,
|
||||
created_by,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
updated_at_unix_secs
|
||||
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
||||
ON DUPLICATE KEY UPDATE
|
||||
task_key = VALUES(task_key),
|
||||
kind = VALUES(kind),
|
||||
`trigger` = VALUES(`trigger`),
|
||||
status = VALUES(status),
|
||||
attempt = VALUES(attempt),
|
||||
max_attempts = VALUES(max_attempts),
|
||||
owner_instance = VALUES(owner_instance),
|
||||
progress_percent = VALUES(progress_percent),
|
||||
progress_message = VALUES(progress_message),
|
||||
payload_json = VALUES(payload_json),
|
||||
result_json = VALUES(result_json),
|
||||
error_message = VALUES(error_message),
|
||||
cancel_requested = VALUES(cancel_requested),
|
||||
created_by = VALUES(created_by),
|
||||
created_at_unix_secs = VALUES(created_at_unix_secs),
|
||||
started_at_unix_secs = VALUES(started_at_unix_secs),
|
||||
finished_at_unix_secs = VALUES(finished_at_unix_secs),
|
||||
updated_at_unix_secs = VALUES(updated_at_unix_secs)
|
||||
"#,
|
||||
)
|
||||
.bind(&run.id)
|
||||
.bind(&run.task_key)
|
||||
.bind(run.kind.as_database())
|
||||
.bind(&run.trigger)
|
||||
.bind(run.status.as_database())
|
||||
.bind(i64::from(run.attempt))
|
||||
.bind(i64::from(run.max_attempts))
|
||||
.bind(run.owner_instance.as_deref())
|
||||
.bind(i32::from(run.progress_percent))
|
||||
.bind(run.progress_message.as_deref())
|
||||
.bind(json_to_string(&run.payload_json, "payload_json")?)
|
||||
.bind(json_to_string(&run.result_json, "result_json")?)
|
||||
.bind(run.error_message.as_deref())
|
||||
.bind(run.cancel_requested)
|
||||
.bind(run.created_by.as_deref())
|
||||
.bind(u64_to_i64(
|
||||
run.created_at_unix_secs,
|
||||
"created_at_unix_secs",
|
||||
)?)
|
||||
.bind(run.started_at_unix_secs.map(|value| value as i64))
|
||||
.bind(run.finished_at_unix_secs.map(|value| value as i64))
|
||||
.bind(u64_to_i64(
|
||||
run.updated_at_unix_secs,
|
||||
"updated_at_unix_secs",
|
||||
)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
self.find_run(&run.id).await?.ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue("background task run missing after upsert".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
async fn request_cancel(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let affected = sqlx::query(
|
||||
"UPDATE background_task_runs SET cancel_requested = TRUE, updated_at_unix_secs = ? WHERE id = ?",
|
||||
)
|
||||
.bind(u64_to_i64(updated_at_unix_secs, "updated_at_unix_secs")?)
|
||||
.bind(run_id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
Ok(affected > 0)
|
||||
}
|
||||
|
||||
async fn upsert_event(
|
||||
&self,
|
||||
event: UpsertBackgroundTaskEvent,
|
||||
) -> Result<StoredBackgroundTaskEvent, DataLayerError> {
|
||||
event.validate()?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO background_task_events (
|
||||
id, run_id, event_type, message, payload_json, created_at_unix_secs
|
||||
) VALUES (?, ?, ?, ?, ?, ?)
|
||||
ON DUPLICATE KEY UPDATE
|
||||
run_id = VALUES(run_id),
|
||||
event_type = VALUES(event_type),
|
||||
message = VALUES(message),
|
||||
payload_json = VALUES(payload_json),
|
||||
created_at_unix_secs = VALUES(created_at_unix_secs)
|
||||
"#,
|
||||
)
|
||||
.bind(&event.id)
|
||||
.bind(&event.run_id)
|
||||
.bind(&event.event_type)
|
||||
.bind(&event.message)
|
||||
.bind(json_to_string(&event.payload_json, "payload_json")?)
|
||||
.bind(u64_to_i64(
|
||||
event.created_at_unix_secs,
|
||||
"created_at_unix_secs",
|
||||
)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
let row = sqlx::query(&format!("{EVENT_COLUMNS} WHERE id = ? LIMIT 1"))
|
||||
.bind(&event.id)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
map_event_row(&row)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_run_row(row: &MySqlRow) -> Result<StoredBackgroundTaskRun, DataLayerError> {
|
||||
let kind: String = row.try_get("kind").map_sql_err()?;
|
||||
let status: String = row.try_get("status").map_sql_err()?;
|
||||
let attempt: i64 = row.try_get("attempt").map_sql_err()?;
|
||||
let max_attempts: i64 = row.try_get("max_attempts").map_sql_err()?;
|
||||
let progress_percent: i32 = row.try_get("progress_percent").map_sql_err()?;
|
||||
let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?;
|
||||
let started_at_unix_secs: Option<i64> = row.try_get("started_at_unix_secs").map_sql_err()?;
|
||||
let finished_at_unix_secs: Option<i64> = row.try_get("finished_at_unix_secs").map_sql_err()?;
|
||||
let updated_at_unix_secs: i64 = row.try_get("updated_at_unix_secs").map_sql_err()?;
|
||||
|
||||
Ok(StoredBackgroundTaskRun {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
task_key: row.try_get("task_key").map_sql_err()?,
|
||||
kind: BackgroundTaskKind::from_database(&kind)?,
|
||||
trigger: row.try_get("trigger").map_sql_err()?,
|
||||
status: BackgroundTaskStatus::from_database(&status)?,
|
||||
attempt: u32::try_from(attempt).unwrap_or_default(),
|
||||
max_attempts: u32::try_from(max_attempts).unwrap_or_default(),
|
||||
owner_instance: row.try_get("owner_instance").map_sql_err()?,
|
||||
progress_percent: u16::try_from(progress_percent).unwrap_or_default(),
|
||||
progress_message: row.try_get("progress_message").map_sql_err()?,
|
||||
payload_json: parse_optional_json(
|
||||
row.try_get("payload_json").ok().flatten(),
|
||||
"payload_json",
|
||||
)?,
|
||||
result_json: parse_optional_json(row.try_get("result_json").ok().flatten(), "result_json")?,
|
||||
error_message: row.try_get("error_message").map_sql_err()?,
|
||||
cancel_requested: row.try_get("cancel_requested").map_sql_err()?,
|
||||
created_by: row.try_get("created_by").map_sql_err()?,
|
||||
created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(),
|
||||
started_at_unix_secs: started_at_unix_secs.and_then(|value| u64::try_from(value).ok()),
|
||||
finished_at_unix_secs: finished_at_unix_secs.and_then(|value| u64::try_from(value).ok()),
|
||||
updated_at_unix_secs: u64::try_from(updated_at_unix_secs).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn map_event_row(row: &MySqlRow) -> Result<StoredBackgroundTaskEvent, DataLayerError> {
|
||||
let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?;
|
||||
Ok(StoredBackgroundTaskEvent {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
run_id: row.try_get("run_id").map_sql_err()?,
|
||||
event_type: row.try_get("event_type").map_sql_err()?,
|
||||
message: row.try_get("message").map_sql_err()?,
|
||||
payload_json: parse_optional_json(
|
||||
row.try_get("payload_json").ok().flatten(),
|
||||
"payload_json",
|
||||
)?,
|
||||
created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn i64_from_usize(value: usize, label: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn u64_to_i64(value: u64, label: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn json_to_string(
|
||||
value: &Option<serde_json::Value>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<String>, DataLayerError> {
|
||||
value
|
||||
.as_ref()
|
||||
.map(|value| {
|
||||
serde_json::to_string(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"background task {field_name} is unserializable: {err}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn parse_optional_json(
|
||||
value: Option<String>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
value
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.map(|value| {
|
||||
serde_json::from_str(&value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"background task {field_name} contains invalid JSON: {err}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
426
crates/aether-data/src/repository/background_tasks/postgres.rs
Normal file
426
crates/aether-data/src/repository/background_tasks/postgres.rs
Normal file
@@ -0,0 +1,426 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||
|
||||
use super::{
|
||||
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskReadRepository,
|
||||
BackgroundTaskStatus, BackgroundTaskSummary, BackgroundTaskWriteRepository,
|
||||
StoredBackgroundTaskEvent, StoredBackgroundTaskRun, StoredBackgroundTaskRunPage,
|
||||
UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
|
||||
};
|
||||
use crate::error::SqlxResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const RUN_COLUMNS: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
task_key,
|
||||
kind,
|
||||
"trigger",
|
||||
status,
|
||||
attempt,
|
||||
max_attempts,
|
||||
owner_instance,
|
||||
progress_percent,
|
||||
progress_message,
|
||||
payload_json,
|
||||
result_json,
|
||||
error_message,
|
||||
cancel_requested,
|
||||
created_by,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
updated_at_unix_secs
|
||||
FROM background_task_runs
|
||||
"#;
|
||||
|
||||
const EVENT_COLUMNS: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
run_id,
|
||||
event_type,
|
||||
message,
|
||||
payload_json,
|
||||
created_at_unix_secs
|
||||
FROM background_task_events
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxBackgroundTaskRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxBackgroundTaskRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
fn apply_run_filter(
|
||||
builder: &mut QueryBuilder<'_, Postgres>,
|
||||
query: &BackgroundTaskListQuery,
|
||||
include_where: bool,
|
||||
) {
|
||||
let mut has_where = include_where;
|
||||
let mut push_where = |builder: &mut QueryBuilder<'_, Postgres>| {
|
||||
if has_where {
|
||||
builder.push(" AND ");
|
||||
} else {
|
||||
builder.push(" WHERE ");
|
||||
has_where = true;
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(kind) = query.kind {
|
||||
push_where(builder);
|
||||
builder.push("kind = ").push_bind(kind.as_database());
|
||||
}
|
||||
if let Some(status) = query.status {
|
||||
push_where(builder);
|
||||
builder.push("status = ").push_bind(status.as_database());
|
||||
}
|
||||
if let Some(trigger) = query.trigger.as_deref() {
|
||||
push_where(builder);
|
||||
builder
|
||||
.push("\"trigger\" = ")
|
||||
.push_bind(trigger.to_string());
|
||||
}
|
||||
if let Some(task_key_substring) = query.task_key_substring.as_deref() {
|
||||
push_where(builder);
|
||||
builder
|
||||
.push("task_key ILIKE ")
|
||||
.push_bind(format!("%{}%", task_key_substring.trim()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskReadRepository for SqlxBackgroundTaskRepository {
|
||||
async fn find_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
||||
let row = sqlx::query(&format!("{RUN_COLUMNS} WHERE id = $1 LIMIT 1"))
|
||||
.bind(run_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
row.as_ref().map(map_run_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_runs(
|
||||
&self,
|
||||
query: &BackgroundTaskListQuery,
|
||||
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
|
||||
let limit = query.limit.max(1);
|
||||
|
||||
let mut count_builder =
|
||||
QueryBuilder::<Postgres>::new("SELECT COUNT(id) AS total FROM background_task_runs");
|
||||
Self::apply_run_filter(&mut count_builder, query, false);
|
||||
let total = count_builder
|
||||
.build_query_scalar::<i64>()
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
let mut builder = QueryBuilder::<Postgres>::new(RUN_COLUMNS);
|
||||
Self::apply_run_filter(&mut builder, query, false);
|
||||
builder
|
||||
.push(" ORDER BY created_at_unix_secs DESC, updated_at_unix_secs DESC")
|
||||
.push(" LIMIT ")
|
||||
.push_bind(i64_from_usize(limit, "background task run limit")?)
|
||||
.push(" OFFSET ")
|
||||
.push_bind(i64_from_usize(query.offset, "background task run offset")?);
|
||||
let rows = builder
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(map_run_row)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
|
||||
Ok(StoredBackgroundTaskRunPage {
|
||||
items,
|
||||
total: usize::try_from(total).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
|
||||
let limit = limit.max(1);
|
||||
let rows = sqlx::query(&format!(
|
||||
"{EVENT_COLUMNS} WHERE run_id = $1 ORDER BY created_at_unix_secs ASC, id ASC LIMIT $2 OFFSET $3"
|
||||
))
|
||||
.bind(run_id)
|
||||
.bind(i64_from_usize(limit, "background task event limit")?)
|
||||
.bind(i64_from_usize(offset, "background task event offset")?)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
rows.iter().map(map_event_row).collect()
|
||||
}
|
||||
|
||||
async fn summarize_runs(&self) -> Result<BackgroundTaskSummary, DataLayerError> {
|
||||
let total = sqlx::query_scalar::<_, i64>("SELECT COUNT(id) FROM background_task_runs")
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let running_count = sqlx::query_scalar::<_, i64>(
|
||||
"SELECT COUNT(id) FROM background_task_runs WHERE status = 'running'",
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let status_rows = sqlx::query(
|
||||
"SELECT status, COUNT(id) AS total FROM background_task_runs GROUP BY status",
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let kind_rows =
|
||||
sqlx::query("SELECT kind, COUNT(id) AS total FROM background_task_runs GROUP BY kind")
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
let mut by_status = std::collections::BTreeMap::new();
|
||||
for row in status_rows {
|
||||
let key: String = row.try_get("status").map_postgres_err()?;
|
||||
let count: i64 = row.try_get("total").map_postgres_err()?;
|
||||
by_status.insert(key, u64::try_from(count).unwrap_or_default());
|
||||
}
|
||||
let mut by_kind = std::collections::BTreeMap::new();
|
||||
for row in kind_rows {
|
||||
let key: String = row.try_get("kind").map_postgres_err()?;
|
||||
let count: i64 = row.try_get("total").map_postgres_err()?;
|
||||
by_kind.insert(key, u64::try_from(count).unwrap_or_default());
|
||||
}
|
||||
|
||||
Ok(BackgroundTaskSummary {
|
||||
total: u64::try_from(total).unwrap_or_default(),
|
||||
running_count: u64::try_from(running_count).unwrap_or_default(),
|
||||
by_status,
|
||||
by_kind,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskWriteRepository for SqlxBackgroundTaskRepository {
|
||||
async fn upsert_run(
|
||||
&self,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Result<StoredBackgroundTaskRun, DataLayerError> {
|
||||
run.validate()?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO background_task_runs (
|
||||
id,
|
||||
task_key,
|
||||
kind,
|
||||
"trigger",
|
||||
status,
|
||||
attempt,
|
||||
max_attempts,
|
||||
owner_instance,
|
||||
progress_percent,
|
||||
progress_message,
|
||||
payload_json,
|
||||
result_json,
|
||||
error_message,
|
||||
cancel_requested,
|
||||
created_by,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
updated_at_unix_secs
|
||||
) VALUES (
|
||||
$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19
|
||||
)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
task_key = EXCLUDED.task_key,
|
||||
kind = EXCLUDED.kind,
|
||||
"trigger" = EXCLUDED."trigger",
|
||||
status = EXCLUDED.status,
|
||||
attempt = EXCLUDED.attempt,
|
||||
max_attempts = EXCLUDED.max_attempts,
|
||||
owner_instance = EXCLUDED.owner_instance,
|
||||
progress_percent = EXCLUDED.progress_percent,
|
||||
progress_message = EXCLUDED.progress_message,
|
||||
payload_json = EXCLUDED.payload_json,
|
||||
result_json = EXCLUDED.result_json,
|
||||
error_message = EXCLUDED.error_message,
|
||||
cancel_requested = EXCLUDED.cancel_requested,
|
||||
created_by = EXCLUDED.created_by,
|
||||
created_at_unix_secs = EXCLUDED.created_at_unix_secs,
|
||||
started_at_unix_secs = EXCLUDED.started_at_unix_secs,
|
||||
finished_at_unix_secs = EXCLUDED.finished_at_unix_secs,
|
||||
updated_at_unix_secs = EXCLUDED.updated_at_unix_secs
|
||||
"#,
|
||||
)
|
||||
.bind(&run.id)
|
||||
.bind(&run.task_key)
|
||||
.bind(run.kind.as_database())
|
||||
.bind(&run.trigger)
|
||||
.bind(run.status.as_database())
|
||||
.bind(u32_to_i32(run.attempt, "attempt")?)
|
||||
.bind(u32_to_i32(run.max_attempts, "max_attempts")?)
|
||||
.bind(run.owner_instance.as_deref())
|
||||
.bind(i32::from(run.progress_percent))
|
||||
.bind(run.progress_message.as_deref())
|
||||
.bind(run.payload_json.clone())
|
||||
.bind(run.result_json.clone())
|
||||
.bind(run.error_message.as_deref())
|
||||
.bind(run.cancel_requested)
|
||||
.bind(run.created_by.as_deref())
|
||||
.bind(u64_to_i64(
|
||||
run.created_at_unix_secs,
|
||||
"created_at_unix_secs",
|
||||
)?)
|
||||
.bind(run.started_at_unix_secs.map(|value| value as i64))
|
||||
.bind(run.finished_at_unix_secs.map(|value| value as i64))
|
||||
.bind(u64_to_i64(
|
||||
run.updated_at_unix_secs,
|
||||
"updated_at_unix_secs",
|
||||
)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
self.find_run(&run.id).await?.ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue("background task run missing after upsert".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
async fn request_cancel(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let affected = sqlx::query(
|
||||
"UPDATE background_task_runs SET cancel_requested = TRUE, updated_at_unix_secs = $2 WHERE id = $1",
|
||||
)
|
||||
.bind(run_id)
|
||||
.bind(u64_to_i64(updated_at_unix_secs, "updated_at_unix_secs")?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
Ok(affected > 0)
|
||||
}
|
||||
|
||||
async fn upsert_event(
|
||||
&self,
|
||||
event: UpsertBackgroundTaskEvent,
|
||||
) -> Result<StoredBackgroundTaskEvent, DataLayerError> {
|
||||
event.validate()?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO background_task_events (
|
||||
id,
|
||||
run_id,
|
||||
event_type,
|
||||
message,
|
||||
payload_json,
|
||||
created_at_unix_secs
|
||||
) VALUES ($1,$2,$3,$4,$5,$6)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
run_id = EXCLUDED.run_id,
|
||||
event_type = EXCLUDED.event_type,
|
||||
message = EXCLUDED.message,
|
||||
payload_json = EXCLUDED.payload_json,
|
||||
created_at_unix_secs = EXCLUDED.created_at_unix_secs
|
||||
"#,
|
||||
)
|
||||
.bind(&event.id)
|
||||
.bind(&event.run_id)
|
||||
.bind(&event.event_type)
|
||||
.bind(&event.message)
|
||||
.bind(event.payload_json.clone())
|
||||
.bind(u64_to_i64(
|
||||
event.created_at_unix_secs,
|
||||
"created_at_unix_secs",
|
||||
)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
let row = sqlx::query(&format!("{EVENT_COLUMNS} WHERE id = $1 LIMIT 1"))
|
||||
.bind(&event.id)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
map_event_row(&row)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_run_row(row: &PgRow) -> Result<StoredBackgroundTaskRun, DataLayerError> {
|
||||
let kind: String = row.try_get("kind").map_postgres_err()?;
|
||||
let status: String = row.try_get("status").map_postgres_err()?;
|
||||
let attempt: i32 = row.try_get("attempt").map_postgres_err()?;
|
||||
let max_attempts: i32 = row.try_get("max_attempts").map_postgres_err()?;
|
||||
let progress_percent: i32 = row.try_get("progress_percent").map_postgres_err()?;
|
||||
let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_postgres_err()?;
|
||||
let started_at_unix_secs: Option<i64> =
|
||||
row.try_get("started_at_unix_secs").map_postgres_err()?;
|
||||
let finished_at_unix_secs: Option<i64> =
|
||||
row.try_get("finished_at_unix_secs").map_postgres_err()?;
|
||||
let updated_at_unix_secs: i64 = row.try_get("updated_at_unix_secs").map_postgres_err()?;
|
||||
|
||||
Ok(StoredBackgroundTaskRun {
|
||||
id: row.try_get("id").map_postgres_err()?,
|
||||
task_key: row.try_get("task_key").map_postgres_err()?,
|
||||
kind: BackgroundTaskKind::from_database(&kind)?,
|
||||
trigger: row.try_get("trigger").map_postgres_err()?,
|
||||
status: BackgroundTaskStatus::from_database(&status)?,
|
||||
attempt: u32::try_from(attempt).unwrap_or_default(),
|
||||
max_attempts: u32::try_from(max_attempts).unwrap_or_default(),
|
||||
owner_instance: row.try_get("owner_instance").map_postgres_err()?,
|
||||
progress_percent: u16::try_from(progress_percent).unwrap_or_default(),
|
||||
progress_message: row.try_get("progress_message").map_postgres_err()?,
|
||||
payload_json: row.try_get("payload_json").map_postgres_err()?,
|
||||
result_json: row.try_get("result_json").map_postgres_err()?,
|
||||
error_message: row.try_get("error_message").map_postgres_err()?,
|
||||
cancel_requested: row.try_get("cancel_requested").map_postgres_err()?,
|
||||
created_by: row.try_get("created_by").map_postgres_err()?,
|
||||
created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(),
|
||||
started_at_unix_secs: started_at_unix_secs.and_then(|value| u64::try_from(value).ok()),
|
||||
finished_at_unix_secs: finished_at_unix_secs.and_then(|value| u64::try_from(value).ok()),
|
||||
updated_at_unix_secs: u64::try_from(updated_at_unix_secs).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn map_event_row(row: &PgRow) -> Result<StoredBackgroundTaskEvent, DataLayerError> {
|
||||
let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_postgres_err()?;
|
||||
Ok(StoredBackgroundTaskEvent {
|
||||
id: row.try_get("id").map_postgres_err()?,
|
||||
run_id: row.try_get("run_id").map_postgres_err()?,
|
||||
event_type: row.try_get("event_type").map_postgres_err()?,
|
||||
message: row.try_get("message").map_postgres_err()?,
|
||||
payload_json: row.try_get("payload_json").map_postgres_err()?,
|
||||
created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn i64_from_usize(value: usize, label: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn u64_to_i64(value: u64, label: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn u32_to_i32(value: u32, label: &str) -> Result<i32, DataLayerError> {
|
||||
i32::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}"))
|
||||
})
|
||||
}
|
||||
430
crates/aether-data/src/repository/background_tasks/sqlite.rs
Normal file
430
crates/aether-data/src/repository/background_tasks/sqlite.rs
Normal file
@@ -0,0 +1,430 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
|
||||
|
||||
use super::{
|
||||
BackgroundTaskKind, BackgroundTaskListQuery, BackgroundTaskReadRepository,
|
||||
BackgroundTaskStatus, BackgroundTaskSummary, BackgroundTaskWriteRepository,
|
||||
StoredBackgroundTaskEvent, StoredBackgroundTaskRun, StoredBackgroundTaskRunPage,
|
||||
UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const RUN_COLUMNS: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
task_key,
|
||||
kind,
|
||||
"trigger",
|
||||
status,
|
||||
attempt,
|
||||
max_attempts,
|
||||
owner_instance,
|
||||
progress_percent,
|
||||
progress_message,
|
||||
payload_json,
|
||||
result_json,
|
||||
error_message,
|
||||
cancel_requested,
|
||||
created_by,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
updated_at_unix_secs
|
||||
FROM background_task_runs
|
||||
"#;
|
||||
|
||||
const EVENT_COLUMNS: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
run_id,
|
||||
event_type,
|
||||
message,
|
||||
payload_json,
|
||||
created_at_unix_secs
|
||||
FROM background_task_events
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqliteBackgroundTaskRepository {
|
||||
pool: SqlitePool,
|
||||
}
|
||||
|
||||
impl SqliteBackgroundTaskRepository {
|
||||
pub fn new(pool: SqlitePool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
fn apply_run_filter(builder: &mut QueryBuilder<'_, Sqlite>, query: &BackgroundTaskListQuery) {
|
||||
let mut has_where = false;
|
||||
if let Some(kind) = query.kind {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
has_where = true;
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder.push("kind = ").push_bind(kind.as_database());
|
||||
}
|
||||
if let Some(status) = query.status {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
has_where = true;
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder.push("status = ").push_bind(status.as_database());
|
||||
}
|
||||
if let Some(trigger) = query.trigger.as_deref() {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
has_where = true;
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder
|
||||
.push("\"trigger\" = ")
|
||||
.push_bind(trigger.to_string());
|
||||
}
|
||||
if let Some(task_key_substring) = query.task_key_substring.as_deref() {
|
||||
if !has_where {
|
||||
builder.push(" WHERE ");
|
||||
} else {
|
||||
builder.push(" AND ");
|
||||
}
|
||||
builder.push("LOWER(task_key) LIKE ").push_bind(format!(
|
||||
"%{}%",
|
||||
task_key_substring.trim().to_ascii_lowercase()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskReadRepository for SqliteBackgroundTaskRepository {
|
||||
async fn find_run(
|
||||
&self,
|
||||
run_id: &str,
|
||||
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
||||
let row = sqlx::query(&format!("{RUN_COLUMNS} WHERE id = ? LIMIT 1"))
|
||||
.bind(run_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row.as_ref().map(map_run_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_runs(
|
||||
&self,
|
||||
query: &BackgroundTaskListQuery,
|
||||
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
|
||||
let limit = query.limit.max(1);
|
||||
let mut count_builder =
|
||||
QueryBuilder::<Sqlite>::new("SELECT COUNT(id) AS total FROM background_task_runs");
|
||||
Self::apply_run_filter(&mut count_builder, query);
|
||||
let total = count_builder
|
||||
.build_query_scalar::<i64>()
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(RUN_COLUMNS);
|
||||
Self::apply_run_filter(&mut builder, query);
|
||||
builder
|
||||
.push(" ORDER BY created_at_unix_secs DESC, updated_at_unix_secs DESC")
|
||||
.push(" LIMIT ")
|
||||
.push_bind(i64_from_usize(limit, "run limit")?)
|
||||
.push(" OFFSET ")
|
||||
.push_bind(i64_from_usize(query.offset, "run offset")?);
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(map_run_row)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(StoredBackgroundTaskRunPage {
|
||||
items,
|
||||
total: usize::try_from(total).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_events(
|
||||
&self,
|
||||
run_id: &str,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
|
||||
let limit = limit.max(1);
|
||||
let rows = sqlx::query(&format!(
|
||||
"{EVENT_COLUMNS} WHERE run_id = ? ORDER BY created_at_unix_secs ASC, id ASC LIMIT ? OFFSET ?"
|
||||
))
|
||||
.bind(run_id)
|
||||
.bind(i64_from_usize(limit, "event limit")?)
|
||||
.bind(i64_from_usize(offset, "event offset")?)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_event_row).collect()
|
||||
}
|
||||
|
||||
async fn summarize_runs(&self) -> Result<BackgroundTaskSummary, DataLayerError> {
|
||||
let total = sqlx::query_scalar::<_, i64>("SELECT COUNT(id) FROM background_task_runs")
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let running_count = sqlx::query_scalar::<_, i64>(
|
||||
"SELECT COUNT(id) FROM background_task_runs WHERE status = 'running'",
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let status_rows = sqlx::query(
|
||||
"SELECT status, COUNT(id) AS total FROM background_task_runs GROUP BY status",
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let kind_rows =
|
||||
sqlx::query("SELECT kind, COUNT(id) AS total FROM background_task_runs GROUP BY kind")
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
let mut by_status = std::collections::BTreeMap::new();
|
||||
for row in status_rows {
|
||||
let key: String = row.try_get("status").map_sql_err()?;
|
||||
let count: i64 = row.try_get("total").map_sql_err()?;
|
||||
by_status.insert(key, u64::try_from(count).unwrap_or_default());
|
||||
}
|
||||
let mut by_kind = std::collections::BTreeMap::new();
|
||||
for row in kind_rows {
|
||||
let key: String = row.try_get("kind").map_sql_err()?;
|
||||
let count: i64 = row.try_get("total").map_sql_err()?;
|
||||
by_kind.insert(key, u64::try_from(count).unwrap_or_default());
|
||||
}
|
||||
|
||||
Ok(BackgroundTaskSummary {
|
||||
total: u64::try_from(total).unwrap_or_default(),
|
||||
running_count: u64::try_from(running_count).unwrap_or_default(),
|
||||
by_status,
|
||||
by_kind,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BackgroundTaskWriteRepository for SqliteBackgroundTaskRepository {
|
||||
async fn upsert_run(
|
||||
&self,
|
||||
run: UpsertBackgroundTaskRun,
|
||||
) -> Result<StoredBackgroundTaskRun, DataLayerError> {
|
||||
run.validate()?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO background_task_runs (
|
||||
id,
|
||||
task_key,
|
||||
kind,
|
||||
"trigger",
|
||||
status,
|
||||
attempt,
|
||||
max_attempts,
|
||||
owner_instance,
|
||||
progress_percent,
|
||||
progress_message,
|
||||
payload_json,
|
||||
result_json,
|
||||
error_message,
|
||||
cancel_requested,
|
||||
created_by,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
updated_at_unix_secs
|
||||
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
task_key = excluded.task_key,
|
||||
kind = excluded.kind,
|
||||
"trigger" = excluded."trigger",
|
||||
status = excluded.status,
|
||||
attempt = excluded.attempt,
|
||||
max_attempts = excluded.max_attempts,
|
||||
owner_instance = excluded.owner_instance,
|
||||
progress_percent = excluded.progress_percent,
|
||||
progress_message = excluded.progress_message,
|
||||
payload_json = excluded.payload_json,
|
||||
result_json = excluded.result_json,
|
||||
error_message = excluded.error_message,
|
||||
cancel_requested = excluded.cancel_requested,
|
||||
created_by = excluded.created_by,
|
||||
created_at_unix_secs = excluded.created_at_unix_secs,
|
||||
started_at_unix_secs = excluded.started_at_unix_secs,
|
||||
finished_at_unix_secs = excluded.finished_at_unix_secs,
|
||||
updated_at_unix_secs = excluded.updated_at_unix_secs
|
||||
"#,
|
||||
)
|
||||
.bind(&run.id)
|
||||
.bind(&run.task_key)
|
||||
.bind(run.kind.as_database())
|
||||
.bind(&run.trigger)
|
||||
.bind(run.status.as_database())
|
||||
.bind(i64::from(run.attempt))
|
||||
.bind(i64::from(run.max_attempts))
|
||||
.bind(run.owner_instance.as_deref())
|
||||
.bind(i32::from(run.progress_percent))
|
||||
.bind(run.progress_message.as_deref())
|
||||
.bind(run.payload_json.as_ref().map(serde_json::Value::to_string))
|
||||
.bind(run.result_json.as_ref().map(serde_json::Value::to_string))
|
||||
.bind(run.error_message.as_deref())
|
||||
.bind(run.cancel_requested)
|
||||
.bind(run.created_by.as_deref())
|
||||
.bind(u64_to_i64(
|
||||
run.created_at_unix_secs,
|
||||
"created_at_unix_secs",
|
||||
)?)
|
||||
.bind(run.started_at_unix_secs.map(|value| value as i64))
|
||||
.bind(run.finished_at_unix_secs.map(|value| value as i64))
|
||||
.bind(u64_to_i64(
|
||||
run.updated_at_unix_secs,
|
||||
"updated_at_unix_secs",
|
||||
)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
self.find_run(&run.id).await?.ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue("background task run missing after upsert".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
async fn request_cancel(
|
||||
&self,
|
||||
run_id: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let affected = sqlx::query(
|
||||
"UPDATE background_task_runs SET cancel_requested = 1, updated_at_unix_secs = ? WHERE id = ?",
|
||||
)
|
||||
.bind(u64_to_i64(updated_at_unix_secs, "updated_at_unix_secs")?)
|
||||
.bind(run_id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
Ok(affected > 0)
|
||||
}
|
||||
|
||||
async fn upsert_event(
|
||||
&self,
|
||||
event: UpsertBackgroundTaskEvent,
|
||||
) -> Result<StoredBackgroundTaskEvent, DataLayerError> {
|
||||
event.validate()?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO background_task_events (
|
||||
id, run_id, event_type, message, payload_json, created_at_unix_secs
|
||||
) VALUES (?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
run_id = excluded.run_id,
|
||||
event_type = excluded.event_type,
|
||||
message = excluded.message,
|
||||
payload_json = excluded.payload_json,
|
||||
created_at_unix_secs = excluded.created_at_unix_secs
|
||||
"#,
|
||||
)
|
||||
.bind(&event.id)
|
||||
.bind(&event.run_id)
|
||||
.bind(&event.event_type)
|
||||
.bind(&event.message)
|
||||
.bind(
|
||||
event
|
||||
.payload_json
|
||||
.as_ref()
|
||||
.map(serde_json::Value::to_string),
|
||||
)
|
||||
.bind(u64_to_i64(
|
||||
event.created_at_unix_secs,
|
||||
"created_at_unix_secs",
|
||||
)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
let row = sqlx::query(&format!("{EVENT_COLUMNS} WHERE id = ? LIMIT 1"))
|
||||
.bind(&event.id)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
map_event_row(&row)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_run_row(row: &SqliteRow) -> Result<StoredBackgroundTaskRun, DataLayerError> {
|
||||
let kind: String = row.try_get("kind").map_sql_err()?;
|
||||
let status: String = row.try_get("status").map_sql_err()?;
|
||||
let attempt: i64 = row.try_get("attempt").map_sql_err()?;
|
||||
let max_attempts: i64 = row.try_get("max_attempts").map_sql_err()?;
|
||||
let progress_percent: i32 = row.try_get("progress_percent").map_sql_err()?;
|
||||
let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?;
|
||||
let started_at_unix_secs: Option<i64> = row.try_get("started_at_unix_secs").map_sql_err()?;
|
||||
let finished_at_unix_secs: Option<i64> = row.try_get("finished_at_unix_secs").map_sql_err()?;
|
||||
let updated_at_unix_secs: i64 = row.try_get("updated_at_unix_secs").map_sql_err()?;
|
||||
|
||||
Ok(StoredBackgroundTaskRun {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
task_key: row.try_get("task_key").map_sql_err()?,
|
||||
kind: BackgroundTaskKind::from_database(&kind)?,
|
||||
trigger: row.try_get("trigger").map_sql_err()?,
|
||||
status: BackgroundTaskStatus::from_database(&status)?,
|
||||
attempt: u32::try_from(attempt).unwrap_or_default(),
|
||||
max_attempts: u32::try_from(max_attempts).unwrap_or_default(),
|
||||
owner_instance: row.try_get("owner_instance").map_sql_err()?,
|
||||
progress_percent: u16::try_from(progress_percent).unwrap_or_default(),
|
||||
progress_message: row.try_get("progress_message").map_sql_err()?,
|
||||
payload_json: parse_optional_json(row.try_get("payload_json").map_sql_err()?)?,
|
||||
result_json: parse_optional_json(row.try_get("result_json").map_sql_err()?)?,
|
||||
error_message: row.try_get("error_message").map_sql_err()?,
|
||||
cancel_requested: row.try_get("cancel_requested").map_sql_err()?,
|
||||
created_by: row.try_get("created_by").map_sql_err()?,
|
||||
created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(),
|
||||
started_at_unix_secs: started_at_unix_secs.and_then(|value| u64::try_from(value).ok()),
|
||||
finished_at_unix_secs: finished_at_unix_secs.and_then(|value| u64::try_from(value).ok()),
|
||||
updated_at_unix_secs: u64::try_from(updated_at_unix_secs).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn map_event_row(row: &SqliteRow) -> Result<StoredBackgroundTaskEvent, DataLayerError> {
|
||||
let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?;
|
||||
Ok(StoredBackgroundTaskEvent {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
run_id: row.try_get("run_id").map_sql_err()?,
|
||||
event_type: row.try_get("event_type").map_sql_err()?,
|
||||
message: row.try_get("message").map_sql_err()?,
|
||||
payload_json: parse_optional_json(row.try_get("payload_json").map_sql_err()?)?,
|
||||
created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_json(value: Option<String>) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
value
|
||||
.map(|raw| {
|
||||
serde_json::from_str::<serde_json::Value>(&raw).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid background task json payload: {err}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn i64_from_usize(value: usize, label: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn u64_to_i64(value: u64, label: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}"))
|
||||
})
|
||||
}
|
||||
@@ -8,6 +8,7 @@ pub mod announcements;
|
||||
pub mod audit;
|
||||
pub mod auth;
|
||||
pub mod auth_modules;
|
||||
pub mod background_tasks;
|
||||
pub mod billing;
|
||||
pub mod candidate_selection;
|
||||
pub mod candidates;
|
||||
|
||||
@@ -41,6 +41,12 @@ impl AdminSystemPurgeSummary {
|
||||
*self.affected.entry(key.into()).or_insert(0) += count;
|
||||
}
|
||||
|
||||
pub fn merge(&mut self, other: &Self) {
|
||||
for (key, count) in &other.affected {
|
||||
self.add(key.clone(), *count);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn total(&self) -> u64 {
|
||||
self.affected.values().copied().sum()
|
||||
}
|
||||
|
||||
15
crates/aether-task-runtime/Cargo.toml
Normal file
15
crates/aether-task-runtime/Cargo.toml
Normal file
@@ -0,0 +1,15 @@
|
||||
[package]
|
||||
name = "aether-task-runtime"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Shared async task runtime primitives for Aether services"
|
||||
|
||||
[dependencies]
|
||||
aether-runtime.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-util.workspace = true
|
||||
tracing.workspace = true
|
||||
227
crates/aether-task-runtime/src/lib.rs
Normal file
227
crates/aether-task-runtime/src/lib.rs
Normal file
@@ -0,0 +1,227 @@
|
||||
use std::future::Future;
|
||||
|
||||
use aether_runtime::task::spawn_named;
|
||||
use tokio::task::{JoinHandle, JoinSet};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
#[derive(
|
||||
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
|
||||
)]
|
||||
pub enum TaskKind {
|
||||
Scheduled,
|
||||
Daemon,
|
||||
OnDemand,
|
||||
FireAndForget,
|
||||
}
|
||||
|
||||
impl TaskKind {
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Scheduled => "scheduled",
|
||||
Self::Daemon => "daemon",
|
||||
Self::OnDemand => "on_demand",
|
||||
Self::FireAndForget => "fire_and_forget",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(
|
||||
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
|
||||
)]
|
||||
pub enum TaskStatus {
|
||||
Queued,
|
||||
Running,
|
||||
Retrying,
|
||||
Succeeded,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Skipped,
|
||||
}
|
||||
|
||||
impl TaskStatus {
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Queued => "queued",
|
||||
Self::Running => "running",
|
||||
Self::Retrying => "retrying",
|
||||
Self::Succeeded => "succeeded",
|
||||
Self::Failed => "failed",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::Skipped => "skipped",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct RetryPolicy {
|
||||
pub max_attempts: u32,
|
||||
}
|
||||
|
||||
impl Default for RetryPolicy {
|
||||
fn default() -> Self {
|
||||
Self { max_attempts: 1 }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TaskDefinition {
|
||||
pub key: &'static str,
|
||||
pub kind: TaskKind,
|
||||
pub trigger: &'static str,
|
||||
pub singleton: bool,
|
||||
pub persist_history: bool,
|
||||
pub retry_policy: RetryPolicy,
|
||||
}
|
||||
|
||||
impl TaskDefinition {
|
||||
pub const fn new(
|
||||
key: &'static str,
|
||||
kind: TaskKind,
|
||||
trigger: &'static str,
|
||||
singleton: bool,
|
||||
persist_history: bool,
|
||||
retry_policy: RetryPolicy,
|
||||
) -> Self {
|
||||
Self {
|
||||
key,
|
||||
kind,
|
||||
trigger,
|
||||
singleton,
|
||||
persist_history,
|
||||
retry_policy,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TaskContext<TPayload = serde_json::Value> {
|
||||
run_id: String,
|
||||
task_key: String,
|
||||
payload: Option<TPayload>,
|
||||
cancellation_token: CancellationToken,
|
||||
}
|
||||
|
||||
impl<TPayload> TaskContext<TPayload> {
|
||||
pub fn new(
|
||||
run_id: impl Into<String>,
|
||||
task_key: impl Into<String>,
|
||||
payload: Option<TPayload>,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> Self {
|
||||
Self {
|
||||
run_id: run_id.into(),
|
||||
task_key: task_key.into(),
|
||||
payload,
|
||||
cancellation_token,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn run_id(&self) -> &str {
|
||||
&self.run_id
|
||||
}
|
||||
|
||||
pub fn task_key(&self) -> &str {
|
||||
&self.task_key
|
||||
}
|
||||
|
||||
pub fn payload(&self) -> Option<&TPayload> {
|
||||
self.payload.as_ref()
|
||||
}
|
||||
|
||||
pub fn cancellation_token(&self) -> CancellationToken {
|
||||
self.cancellation_token.clone()
|
||||
}
|
||||
|
||||
pub fn is_cancelled(&self) -> bool {
|
||||
self.cancellation_token.is_cancelled()
|
||||
}
|
||||
|
||||
pub async fn cancelled(&self) {
|
||||
self.cancellation_token.cancelled().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TaskSupervisor {
|
||||
cancellation_token: CancellationToken,
|
||||
join_set: JoinSet<()>,
|
||||
supervised_task_count: usize,
|
||||
}
|
||||
|
||||
impl TaskSupervisor {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
cancellation_token: CancellationToken::new(),
|
||||
join_set: JoinSet::new(),
|
||||
supervised_task_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cancellation_token(&self) -> CancellationToken {
|
||||
self.cancellation_token.clone()
|
||||
}
|
||||
|
||||
pub fn spawn_named<F>(&mut self, task_name: &'static str, future: F)
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
self.supervised_task_count = self.supervised_task_count.saturating_add(1);
|
||||
let cancellation_token = self.cancellation_token.clone();
|
||||
self.join_set.spawn(async move {
|
||||
let mut handle = spawn_named(task_name, future);
|
||||
tokio::select! {
|
||||
_ = cancellation_token.cancelled() => {
|
||||
handle.abort();
|
||||
let _ = handle.await;
|
||||
}
|
||||
result = &mut handle => {
|
||||
if let Err(error) = result {
|
||||
warn!(task = task_name, error = ?error, "supervised task failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub fn supervise_handle(&mut self, task_name: &'static str, mut handle: JoinHandle<()>) {
|
||||
self.supervised_task_count = self.supervised_task_count.saturating_add(1);
|
||||
let cancellation_token = self.cancellation_token.clone();
|
||||
self.join_set.spawn(async move {
|
||||
tokio::select! {
|
||||
_ = cancellation_token.cancelled() => {
|
||||
handle.abort();
|
||||
let _ = handle.await;
|
||||
}
|
||||
result = &mut handle => {
|
||||
if let Err(error) = result {
|
||||
warn!(task = task_name, error = ?error, "supervised task failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.supervised_task_count == 0
|
||||
}
|
||||
|
||||
pub fn task_count(&self) -> usize {
|
||||
self.supervised_task_count
|
||||
}
|
||||
|
||||
pub fn cancel(&self) {
|
||||
self.cancellation_token.cancel();
|
||||
}
|
||||
|
||||
pub async fn shutdown(mut self) {
|
||||
self.cancel();
|
||||
while self.join_set.join_next().await.is_some() {}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TaskSupervisor {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user