registry: extract platform and usage modules
This commit is contained in:
a.tolmachev
@@ -0,0 +1,180 @@
|
||||
use super::*;
|
||||
|
||||
impl PostgresRegistry {
|
||||
pub async fn list_platform_api_keys(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
) -> Result<Vec<PlatformApiKeyRecord>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
name,
|
||||
prefix,
|
||||
scopes_json,
|
||||
status,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(last_used_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_used_at
|
||||
from platform_api_keys
|
||||
where workspace_id = $1
|
||||
order by created_at desc",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_platform_api_key_record).collect()
|
||||
}
|
||||
|
||||
pub async fn get_platform_api_key_by_secret_for_workspace_slug(
|
||||
&self,
|
||||
workspace_slug: &str,
|
||||
secret_hash: &str,
|
||||
) -> Result<Option<PlatformApiKeyRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
k.id,
|
||||
k.workspace_id,
|
||||
k.name,
|
||||
k.prefix,
|
||||
k.scopes_json,
|
||||
k.status,
|
||||
to_char(k.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(k.last_used_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_used_at
|
||||
from platform_api_keys k
|
||||
join workspaces w on w.id = k.workspace_id
|
||||
where w.slug = $1
|
||||
and k.secret_hash = $2
|
||||
and k.status = 'active'
|
||||
limit 1",
|
||||
)
|
||||
.bind(workspace_slug)
|
||||
.bind(secret_hash)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_platform_api_key_record).transpose()
|
||||
}
|
||||
|
||||
pub async fn create_platform_api_key(
|
||||
&self,
|
||||
request: CreatePlatformApiKeyRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"insert into platform_api_keys (
|
||||
id,
|
||||
workspace_id,
|
||||
name,
|
||||
prefix,
|
||||
secret_hash,
|
||||
scopes_json,
|
||||
status,
|
||||
created_at,
|
||||
last_used_at,
|
||||
revoked_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8::timestamptz, $9::timestamptz, $10::timestamptz
|
||||
)",
|
||||
)
|
||||
.bind(request.api_key.id.as_str())
|
||||
.bind(request.api_key.workspace_id.as_str())
|
||||
.bind(&request.api_key.name)
|
||||
.bind(&request.api_key.prefix)
|
||||
.bind(request.secret_hash)
|
||||
.bind(Json(serialize_json_value(&request.api_key.scopes)?))
|
||||
.bind(serialize_enum_text(&request.api_key.status, "status")?)
|
||||
.bind(&request.api_key.created_at)
|
||||
.bind(request.api_key.last_used_at.as_deref())
|
||||
.bind(Option::<&str>::None)
|
||||
.execute(&self.pool)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(_) => Ok(()),
|
||||
Err(sqlx::Error::Database(error))
|
||||
if error.constraint() == Some("platform_api_keys_workspace_name_idx") =>
|
||||
{
|
||||
Err(RegistryError::Storage(sqlx::Error::Database(error)))
|
||||
}
|
||||
Err(error) => Err(RegistryError::Storage(error)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn revoke_platform_api_key(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
key_id: &PlatformApiKeyId,
|
||||
revoked_at: &str,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"update platform_api_keys
|
||||
set status = $1,
|
||||
revoked_at = $2::timestamptz
|
||||
where workspace_id = $3 and id = $4",
|
||||
)
|
||||
.bind(serialize_enum_text(
|
||||
&PlatformApiKeyStatus::Revoked,
|
||||
"status",
|
||||
)?)
|
||||
.bind(revoked_at)
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(key_id.as_str())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::PlatformApiKeyNotFound {
|
||||
key_id: key_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn delete_platform_api_key(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
key_id: &PlatformApiKeyId,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result =
|
||||
sqlx::query("delete from platform_api_keys where workspace_id = $1 and id = $2")
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(key_id.as_str())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::PlatformApiKeyNotFound {
|
||||
key_id: key_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn touch_platform_api_key(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
key_id: &PlatformApiKeyId,
|
||||
used_at: &str,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"update platform_api_keys
|
||||
set last_used_at = $3::timestamptz
|
||||
where workspace_id = $1 and id = $2",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(key_id.as_str())
|
||||
.bind(used_at)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::PlatformApiKeyNotFound {
|
||||
key_id: key_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,10 @@
|
||||
mod api_key;
|
||||
mod auth;
|
||||
mod observability;
|
||||
mod secret;
|
||||
mod stream;
|
||||
mod workspace;
|
||||
mod yaml_import;
|
||||
|
||||
use crank_core::{
|
||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
|
||||
@@ -176,626 +179,6 @@ impl PostgresRegistry {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_platform_api_keys(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
) -> Result<Vec<PlatformApiKeyRecord>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
name,
|
||||
prefix,
|
||||
scopes_json,
|
||||
status,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(last_used_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_used_at
|
||||
from platform_api_keys
|
||||
where workspace_id = $1
|
||||
order by created_at desc",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_platform_api_key_record).collect()
|
||||
}
|
||||
|
||||
pub async fn get_platform_api_key_by_secret_for_workspace_slug(
|
||||
&self,
|
||||
workspace_slug: &str,
|
||||
secret_hash: &str,
|
||||
) -> Result<Option<PlatformApiKeyRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
k.id,
|
||||
k.workspace_id,
|
||||
k.name,
|
||||
k.prefix,
|
||||
k.scopes_json,
|
||||
k.status,
|
||||
to_char(k.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(k.last_used_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_used_at
|
||||
from platform_api_keys k
|
||||
join workspaces w on w.id = k.workspace_id
|
||||
where w.slug = $1
|
||||
and k.secret_hash = $2
|
||||
and k.status = 'active'
|
||||
limit 1",
|
||||
)
|
||||
.bind(workspace_slug)
|
||||
.bind(secret_hash)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_platform_api_key_record).transpose()
|
||||
}
|
||||
|
||||
pub async fn create_platform_api_key(
|
||||
&self,
|
||||
request: CreatePlatformApiKeyRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"insert into platform_api_keys (
|
||||
id,
|
||||
workspace_id,
|
||||
name,
|
||||
prefix,
|
||||
secret_hash,
|
||||
scopes_json,
|
||||
status,
|
||||
created_at,
|
||||
last_used_at,
|
||||
revoked_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8::timestamptz, $9::timestamptz, $10::timestamptz
|
||||
)",
|
||||
)
|
||||
.bind(request.api_key.id.as_str())
|
||||
.bind(request.api_key.workspace_id.as_str())
|
||||
.bind(&request.api_key.name)
|
||||
.bind(&request.api_key.prefix)
|
||||
.bind(request.secret_hash)
|
||||
.bind(Json(serialize_json_value(&request.api_key.scopes)?))
|
||||
.bind(serialize_enum_text(&request.api_key.status, "status")?)
|
||||
.bind(&request.api_key.created_at)
|
||||
.bind(request.api_key.last_used_at.as_deref())
|
||||
.bind(Option::<&str>::None)
|
||||
.execute(&self.pool)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(_) => Ok(()),
|
||||
Err(sqlx::Error::Database(error))
|
||||
if error.constraint() == Some("platform_api_keys_workspace_name_idx") =>
|
||||
{
|
||||
Err(RegistryError::Storage(sqlx::Error::Database(error)))
|
||||
}
|
||||
Err(error) => Err(RegistryError::Storage(error)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn revoke_platform_api_key(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
key_id: &PlatformApiKeyId,
|
||||
revoked_at: &str,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"update platform_api_keys
|
||||
set status = $1,
|
||||
revoked_at = $2::timestamptz
|
||||
where workspace_id = $3 and id = $4",
|
||||
)
|
||||
.bind(serialize_enum_text(
|
||||
&PlatformApiKeyStatus::Revoked,
|
||||
"status",
|
||||
)?)
|
||||
.bind(revoked_at)
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(key_id.as_str())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::PlatformApiKeyNotFound {
|
||||
key_id: key_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn delete_platform_api_key(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
key_id: &PlatformApiKeyId,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result =
|
||||
sqlx::query("delete from platform_api_keys where workspace_id = $1 and id = $2")
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(key_id.as_str())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::PlatformApiKeyNotFound {
|
||||
key_id: key_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn touch_platform_api_key(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
key_id: &PlatformApiKeyId,
|
||||
used_at: &str,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"update platform_api_keys
|
||||
set last_used_at = $3::timestamptz
|
||||
where workspace_id = $1 and id = $2",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(key_id.as_str())
|
||||
.bind(used_at)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::PlatformApiKeyNotFound {
|
||||
key_id: key_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_invocation_log(
|
||||
&self,
|
||||
request: CreateInvocationLogRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
sqlx::query(
|
||||
"insert into invocation_logs (
|
||||
id,
|
||||
workspace_id,
|
||||
agent_id,
|
||||
operation_id,
|
||||
source,
|
||||
level,
|
||||
status,
|
||||
tool_name,
|
||||
message,
|
||||
request_id,
|
||||
status_code,
|
||||
duration_ms,
|
||||
error_kind,
|
||||
request_preview_json,
|
||||
response_preview_json,
|
||||
created_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16::timestamptz
|
||||
)",
|
||||
)
|
||||
.bind(request.log.id.as_str())
|
||||
.bind(request.log.workspace_id.as_str())
|
||||
.bind(request.log.agent_id.as_ref().map(|value| value.as_str()))
|
||||
.bind(request.log.operation_id.as_str())
|
||||
.bind(serialize_enum_text(&request.log.source, "source")?)
|
||||
.bind(serialize_enum_text(&request.log.level, "level")?)
|
||||
.bind(serialize_enum_text(&request.log.status, "status")?)
|
||||
.bind(&request.log.tool_name)
|
||||
.bind(&request.log.message)
|
||||
.bind(&request.log.request_id)
|
||||
.bind(request.log.status_code.map(i32::from))
|
||||
.bind(i64::try_from(request.log.duration_ms).map_err(|_| {
|
||||
RegistryError::InvalidNumericValue {
|
||||
field: "duration_ms",
|
||||
value: i64::MAX,
|
||||
}
|
||||
})?)
|
||||
.bind(&request.log.error_kind)
|
||||
.bind(Json(request.log.request_preview.clone()))
|
||||
.bind(Json(request.log.response_preview.clone()))
|
||||
.bind(&request.log.created_at)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_invocation_logs(
|
||||
&self,
|
||||
query: ListInvocationLogsQuery<'_>,
|
||||
) -> Result<Vec<InvocationLogRecord>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
l.id,
|
||||
l.workspace_id,
|
||||
l.agent_id,
|
||||
l.operation_id,
|
||||
l.source,
|
||||
l.level,
|
||||
l.status,
|
||||
l.tool_name,
|
||||
l.message,
|
||||
l.request_id,
|
||||
l.status_code,
|
||||
l.duration_ms,
|
||||
l.error_kind,
|
||||
l.request_preview_json,
|
||||
l.response_preview_json,
|
||||
to_char(l.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.MS\"Z\"') as created_at,
|
||||
o.name as operation_name,
|
||||
o.display_name as operation_display_name,
|
||||
a.slug as agent_slug,
|
||||
a.display_name as agent_display_name
|
||||
from invocation_logs l
|
||||
join operations o on o.id = l.operation_id
|
||||
left join agents a on a.id = l.agent_id
|
||||
where l.workspace_id = $1
|
||||
and ($2::text is null or l.level = $2)
|
||||
and ($3::text is null or l.source = $3)
|
||||
and ($4::text is null or l.operation_id = $4)
|
||||
and ($5::text is null or l.agent_id = $5)
|
||||
and ($6::timestamptz is null or l.created_at >= $6::timestamptz)
|
||||
and (
|
||||
$7::text is null
|
||||
or l.tool_name ilike '%' || $7 || '%'
|
||||
or l.message ilike '%' || $7 || '%'
|
||||
or o.name ilike '%' || $7 || '%'
|
||||
or o.display_name ilike '%' || $7 || '%'
|
||||
)
|
||||
order by l.created_at desc
|
||||
limit $8",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(
|
||||
query
|
||||
.level
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "level"))
|
||||
.transpose()?,
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.bind(query.operation_id.map(|value| value.as_str()))
|
||||
.bind(query.agent_id.map(|value| value.as_str()))
|
||||
.bind(query.created_after)
|
||||
.bind(query.search_text)
|
||||
.bind(i64::from(query.limit))
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_invocation_log_record).collect()
|
||||
}
|
||||
|
||||
pub async fn get_invocation_log(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
log_id: &InvocationLogId,
|
||||
) -> Result<Option<InvocationLogRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
l.id,
|
||||
l.workspace_id,
|
||||
l.agent_id,
|
||||
l.operation_id,
|
||||
l.source,
|
||||
l.level,
|
||||
l.status,
|
||||
l.tool_name,
|
||||
l.message,
|
||||
l.request_id,
|
||||
l.status_code,
|
||||
l.duration_ms,
|
||||
l.error_kind,
|
||||
l.request_preview_json,
|
||||
l.response_preview_json,
|
||||
to_char(l.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.MS\"Z\"') as created_at,
|
||||
o.name as operation_name,
|
||||
o.display_name as operation_display_name,
|
||||
a.slug as agent_slug,
|
||||
a.display_name as agent_display_name
|
||||
from invocation_logs l
|
||||
join operations o on o.id = l.operation_id
|
||||
left join agents a on a.id = l.agent_id
|
||||
where l.workspace_id = $1 and l.id = $2",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(log_id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_invocation_log_record).transpose()
|
||||
}
|
||||
|
||||
pub async fn summarize_usage(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<UsageSummary, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and created_at >= $2::timestamptz
|
||||
and ($3::text is null or source = $3)",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
let calls_total = row.try_get::<i64, _>("calls_total")?;
|
||||
let calls_ok = row.try_get::<i64, _>("calls_ok")?;
|
||||
let calls_error = row.try_get::<i64, _>("calls_error")?;
|
||||
let rollup = UsageRollup {
|
||||
workspace_id: query.workspace_id.clone(),
|
||||
agent_id: None,
|
||||
operation_id: None,
|
||||
period: query.period,
|
||||
calls_total: to_u64(calls_total, "calls_total")?,
|
||||
calls_ok: to_u64(calls_ok, "calls_ok")?,
|
||||
calls_error: to_u64(calls_error, "calls_error")?,
|
||||
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
|
||||
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
|
||||
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
|
||||
};
|
||||
|
||||
let success_rate = if rollup.calls_total == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(rollup.calls_ok as f64 / rollup.calls_total as f64) * 100.0
|
||||
};
|
||||
|
||||
Ok(UsageSummary {
|
||||
rollup,
|
||||
success_rate,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn list_usage_timeline(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<Vec<UsageTimelinePoint>, RegistryError> {
|
||||
let bucket = usage_bucket_sql(query.bucket);
|
||||
let sql = format!(
|
||||
"select
|
||||
to_char(date_trunc('{bucket}', created_at at time zone 'UTC'), 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as bucket_start,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and created_at >= $2::timestamptz
|
||||
and ($3::text is null or source = $3)
|
||||
group by 1
|
||||
order by 1 asc"
|
||||
);
|
||||
let rows = sqlx::query(&sql)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter()
|
||||
.map(|row| {
|
||||
Ok(UsageTimelinePoint {
|
||||
bucket_start: row.try_get("bucket_start")?,
|
||||
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
|
||||
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn list_usage_by_operation(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<Vec<UsageOperationBreakdown>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
o.id as operation_id,
|
||||
o.name as operation_name,
|
||||
o.display_name as operation_display_name,
|
||||
o.protocol,
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where l.status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by l.duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by l.duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by l.duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs l
|
||||
join operations o on o.id = l.operation_id
|
||||
where l.workspace_id = $1
|
||||
and l.created_at >= $2::timestamptz
|
||||
and ($3::text is null or l.source = $3)
|
||||
group by o.id, o.name, o.display_name, o.protocol
|
||||
order by calls_total desc, o.name asc",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_usage_operation_breakdown).collect()
|
||||
}
|
||||
|
||||
pub async fn get_usage_for_operation(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
operation_id: &OperationId,
|
||||
) -> Result<Option<UsageRollupRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and operation_id = $2
|
||||
and created_at >= $3::timestamptz
|
||||
and ($4::text is null or source = $4)",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(operation_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
let calls_total = to_u64(row.try_get::<i64, _>("calls_total")?, "calls_total")?;
|
||||
if calls_total == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(UsageRollupRecord {
|
||||
rollup: UsageRollup {
|
||||
workspace_id: query.workspace_id.clone(),
|
||||
agent_id: None,
|
||||
operation_id: Some(operation_id.clone()),
|
||||
period: query.period,
|
||||
calls_total,
|
||||
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
|
||||
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
|
||||
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
|
||||
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
|
||||
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn list_usage_by_agent(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<Vec<UsageAgentBreakdown>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
a.id as agent_id,
|
||||
a.slug as agent_slug,
|
||||
a.display_name as agent_display_name,
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where l.status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by l.duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by l.duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by l.duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs l
|
||||
join agents a on a.id = l.agent_id
|
||||
where l.workspace_id = $1
|
||||
and l.created_at >= $2::timestamptz
|
||||
and ($3::text is null or l.source = $3)
|
||||
group by a.id, a.slug, a.display_name
|
||||
order by calls_total desc, a.slug asc",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_usage_agent_breakdown).collect()
|
||||
}
|
||||
|
||||
pub async fn get_usage_for_agent(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
agent_id: &AgentId,
|
||||
) -> Result<Option<UsageRollupRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and agent_id = $2
|
||||
and created_at >= $3::timestamptz
|
||||
and ($4::text is null or source = $4)",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(agent_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
let calls_total = to_u64(row.try_get::<i64, _>("calls_total")?, "calls_total")?;
|
||||
if calls_total == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(UsageRollupRecord {
|
||||
rollup: UsageRollup {
|
||||
workspace_id: query.workspace_id.clone(),
|
||||
agent_id: Some(agent_id.clone()),
|
||||
operation_id: None,
|
||||
period: query.period,
|
||||
calls_total,
|
||||
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
|
||||
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
|
||||
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
|
||||
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
|
||||
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn list_agents(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
@@ -2091,102 +1474,6 @@ impl PostgresRegistry {
|
||||
rows.iter().map(map_descriptor_metadata).collect()
|
||||
}
|
||||
|
||||
pub async fn create_yaml_import_job(
|
||||
&self,
|
||||
request: CreateYamlImportJobRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
sqlx::query(
|
||||
"insert into yaml_import_jobs (
|
||||
id,
|
||||
source_sample_id,
|
||||
status,
|
||||
format_version,
|
||||
mode,
|
||||
result_operation_id,
|
||||
result_version,
|
||||
error_text,
|
||||
created_at,
|
||||
finished_at
|
||||
) values ($1, $2, $3, $4, $5, null, null, null, $6::timestamptz, null)",
|
||||
)
|
||||
.bind(request.id.as_str())
|
||||
.bind(request.source_sample_id.map(|value| value.as_str()))
|
||||
.bind(serialize_enum_text(
|
||||
&YamlImportJobStatus::Pending,
|
||||
"status",
|
||||
)?)
|
||||
.bind(request.format_version)
|
||||
.bind(serialize_enum_text(&request.mode, "mode")?)
|
||||
.bind(request.created_at)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn finish_yaml_import_job(
|
||||
&self,
|
||||
job_id: &YamlImportJobId,
|
||||
completion: &YamlImportJobCompletion,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"update yaml_import_jobs
|
||||
set status = $1,
|
||||
result_operation_id = $2,
|
||||
result_version = $3,
|
||||
error_text = $4,
|
||||
finished_at = $5::timestamptz
|
||||
where id = $6",
|
||||
)
|
||||
.bind(serialize_enum_text(&completion.status, "status")?)
|
||||
.bind(
|
||||
completion
|
||||
.result_operation_id
|
||||
.as_ref()
|
||||
.map(|value| value.as_str()),
|
||||
)
|
||||
.bind(completion.result_version.map(to_db_version))
|
||||
.bind(completion.error_text.as_deref())
|
||||
.bind(&completion.finished_at)
|
||||
.bind(job_id.as_str())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::YamlImportJobNotFound {
|
||||
job_id: job_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_yaml_import_job(
|
||||
&self,
|
||||
job_id: &YamlImportJobId,
|
||||
) -> Result<Option<YamlImportJob>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
id,
|
||||
source_sample_id,
|
||||
status,
|
||||
format_version,
|
||||
mode,
|
||||
result_operation_id,
|
||||
result_version,
|
||||
error_text,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as finished_at
|
||||
from yaml_import_jobs
|
||||
where id = $1",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_yaml_import_job).transpose()
|
||||
}
|
||||
|
||||
async fn list_agent_bindings(
|
||||
&self,
|
||||
agent_id: &AgentId,
|
||||
|
||||
@@ -0,0 +1,446 @@
|
||||
use super::*;
|
||||
|
||||
impl PostgresRegistry {
|
||||
pub async fn create_invocation_log(
|
||||
&self,
|
||||
request: CreateInvocationLogRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
sqlx::query(
|
||||
"insert into invocation_logs (
|
||||
id,
|
||||
workspace_id,
|
||||
agent_id,
|
||||
operation_id,
|
||||
source,
|
||||
level,
|
||||
status,
|
||||
tool_name,
|
||||
message,
|
||||
request_id,
|
||||
status_code,
|
||||
duration_ms,
|
||||
error_kind,
|
||||
request_preview_json,
|
||||
response_preview_json,
|
||||
created_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16::timestamptz
|
||||
)",
|
||||
)
|
||||
.bind(request.log.id.as_str())
|
||||
.bind(request.log.workspace_id.as_str())
|
||||
.bind(request.log.agent_id.as_ref().map(|value| value.as_str()))
|
||||
.bind(request.log.operation_id.as_str())
|
||||
.bind(serialize_enum_text(&request.log.source, "source")?)
|
||||
.bind(serialize_enum_text(&request.log.level, "level")?)
|
||||
.bind(serialize_enum_text(&request.log.status, "status")?)
|
||||
.bind(&request.log.tool_name)
|
||||
.bind(&request.log.message)
|
||||
.bind(&request.log.request_id)
|
||||
.bind(request.log.status_code.map(i32::from))
|
||||
.bind(i64::try_from(request.log.duration_ms).map_err(|_| {
|
||||
RegistryError::InvalidNumericValue {
|
||||
field: "duration_ms",
|
||||
value: i64::MAX,
|
||||
}
|
||||
})?)
|
||||
.bind(&request.log.error_kind)
|
||||
.bind(Json(request.log.request_preview.clone()))
|
||||
.bind(Json(request.log.response_preview.clone()))
|
||||
.bind(&request.log.created_at)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_invocation_logs(
|
||||
&self,
|
||||
query: ListInvocationLogsQuery<'_>,
|
||||
) -> Result<Vec<InvocationLogRecord>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
l.id,
|
||||
l.workspace_id,
|
||||
l.agent_id,
|
||||
l.operation_id,
|
||||
l.source,
|
||||
l.level,
|
||||
l.status,
|
||||
l.tool_name,
|
||||
l.message,
|
||||
l.request_id,
|
||||
l.status_code,
|
||||
l.duration_ms,
|
||||
l.error_kind,
|
||||
l.request_preview_json,
|
||||
l.response_preview_json,
|
||||
to_char(l.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.MS\"Z\"') as created_at,
|
||||
o.name as operation_name,
|
||||
o.display_name as operation_display_name,
|
||||
a.slug as agent_slug,
|
||||
a.display_name as agent_display_name
|
||||
from invocation_logs l
|
||||
join operations o on o.id = l.operation_id
|
||||
left join agents a on a.id = l.agent_id
|
||||
where l.workspace_id = $1
|
||||
and ($2::text is null or l.level = $2)
|
||||
and ($3::text is null or l.source = $3)
|
||||
and ($4::text is null or l.operation_id = $4)
|
||||
and ($5::text is null or l.agent_id = $5)
|
||||
and ($6::timestamptz is null or l.created_at >= $6::timestamptz)
|
||||
and (
|
||||
$7::text is null
|
||||
or l.tool_name ilike '%' || $7 || '%'
|
||||
or l.message ilike '%' || $7 || '%'
|
||||
or o.name ilike '%' || $7 || '%'
|
||||
or o.display_name ilike '%' || $7 || '%'
|
||||
)
|
||||
order by l.created_at desc
|
||||
limit $8",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(
|
||||
query
|
||||
.level
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "level"))
|
||||
.transpose()?,
|
||||
)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.bind(query.operation_id.map(|value| value.as_str()))
|
||||
.bind(query.agent_id.map(|value| value.as_str()))
|
||||
.bind(query.created_after)
|
||||
.bind(query.search_text)
|
||||
.bind(i64::from(query.limit))
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_invocation_log_record).collect()
|
||||
}
|
||||
|
||||
pub async fn get_invocation_log(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
log_id: &InvocationLogId,
|
||||
) -> Result<Option<InvocationLogRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
l.id,
|
||||
l.workspace_id,
|
||||
l.agent_id,
|
||||
l.operation_id,
|
||||
l.source,
|
||||
l.level,
|
||||
l.status,
|
||||
l.tool_name,
|
||||
l.message,
|
||||
l.request_id,
|
||||
l.status_code,
|
||||
l.duration_ms,
|
||||
l.error_kind,
|
||||
l.request_preview_json,
|
||||
l.response_preview_json,
|
||||
to_char(l.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.MS\"Z\"') as created_at,
|
||||
o.name as operation_name,
|
||||
o.display_name as operation_display_name,
|
||||
a.slug as agent_slug,
|
||||
a.display_name as agent_display_name
|
||||
from invocation_logs l
|
||||
join operations o on o.id = l.operation_id
|
||||
left join agents a on a.id = l.agent_id
|
||||
where l.workspace_id = $1 and l.id = $2",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(log_id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_invocation_log_record).transpose()
|
||||
}
|
||||
|
||||
pub async fn summarize_usage(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<UsageSummary, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and created_at >= $2::timestamptz
|
||||
and ($3::text is null or source = $3)",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
let calls_total = row.try_get::<i64, _>("calls_total")?;
|
||||
let calls_ok = row.try_get::<i64, _>("calls_ok")?;
|
||||
let calls_error = row.try_get::<i64, _>("calls_error")?;
|
||||
let rollup = UsageRollup {
|
||||
workspace_id: query.workspace_id.clone(),
|
||||
agent_id: None,
|
||||
operation_id: None,
|
||||
period: query.period,
|
||||
calls_total: to_u64(calls_total, "calls_total")?,
|
||||
calls_ok: to_u64(calls_ok, "calls_ok")?,
|
||||
calls_error: to_u64(calls_error, "calls_error")?,
|
||||
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
|
||||
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
|
||||
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
|
||||
};
|
||||
|
||||
let success_rate = if rollup.calls_total == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(rollup.calls_ok as f64 / rollup.calls_total as f64) * 100.0
|
||||
};
|
||||
|
||||
Ok(UsageSummary {
|
||||
rollup,
|
||||
success_rate,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn list_usage_timeline(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<Vec<UsageTimelinePoint>, RegistryError> {
|
||||
let bucket = usage_bucket_sql(query.bucket);
|
||||
let sql = format!(
|
||||
"select
|
||||
to_char(date_trunc('{bucket}', created_at at time zone 'UTC'), 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as bucket_start,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and created_at >= $2::timestamptz
|
||||
and ($3::text is null or source = $3)
|
||||
group by 1
|
||||
order by 1 asc"
|
||||
);
|
||||
let rows = sqlx::query(&sql)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter()
|
||||
.map(|row| {
|
||||
Ok(UsageTimelinePoint {
|
||||
bucket_start: row.try_get("bucket_start")?,
|
||||
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
|
||||
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn list_usage_by_operation(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<Vec<UsageOperationBreakdown>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
o.id as operation_id,
|
||||
o.name as operation_name,
|
||||
o.display_name as operation_display_name,
|
||||
o.protocol,
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where l.status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by l.duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by l.duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by l.duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs l
|
||||
join operations o on o.id = l.operation_id
|
||||
where l.workspace_id = $1
|
||||
and l.created_at >= $2::timestamptz
|
||||
and ($3::text is null or l.source = $3)
|
||||
group by o.id, o.name, o.display_name, o.protocol
|
||||
order by calls_total desc, o.name asc",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_usage_operation_breakdown).collect()
|
||||
}
|
||||
|
||||
pub async fn get_usage_for_operation(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
operation_id: &OperationId,
|
||||
) -> Result<Option<UsageRollupRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and operation_id = $2
|
||||
and created_at >= $3::timestamptz
|
||||
and ($4::text is null or source = $4)",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(operation_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
let calls_total = to_u64(row.try_get::<i64, _>("calls_total")?, "calls_total")?;
|
||||
if calls_total == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(UsageRollupRecord {
|
||||
rollup: UsageRollup {
|
||||
workspace_id: query.workspace_id.clone(),
|
||||
agent_id: None,
|
||||
operation_id: Some(operation_id.clone()),
|
||||
period: query.period,
|
||||
calls_total,
|
||||
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
|
||||
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
|
||||
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
|
||||
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
|
||||
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn list_usage_by_agent(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
) -> Result<Vec<UsageAgentBreakdown>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
a.id as agent_id,
|
||||
a.slug as agent_slug,
|
||||
a.display_name as agent_display_name,
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where l.status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by l.duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by l.duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by l.duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs l
|
||||
join agents a on a.id = l.agent_id
|
||||
where l.workspace_id = $1
|
||||
and l.created_at >= $2::timestamptz
|
||||
and ($3::text is null or l.source = $3)
|
||||
group by a.id, a.slug, a.display_name
|
||||
order by calls_total desc, a.slug asc",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_usage_agent_breakdown).collect()
|
||||
}
|
||||
|
||||
pub async fn get_usage_for_agent(
|
||||
&self,
|
||||
query: UsageQuery<'_>,
|
||||
agent_id: &AgentId,
|
||||
) -> Result<Option<UsageRollupRecord>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
count(*)::bigint as calls_total,
|
||||
count(*) filter (where status = 'ok')::bigint as calls_ok,
|
||||
count(*) filter (where status = 'error')::bigint as calls_error,
|
||||
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
|
||||
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
|
||||
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
|
||||
from invocation_logs
|
||||
where workspace_id = $1
|
||||
and agent_id = $2
|
||||
and created_at >= $3::timestamptz
|
||||
and ($4::text is null or source = $4)",
|
||||
)
|
||||
.bind(query.workspace_id.as_str())
|
||||
.bind(agent_id.as_str())
|
||||
.bind(query.created_after)
|
||||
.bind(
|
||||
query
|
||||
.source
|
||||
.as_ref()
|
||||
.map(|value| serialize_enum_text(value, "source"))
|
||||
.transpose()?,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
let calls_total = to_u64(row.try_get::<i64, _>("calls_total")?, "calls_total")?;
|
||||
if calls_total == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(UsageRollupRecord {
|
||||
rollup: UsageRollup {
|
||||
workspace_id: query.workspace_id.clone(),
|
||||
agent_id: Some(agent_id.clone()),
|
||||
operation_id: None,
|
||||
period: query.period,
|
||||
calls_total,
|
||||
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
|
||||
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
|
||||
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
|
||||
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
|
||||
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
|
||||
},
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
use super::*;
|
||||
|
||||
impl PostgresRegistry {
|
||||
pub async fn create_yaml_import_job(
|
||||
&self,
|
||||
request: CreateYamlImportJobRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
sqlx::query(
|
||||
"insert into yaml_import_jobs (
|
||||
id,
|
||||
source_sample_id,
|
||||
status,
|
||||
format_version,
|
||||
mode,
|
||||
result_operation_id,
|
||||
result_version,
|
||||
error_text,
|
||||
created_at,
|
||||
finished_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, null, null, null, $6::timestamptz, null
|
||||
)",
|
||||
)
|
||||
.bind(request.id.as_str())
|
||||
.bind(request.source_sample_id.map(|value| value.as_str()))
|
||||
.bind(serialize_enum_text(
|
||||
&YamlImportJobStatus::Pending,
|
||||
"status",
|
||||
)?)
|
||||
.bind(request.format_version)
|
||||
.bind(serialize_enum_text(&request.mode, "mode")?)
|
||||
.bind(request.created_at)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn finish_yaml_import_job(
|
||||
&self,
|
||||
job_id: &YamlImportJobId,
|
||||
completion: &YamlImportJobCompletion,
|
||||
) -> Result<(), RegistryError> {
|
||||
let result = sqlx::query(
|
||||
"update yaml_import_jobs
|
||||
set status = $2,
|
||||
result_operation_id = $3,
|
||||
result_version = $4,
|
||||
error_text = $5,
|
||||
finished_at = $6::timestamptz
|
||||
where id = $1",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.bind(serialize_enum_text(&completion.status, "status")?)
|
||||
.bind(
|
||||
completion
|
||||
.result_operation_id
|
||||
.as_ref()
|
||||
.map(|value| value.as_str()),
|
||||
)
|
||||
.bind(completion.result_version.map(to_db_version))
|
||||
.bind(completion.error_text.as_deref())
|
||||
.bind(&completion.finished_at)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(RegistryError::YamlImportJobNotFound {
|
||||
job_id: job_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_yaml_import_job(
|
||||
&self,
|
||||
job_id: &YamlImportJobId,
|
||||
) -> Result<Option<YamlImportJob>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
id,
|
||||
source_sample_id,
|
||||
status,
|
||||
format_version,
|
||||
mode,
|
||||
result_operation_id,
|
||||
result_version,
|
||||
error_text,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as finished_at
|
||||
from yaml_import_jobs
|
||||
where id = $1",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_yaml_import_job).transpose()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user