registry: add sqlx checks for workspace and stream reads
This commit is contained in:
a.tolmachev
@@ -54,7 +54,7 @@ impl PostgresRegistry {
|
||||
&self,
|
||||
id: &StreamSessionId,
|
||||
) -> Result<Option<StreamSession>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
let row = sqlx::query!(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
@@ -65,18 +65,35 @@ impl PostgresRegistry {
|
||||
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(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(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",
|
||||
id.as_str(),
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_stream_session).transpose()
|
||||
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(
|
||||
@@ -196,7 +213,7 @@ impl PostgresRegistry {
|
||||
.transpose()?;
|
||||
let limit = i64::from(filter.limit);
|
||||
|
||||
let rows = sqlx::query(
|
||||
let rows = sqlx::query!(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
@@ -207,9 +224,9 @@ impl PostgresRegistry {
|
||||
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(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(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
|
||||
@@ -219,20 +236,36 @@ impl PostgresRegistry {
|
||||
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,
|
||||
)
|
||||
.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<_>, _>>()?;
|
||||
.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 })
|
||||
@@ -295,7 +328,7 @@ impl PostgresRegistry {
|
||||
&self,
|
||||
id: &AsyncJobId,
|
||||
) -> Result<Option<AsyncJobHandle>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
let row = sqlx::query!(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
@@ -306,17 +339,33 @@ impl PostgresRegistry {
|
||||
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(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",
|
||||
id.as_str(),
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_async_job).transpose()
|
||||
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,
|
||||
created_at: row.created_at,
|
||||
updated_at: row.updated_at,
|
||||
finished_at: row.finished_at,
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
pub async fn update_async_job_status(
|
||||
@@ -424,7 +473,7 @@ impl PostgresRegistry {
|
||||
.transpose()?;
|
||||
let limit = i64::from(filter.limit);
|
||||
|
||||
let rows = sqlx::query(
|
||||
let rows = sqlx::query!(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
@@ -435,8 +484,8 @@ impl PostgresRegistry {
|
||||
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(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
|
||||
@@ -445,19 +494,34 @@ impl PostgresRegistry {
|
||||
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,
|
||||
)
|
||||
.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<_>, _>>()?;
|
||||
.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,
|
||||
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 })
|
||||
|
||||
Reference in New Issue
Block a user