a.tolmachev
635 lines
21 KiB
Rust
635 lines
21 KiB
Rust
use super::*;
|
|
use time::OffsetDateTime;
|
|
|
|
impl PostgresRegistry {
|
|
pub async fn create_stream_session(
|
|
&self,
|
|
request: CreateStreamSessionRequest<'_>,
|
|
) -> Result<StreamSession, RegistryError> {
|
|
sqlx::query(
|
|
"insert into stream_sessions (
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
protocol,
|
|
mode,
|
|
status,
|
|
cursor_json,
|
|
state_json,
|
|
expires_at,
|
|
last_poll_at,
|
|
created_at,
|
|
closed_at
|
|
) values (
|
|
$1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9::jsonb, $10::timestamptz,
|
|
$11::timestamptz, $12::timestamptz, $13::timestamptz
|
|
)",
|
|
)
|
|
.bind(request.session.id.as_str())
|
|
.bind(request.session.workspace_id.as_str())
|
|
.bind(
|
|
request
|
|
.session
|
|
.agent_id
|
|
.as_ref()
|
|
.map(|value| value.as_str()),
|
|
)
|
|
.bind(request.session.operation_id.as_str())
|
|
.bind(serialize_enum_text(&request.session.protocol, "protocol")?)
|
|
.bind(serialize_enum_text(&request.session.mode, "mode")?)
|
|
.bind(serialize_enum_text(&request.session.status, "status")?)
|
|
.bind(request.session.cursor.clone().map(Json))
|
|
.bind(Json(request.session.state.clone()))
|
|
.bind(request.session.expires_at)
|
|
.bind(request.session.last_poll_at)
|
|
.bind(request.session.created_at)
|
|
.bind(request.session.closed_at)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(request.session.clone())
|
|
}
|
|
|
|
pub async fn get_stream_session(
|
|
&self,
|
|
id: &StreamSessionId,
|
|
) -> Result<Option<StreamSession>, RegistryError> {
|
|
let row = sqlx::query!(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
protocol,
|
|
mode,
|
|
status,
|
|
cursor_json,
|
|
state_json,
|
|
expires_at as \"expires_at!: OffsetDateTime\",
|
|
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
|
created_at as \"created_at!: OffsetDateTime\",
|
|
closed_at as \"closed_at: OffsetDateTime\"
|
|
from stream_sessions
|
|
where id = $1",
|
|
id.as_str(),
|
|
)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(|row| {
|
|
Ok(StreamSession {
|
|
id: StreamSessionId::new(row.id),
|
|
workspace_id: WorkspaceId::new(row.workspace_id),
|
|
agent_id: row.agent_id.map(AgentId::new),
|
|
operation_id: OperationId::new(row.operation_id),
|
|
protocol: deserialize_enum_text(&row.protocol, "protocol")?,
|
|
mode: deserialize_enum_text(&row.mode, "mode")?,
|
|
status: deserialize_enum_text(&row.status, "status")?,
|
|
cursor: row.cursor_json,
|
|
state: row.state_json,
|
|
expires_at: row.expires_at,
|
|
last_poll_at: row.last_poll_at,
|
|
created_at: row.created_at,
|
|
closed_at: row.closed_at,
|
|
})
|
|
})
|
|
.transpose()
|
|
}
|
|
|
|
pub async fn update_stream_session_state(
|
|
&self,
|
|
request: UpdateStreamSessionStateRequest<'_>,
|
|
) -> Result<StreamSession, RegistryError> {
|
|
validate_stream_session_transition(
|
|
request.session_id,
|
|
request.current_status,
|
|
request.next_status,
|
|
)?;
|
|
|
|
let row = sqlx::query(
|
|
"update stream_sessions
|
|
set status = $3,
|
|
cursor_json = $4::jsonb,
|
|
state_json = $5::jsonb,
|
|
expires_at = coalesce($6::timestamptz, expires_at),
|
|
last_poll_at = coalesce($7::timestamptz, last_poll_at),
|
|
closed_at = case
|
|
when $8::timestamptz is null then closed_at
|
|
else $8::timestamptz
|
|
end
|
|
where id = $1
|
|
and status = $2
|
|
returning
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
protocol,
|
|
mode,
|
|
status,
|
|
cursor_json,
|
|
state_json,
|
|
expires_at,
|
|
last_poll_at,
|
|
created_at,
|
|
closed_at",
|
|
)
|
|
.bind(request.session_id.as_str())
|
|
.bind(serialize_enum_text(&request.current_status, "status")?)
|
|
.bind(serialize_enum_text(&request.next_status, "status")?)
|
|
.bind(request.cursor.cloned().map(Json))
|
|
.bind(Json(request.state.clone()))
|
|
.bind(request.expires_at)
|
|
.bind(request.last_poll_at)
|
|
.bind(request.closed_at)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
match row.as_ref() {
|
|
Some(row) => map_stream_session(row),
|
|
None => {
|
|
if self.get_stream_session(request.session_id).await?.is_some() {
|
|
Err(RegistryError::InvalidStreamSessionTransition {
|
|
session_id: request.session_id.as_str().to_owned(),
|
|
from: serialize_enum_text(&request.current_status, "status")?,
|
|
to: serialize_enum_text(&request.next_status, "status")?,
|
|
})
|
|
} else {
|
|
Err(RegistryError::StreamSessionNotFound {
|
|
session_id: request.session_id.as_str().to_owned(),
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn close_stream_session(
|
|
&self,
|
|
id: &StreamSessionId,
|
|
now: &OffsetDateTime,
|
|
) -> Result<(), RegistryError> {
|
|
let result = sqlx::query(
|
|
"update stream_sessions
|
|
set status = 'stopped',
|
|
last_poll_at = $2::timestamptz,
|
|
closed_at = $2::timestamptz
|
|
where id = $1
|
|
and status in ('created', 'running')",
|
|
)
|
|
.bind(id.as_str())
|
|
.bind(now)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
if result.rows_affected() > 0 {
|
|
return Ok(());
|
|
}
|
|
|
|
match self.get_stream_session(id).await? {
|
|
Some(existing) => Err(RegistryError::InvalidStreamSessionTransition {
|
|
session_id: id.as_str().to_owned(),
|
|
from: serialize_enum_text(&existing.status, "status")?,
|
|
to: "stopped".to_owned(),
|
|
}),
|
|
None => Err(RegistryError::StreamSessionNotFound {
|
|
session_id: id.as_str().to_owned(),
|
|
}),
|
|
}
|
|
}
|
|
|
|
pub async fn touch_stream_session_poll(
|
|
&self,
|
|
id: &StreamSessionId,
|
|
now: &OffsetDateTime,
|
|
) -> Result<StreamSession, RegistryError> {
|
|
let row = sqlx::query(
|
|
"update stream_sessions
|
|
set last_poll_at = $2::timestamptz
|
|
where id = $1
|
|
returning
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
protocol,
|
|
mode,
|
|
status,
|
|
cursor_json,
|
|
state_json,
|
|
expires_at,
|
|
last_poll_at,
|
|
created_at,
|
|
closed_at",
|
|
)
|
|
.bind(id.as_str())
|
|
.bind(now)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
match row.as_ref() {
|
|
Some(row) => map_stream_session(row),
|
|
None => Err(RegistryError::StreamSessionNotFound {
|
|
session_id: id.as_str().to_owned(),
|
|
}),
|
|
}
|
|
}
|
|
|
|
pub async fn list_stream_sessions(
|
|
&self,
|
|
filter: StreamSessionFilter<'_>,
|
|
) -> Result<Page<StreamSession>, RegistryError> {
|
|
let status = filter
|
|
.status
|
|
.as_ref()
|
|
.map(|value| serialize_enum_text(value, "status"))
|
|
.transpose()?;
|
|
let mode = filter
|
|
.mode
|
|
.as_ref()
|
|
.map(|value| serialize_enum_text(value, "mode"))
|
|
.transpose()?;
|
|
let limit = i64::from(filter.limit);
|
|
|
|
let rows = sqlx::query!(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
protocol,
|
|
mode,
|
|
status,
|
|
cursor_json,
|
|
state_json,
|
|
expires_at as \"expires_at!: OffsetDateTime\",
|
|
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
|
created_at as \"created_at!: OffsetDateTime\",
|
|
closed_at as \"closed_at: OffsetDateTime\"
|
|
from stream_sessions
|
|
where workspace_id = $1
|
|
and ($2::text is null or agent_id = $2)
|
|
and ($3::text is null or operation_id = $3)
|
|
and ($4::text is null or status = $4)
|
|
and ($5::text is null or mode = $5)
|
|
order by created_at desc
|
|
limit $6",
|
|
filter.workspace_id.as_str(),
|
|
filter.agent_id.map(|value| value.as_str()),
|
|
filter.operation_id.map(|value| value.as_str()),
|
|
status,
|
|
mode,
|
|
limit,
|
|
)
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
|
|
let items = rows
|
|
.into_iter()
|
|
.map(|row| {
|
|
Ok(StreamSession {
|
|
id: StreamSessionId::new(row.id),
|
|
workspace_id: WorkspaceId::new(row.workspace_id),
|
|
agent_id: row.agent_id.map(AgentId::new),
|
|
operation_id: OperationId::new(row.operation_id),
|
|
protocol: deserialize_enum_text(&row.protocol, "protocol")?,
|
|
mode: deserialize_enum_text(&row.mode, "mode")?,
|
|
status: deserialize_enum_text(&row.status, "status")?,
|
|
cursor: row.cursor_json,
|
|
state: row.state_json,
|
|
expires_at: row.expires_at,
|
|
last_poll_at: row.last_poll_at,
|
|
created_at: row.created_at,
|
|
closed_at: row.closed_at,
|
|
})
|
|
})
|
|
.collect::<Result<Vec<_>, RegistryError>>()?;
|
|
let total = items.len() as u64;
|
|
|
|
Ok(Page { items, total })
|
|
}
|
|
|
|
pub async fn delete_expired_stream_sessions(
|
|
&self,
|
|
now: &OffsetDateTime,
|
|
) -> Result<u64, RegistryError> {
|
|
let result = sqlx::query(
|
|
"delete from stream_sessions
|
|
where expires_at <= $1::timestamptz",
|
|
)
|
|
.bind(now)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result.rows_affected())
|
|
}
|
|
|
|
pub async fn create_async_job(
|
|
&self,
|
|
request: CreateAsyncJobRequest<'_>,
|
|
) -> Result<AsyncJobHandle, RegistryError> {
|
|
sqlx::query(
|
|
"insert into async_jobs (
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
status,
|
|
progress_json,
|
|
result_json,
|
|
error_json,
|
|
expires_at,
|
|
last_poll_at,
|
|
created_at,
|
|
updated_at,
|
|
finished_at
|
|
) values (
|
|
$1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8::jsonb, $9::timestamptz,
|
|
$10::timestamptz, $11::timestamptz, $12::timestamptz, $13::timestamptz
|
|
)",
|
|
)
|
|
.bind(request.job.id.as_str())
|
|
.bind(request.job.workspace_id.as_str())
|
|
.bind(request.job.agent_id.as_ref().map(|value| value.as_str()))
|
|
.bind(request.job.operation_id.as_str())
|
|
.bind(serialize_enum_text(&request.job.status, "status")?)
|
|
.bind(Json(request.job.progress.clone()))
|
|
.bind(request.job.result.clone().map(Json))
|
|
.bind(request.job.error.clone().map(Json))
|
|
.bind(request.job.expires_at)
|
|
.bind(request.job.last_poll_at)
|
|
.bind(request.job.created_at)
|
|
.bind(request.job.updated_at)
|
|
.bind(request.job.finished_at)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(request.job.clone())
|
|
}
|
|
|
|
pub async fn get_async_job(
|
|
&self,
|
|
id: &AsyncJobId,
|
|
) -> Result<Option<AsyncJobHandle>, RegistryError> {
|
|
let row = sqlx::query!(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
status,
|
|
progress_json,
|
|
result_json,
|
|
error_json,
|
|
expires_at as \"expires_at: OffsetDateTime\",
|
|
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
|
created_at as \"created_at!: OffsetDateTime\",
|
|
updated_at as \"updated_at!: OffsetDateTime\",
|
|
finished_at as \"finished_at: OffsetDateTime\"
|
|
from async_jobs
|
|
where id = $1",
|
|
id.as_str(),
|
|
)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(|row| {
|
|
Ok(AsyncJobHandle {
|
|
id: AsyncJobId::new(row.id),
|
|
workspace_id: WorkspaceId::new(row.workspace_id),
|
|
agent_id: row.agent_id.map(AgentId::new),
|
|
operation_id: OperationId::new(row.operation_id),
|
|
status: deserialize_enum_text(&row.status, "status")?,
|
|
progress: row.progress_json,
|
|
result: row.result_json,
|
|
error: row.error_json,
|
|
expires_at: row.expires_at,
|
|
last_poll_at: row.last_poll_at,
|
|
created_at: row.created_at,
|
|
updated_at: row.updated_at,
|
|
finished_at: row.finished_at,
|
|
})
|
|
})
|
|
.transpose()
|
|
}
|
|
|
|
pub async fn update_async_job_status(
|
|
&self,
|
|
request: UpdateAsyncJobStatusRequest<'_>,
|
|
) -> Result<AsyncJobHandle, RegistryError> {
|
|
validate_async_job_transition(request.job_id, request.current_status, request.next_status)?;
|
|
|
|
let row = sqlx::query(
|
|
"update async_jobs
|
|
set status = $3,
|
|
progress_json = $4::jsonb,
|
|
result_json = $5::jsonb,
|
|
error_json = $6::jsonb,
|
|
expires_at = coalesce($7::timestamptz, expires_at),
|
|
updated_at = $8::timestamptz,
|
|
finished_at = case
|
|
when $9::timestamptz is null then finished_at
|
|
else $9::timestamptz
|
|
end
|
|
where id = $1
|
|
and status = $2
|
|
returning
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
status,
|
|
progress_json,
|
|
result_json,
|
|
error_json,
|
|
expires_at,
|
|
last_poll_at,
|
|
created_at,
|
|
updated_at,
|
|
finished_at",
|
|
)
|
|
.bind(request.job_id.as_str())
|
|
.bind(serialize_enum_text(&request.current_status, "status")?)
|
|
.bind(serialize_enum_text(&request.next_status, "status")?)
|
|
.bind(Json(request.progress.clone()))
|
|
.bind(request.result.cloned().map(Json))
|
|
.bind(request.error.cloned().map(Json))
|
|
.bind(request.expires_at)
|
|
.bind(request.updated_at)
|
|
.bind(request.finished_at)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
match row.as_ref() {
|
|
Some(row) => map_async_job(row),
|
|
None => {
|
|
if self.get_async_job(request.job_id).await?.is_some() {
|
|
Err(RegistryError::InvalidAsyncJobTransition {
|
|
job_id: request.job_id.as_str().to_owned(),
|
|
from: serialize_enum_text(&request.current_status, "status")?,
|
|
to: serialize_enum_text(&request.next_status, "status")?,
|
|
})
|
|
} else {
|
|
Err(RegistryError::AsyncJobNotFound {
|
|
job_id: request.job_id.as_str().to_owned(),
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn touch_async_job_poll(
|
|
&self,
|
|
id: &AsyncJobId,
|
|
now: &OffsetDateTime,
|
|
) -> Result<AsyncJobHandle, RegistryError> {
|
|
let row = sqlx::query(
|
|
"update async_jobs
|
|
set last_poll_at = $2::timestamptz
|
|
where id = $1
|
|
returning
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
status,
|
|
progress_json,
|
|
result_json,
|
|
error_json,
|
|
expires_at,
|
|
last_poll_at,
|
|
created_at,
|
|
updated_at,
|
|
finished_at",
|
|
)
|
|
.bind(id.as_str())
|
|
.bind(now)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
match row.as_ref() {
|
|
Some(row) => map_async_job(row),
|
|
None => Err(RegistryError::AsyncJobNotFound {
|
|
job_id: id.as_str().to_owned(),
|
|
}),
|
|
}
|
|
}
|
|
|
|
pub async fn cancel_async_job(
|
|
&self,
|
|
id: &AsyncJobId,
|
|
now: &OffsetDateTime,
|
|
) -> Result<(), RegistryError> {
|
|
let result = sqlx::query(
|
|
"update async_jobs
|
|
set status = 'cancelled',
|
|
updated_at = $2::timestamptz,
|
|
finished_at = $2::timestamptz
|
|
where id = $1
|
|
and status in ('created', 'running')",
|
|
)
|
|
.bind(id.as_str())
|
|
.bind(now)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
if result.rows_affected() > 0 {
|
|
return Ok(());
|
|
}
|
|
|
|
match self.get_async_job(id).await? {
|
|
Some(existing) => Err(RegistryError::InvalidAsyncJobTransition {
|
|
job_id: id.as_str().to_owned(),
|
|
from: serialize_enum_text(&existing.status, "status")?,
|
|
to: "cancelled".to_owned(),
|
|
}),
|
|
None => Err(RegistryError::AsyncJobNotFound {
|
|
job_id: id.as_str().to_owned(),
|
|
}),
|
|
}
|
|
}
|
|
|
|
pub async fn list_async_jobs(
|
|
&self,
|
|
filter: AsyncJobFilter<'_>,
|
|
) -> Result<Page<AsyncJobHandle>, RegistryError> {
|
|
let status = filter
|
|
.status
|
|
.as_ref()
|
|
.map(|value| serialize_enum_text(value, "status"))
|
|
.transpose()?;
|
|
let limit = i64::from(filter.limit);
|
|
|
|
let rows = sqlx::query!(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
status,
|
|
progress_json,
|
|
result_json,
|
|
error_json,
|
|
expires_at as \"expires_at: OffsetDateTime\",
|
|
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
|
created_at as \"created_at!: OffsetDateTime\",
|
|
updated_at as \"updated_at!: OffsetDateTime\",
|
|
finished_at as \"finished_at: OffsetDateTime\"
|
|
from async_jobs
|
|
where workspace_id = $1
|
|
and ($2::text is null or agent_id = $2)
|
|
and ($3::text is null or operation_id = $3)
|
|
and ($4::text is null or status = $4)
|
|
order by updated_at desc
|
|
limit $5",
|
|
filter.workspace_id.as_str(),
|
|
filter.agent_id.map(|value| value.as_str()),
|
|
filter.operation_id.map(|value| value.as_str()),
|
|
status,
|
|
limit,
|
|
)
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
|
|
let items = rows
|
|
.into_iter()
|
|
.map(|row| {
|
|
Ok(AsyncJobHandle {
|
|
id: AsyncJobId::new(row.id),
|
|
workspace_id: WorkspaceId::new(row.workspace_id),
|
|
agent_id: row.agent_id.map(AgentId::new),
|
|
operation_id: OperationId::new(row.operation_id),
|
|
status: deserialize_enum_text(&row.status, "status")?,
|
|
progress: row.progress_json,
|
|
result: row.result_json,
|
|
error: row.error_json,
|
|
expires_at: row.expires_at,
|
|
last_poll_at: row.last_poll_at,
|
|
created_at: row.created_at,
|
|
updated_at: row.updated_at,
|
|
finished_at: row.finished_at,
|
|
})
|
|
})
|
|
.collect::<Result<Vec<_>, RegistryError>>()?;
|
|
let total = items.len() as u64;
|
|
|
|
Ok(Page { items, total })
|
|
}
|
|
|
|
pub async fn delete_expired_async_jobs(
|
|
&self,
|
|
now: &OffsetDateTime,
|
|
) -> Result<u64, RegistryError> {
|
|
let result = sqlx::query(
|
|
"delete from async_jobs
|
|
where expires_at is not null
|
|
and expires_at <= $1::timestamptz",
|
|
)
|
|
.bind(now)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result.rows_affected())
|
|
}
|
|
}
|