registry: extract storage domain modules
This commit is contained in:
a.tolmachev
@@ -0,0 +1,478 @@
|
||||
use super::*;
|
||||
|
||||
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.as_deref())
|
||||
.bind(&request.session.created_at)
|
||||
.bind(request.session.closed_at.as_deref())
|
||||
.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,
|
||||
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
|
||||
to_char(last_poll_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_poll_at,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(closed_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as closed_at
|
||||
from stream_sessions
|
||||
where id = $1",
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_stream_session).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,
|
||||
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
|
||||
to_char(last_poll_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_poll_at,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(closed_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as 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: &str,
|
||||
) -> 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 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,
|
||||
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
|
||||
to_char(last_poll_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_poll_at,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(closed_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as closed_at
|
||||
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",
|
||||
)
|
||||
.bind(filter.workspace_id.as_str())
|
||||
.bind(filter.agent_id.map(|value| value.as_str()))
|
||||
.bind(filter.operation_id.map(|value| value.as_str()))
|
||||
.bind(status)
|
||||
.bind(mode)
|
||||
.bind(limit)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(map_stream_session)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let total = items.len() as u64;
|
||||
|
||||
Ok(Page { items, total })
|
||||
}
|
||||
|
||||
pub async fn delete_expired_stream_sessions(&self, now: &str) -> 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,
|
||||
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
|
||||
)",
|
||||
)
|
||||
.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.as_deref())
|
||||
.bind(&request.job.created_at)
|
||||
.bind(&request.job.updated_at)
|
||||
.bind(request.job.finished_at.as_deref())
|
||||
.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,
|
||||
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
|
||||
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as finished_at
|
||||
from async_jobs
|
||||
where id = $1",
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_async_job).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,
|
||||
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
|
||||
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as 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 cancel_async_job(&self, id: &AsyncJobId, now: &str) -> 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,
|
||||
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
|
||||
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as finished_at
|
||||
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",
|
||||
)
|
||||
.bind(filter.workspace_id.as_str())
|
||||
.bind(filter.agent_id.map(|value| value.as_str()))
|
||||
.bind(filter.operation_id.map(|value| value.as_str()))
|
||||
.bind(status)
|
||||
.bind(limit)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(map_async_job)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let total = items.len() as u64;
|
||||
|
||||
Ok(Page { items, total })
|
||||
}
|
||||
|
||||
pub async fn delete_expired_async_jobs(&self, now: &str) -> 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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user