diff --git a/crates/crank-registry/src/postgres/api_key.rs b/crates/crank-registry/src/postgres/api_key.rs new file mode 100644 index 0000000..fcbcd83 --- /dev/null +++ b/crates/crank-registry/src/postgres/api_key.rs @@ -0,0 +1,180 @@ +use super::*; + +impl PostgresRegistry { + pub async fn list_platform_api_keys( + &self, + workspace_id: &WorkspaceId, + ) -> Result, 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, 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(()) + } +} diff --git a/crates/crank-registry/src/postgres/mod.rs b/crates/crank-registry/src/postgres/mod.rs index 92f60b2..703d040 100644 --- a/crates/crank-registry/src/postgres/mod.rs +++ b/crates/crank-registry/src/postgres/mod.rs @@ -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, 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, 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, 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, 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 { - 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::("calls_total")?; - let calls_ok = row.try_get::("calls_ok")?; - let calls_error = row.try_get::("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::("p50_ms")?, "p50_ms")?, - p95_ms: to_u64(row.try_get::("p95_ms")?, "p95_ms")?, - p99_ms: to_u64(row.try_get::("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, 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::("calls_ok")?, "calls_ok")?, - calls_error: to_u64(row.try_get::("calls_error")?, "calls_error")?, - }) - }) - .collect() - } - - pub async fn list_usage_by_operation( - &self, - query: UsageQuery<'_>, - ) -> Result, 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, 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::("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::("calls_ok")?, "calls_ok")?, - calls_error: to_u64(row.try_get::("calls_error")?, "calls_error")?, - p50_ms: to_u64(row.try_get::("p50_ms")?, "p50_ms")?, - p95_ms: to_u64(row.try_get::("p95_ms")?, "p95_ms")?, - p99_ms: to_u64(row.try_get::("p99_ms")?, "p99_ms")?, - }, - })) - } - - pub async fn list_usage_by_agent( - &self, - query: UsageQuery<'_>, - ) -> Result, 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, 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::("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::("calls_ok")?, "calls_ok")?, - calls_error: to_u64(row.try_get::("calls_error")?, "calls_error")?, - p50_ms: to_u64(row.try_get::("p50_ms")?, "p50_ms")?, - p95_ms: to_u64(row.try_get::("p95_ms")?, "p95_ms")?, - p99_ms: to_u64(row.try_get::("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, 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, diff --git a/crates/crank-registry/src/postgres/observability.rs b/crates/crank-registry/src/postgres/observability.rs new file mode 100644 index 0000000..7903070 --- /dev/null +++ b/crates/crank-registry/src/postgres/observability.rs @@ -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, 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, 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 { + 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::("calls_total")?; + let calls_ok = row.try_get::("calls_ok")?; + let calls_error = row.try_get::("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::("p50_ms")?, "p50_ms")?, + p95_ms: to_u64(row.try_get::("p95_ms")?, "p95_ms")?, + p99_ms: to_u64(row.try_get::("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, 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::("calls_ok")?, "calls_ok")?, + calls_error: to_u64(row.try_get::("calls_error")?, "calls_error")?, + }) + }) + .collect() + } + + pub async fn list_usage_by_operation( + &self, + query: UsageQuery<'_>, + ) -> Result, 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, 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::("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::("calls_ok")?, "calls_ok")?, + calls_error: to_u64(row.try_get::("calls_error")?, "calls_error")?, + p50_ms: to_u64(row.try_get::("p50_ms")?, "p50_ms")?, + p95_ms: to_u64(row.try_get::("p95_ms")?, "p95_ms")?, + p99_ms: to_u64(row.try_get::("p99_ms")?, "p99_ms")?, + }, + })) + } + + pub async fn list_usage_by_agent( + &self, + query: UsageQuery<'_>, + ) -> Result, 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, 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::("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::("calls_ok")?, "calls_ok")?, + calls_error: to_u64(row.try_get::("calls_error")?, "calls_error")?, + p50_ms: to_u64(row.try_get::("p50_ms")?, "p50_ms")?, + p95_ms: to_u64(row.try_get::("p95_ms")?, "p95_ms")?, + p99_ms: to_u64(row.try_get::("p99_ms")?, "p99_ms")?, + }, + })) + } +} diff --git a/crates/crank-registry/src/postgres/yaml_import.rs b/crates/crank-registry/src/postgres/yaml_import.rs new file mode 100644 index 0000000..5a3b3f3 --- /dev/null +++ b/crates/crank-registry/src/postgres/yaml_import.rs @@ -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, 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() + } +}