trending/OpenShell
1
0
Fork 0
mirror of https://github.com/NVIDIA/OpenShell.git synced 2026-11-01 07:53:42 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
33517b32bd3fdacce0ce8a0d862253aa10f0504f
OpenShell/crates/openshell-server/src/persistence/postgres.rs
T
Bryce Wilkinson 33517b32bd feat(server): make the database pool ceiling configurable 
Both stores hardcoded their pool ceiling: 10 connections for Postgres, 5
for on-disk SQLite. One pool is shared by every database-backed RPC, so
that number is also the gateway's ceiling on concurrent database work.
Once every connection is checked out, callers queue on acquire and sqlx
logs "time to acquire exceeded slow threshold"; sandbox creates then time
out and are retried, which adds load rather than shedding it.

The right ceiling is deployment-specific — it depends on how many gateway
replicas share the database and what max_connections the server itself
allows — so it cannot be a single number baked into the binary.

Expose it on the same config surface as the rest of the gateway's
settings: --db-max-connections, OPENSHELL_DB_MAX_CONNECTIONS, and the
TOML key database_max_connections, resolved in that precedence order and
rendered by the Helm chart from server.dbMaxConnections. Omitting it
keeps each backend's previous value, so existing deployments do not move.

Values below 1 are rejected rather than silently replaced by the default:
a zero pool would block every acquire, and quietly ignoring a typo would
reproduce the ceiling the operator is trying to lift. An in-memory SQLite
database stays pinned to one connection, since the database lives inside
that connection.

Refs: #2561
Signed-off-by: Bryce Wilkinson <22760097+bjw123@users.noreply.github.com>
2026-08-20 10:03:46 +00:00

1245 lines
38 KiB
Rust
Raw Blame History

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use super::{
DraftChunkRecord, ObjectRecord, PersistenceError, PersistenceResult, PolicyRecord,
WriteCondition, WriteResult, current_time_ms, map_db_error, map_migrate_error,
};
use crate::policy_store::{
AtomicPolicyRevisionWrite, draft_chunk_payload_from_record, draft_chunk_record_from_parts,
policy_payload_from_record, policy_record_for_atomic_write, policy_record_from_parts,
project_policy_revision_onto_sandbox,
};
use openshell_core::SetResourceVersion;
use openshell_core::proto::Sandbox;
use prost::Message;
use sqlx::postgres::PgPoolOptions;
use sqlx::{Connection, PgPool, Row};
static POSTGRES_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations/postgres");
use super::{DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE};
#[derive(Debug, Clone)]
pub struct PostgresStore {
pool: PgPool,
}
/// Pool ceiling used when the operator has not configured one.
///
/// A single shared pool serves every database-backed RPC, so this value is
/// also the gateway's ceiling on concurrent database work. It is high enough
/// for a single-replica gateway against a default-sized server; fleets that
/// outgrow it raise it through `[openshell.gateway] database_max_connections`.
pub const DEFAULT_MAX_CONNECTIONS: u32 = 10;
impl PostgresStore {
/// Connect and size the pool, falling back to [`DEFAULT_MAX_CONNECTIONS`].
pub async fn connect(url: &str, max_connections: Option<u32>) -> PersistenceResult<Self> {
let max_connections = max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS);
tracing::info!(max_connections, "sizing Postgres connection pool");
let pool = PgPoolOptions::new()
.max_connections(max_connections)
.connect(url)
.await
.map_err(|e| map_db_error(&e))?;
Ok(Self { pool })
}
pub async fn migrate(&self) -> PersistenceResult<()> {
POSTGRES_MIGRATOR
.run(&self.pool)
.await
.map_err(|e| map_migrate_error(&e))
}
/// Verify the database is reachable by acquiring a pooled connection
/// and issuing a protocol-level ping.
pub async fn ping(&self) -> PersistenceResult<()> {
let mut conn = self.pool.acquire().await.map_err(|e| map_db_error(&e))?;
conn.ping().await.map_err(|e| map_db_error(&e))
}
/// Test support only: close the underlying connection pool.
///
/// Do not call from runtime code; this tears down the active pool.
#[cfg(any(test, feature = "test-support"))]
pub async fn close(&self) {
self.pool.close().await;
}
pub async fn put(
&self,
object_type: &str,
id: &str,
name: &str,
workspace: &str,
payload: &[u8],
labels: Option<&str>,
) -> PersistenceResult<()> {
let now_ms = current_time_ms();
let labels_jsonb: Option<serde_json::Value> = labels
.map(serde_json::from_str)
.transpose()
.map_err(|e| PersistenceError::Encode(format!("invalid labels JSON: {e}")))?;
sqlx::query(
r"
INSERT INTO objects (object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels)
VALUES ($1, $2, $3, $4, $5, $6, $6, COALESCE($7, '{}'::jsonb))
ON CONFLICT (object_type, workspace, name) WHERE name IS NOT NULL DO UPDATE SET
payload = EXCLUDED.payload,
updated_at_ms = EXCLUDED.updated_at_ms,
labels = EXCLUDED.labels
",
)
.bind(object_type)
.bind(id)
.bind(name)
.bind(workspace)
.bind(payload)
.bind(now_ms)
.bind(labels_jsonb)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub async fn put_if(
&self,
object_type: &str,
id: &str,
name: &str,
workspace: &str,
payload: &[u8],
labels: Option<&str>,
condition: WriteCondition,
) -> PersistenceResult<WriteResult> {
let now_ms = current_time_ms();
let labels_jsonb: Option<serde_json::Value> = labels
.map(serde_json::from_str)
.transpose()
.map_err(|e| PersistenceError::Encode(format!("invalid labels JSON: {e}")))?;
match condition {
WriteCondition::MustCreate => {
// Insert only - fail if object exists
let row = sqlx::query(
r"
INSERT INTO objects (object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version)
VALUES ($1, $2, $3, $4, $5, $6, $6, COALESCE($7, '{}'::jsonb), 1)
RETURNING resource_version, created_at_ms, updated_at_ms
",
)
.bind(object_type)
.bind(id)
.bind(name)
.bind(workspace)
.bind(payload)
.bind(now_ms)
.bind(labels_jsonb)
.fetch_one(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1);
Ok(WriteResult {
resource_version: resource_version_i64.max(1).cast_unsigned(),
created_at_ms: row.get("created_at_ms"),
updated_at_ms: row.get("updated_at_ms"),
})
}
WriteCondition::MatchResourceVersion(expected_version) => {
// Update with version check using RETURNING
let row_result = sqlx::query(
r"
UPDATE objects
SET payload = $4, labels = COALESCE($5, '{}'::jsonb), updated_at_ms = $6, resource_version = resource_version + 1
WHERE object_type = $1 AND id = $2 AND resource_version = $3
RETURNING resource_version, created_at_ms, updated_at_ms
",
)
.bind(object_type)
.bind(id)
.bind(i64::try_from(expected_version).unwrap_or(i64::MAX))
.bind(payload)
.bind(labels_jsonb)
.bind(now_ms)
.fetch_optional(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
if let Some(row) = row_result {
let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1);
Ok(WriteResult {
resource_version: resource_version_i64.max(1).cast_unsigned(),
created_at_ms: row.get("created_at_ms"),
updated_at_ms: row.get("updated_at_ms"),
})
} else {
// The version-matched UPDATE matched no row. Distinguish a
// version mismatch (row present, different version) from an
// absent row (deleted / never existed). Both are CAS
// precondition failures, so report them as typed `Conflict`
// rather than a backend-dependent error string: absent rows
// carry `current_resource_version: None`.
let existing = self.get(object_type, id).await?;
Err(PersistenceError::Conflict {
current_resource_version: existing.map(|record| record.resource_version),
})
}
}
WriteCondition::Unconditional => {
// Unconditional upsert by name
let row = sqlx::query(
r"
INSERT INTO objects (object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version)
VALUES ($1, $2, $3, $4, $5, $6, $6, COALESCE($7, '{}'::jsonb), 1)
ON CONFLICT (object_type, workspace, name) WHERE name IS NOT NULL DO UPDATE SET
payload = EXCLUDED.payload,
updated_at_ms = EXCLUDED.updated_at_ms,
labels = EXCLUDED.labels,
resource_version = objects.resource_version + 1
RETURNING resource_version, created_at_ms, updated_at_ms
",
)
.bind(object_type)
.bind(id)
.bind(name)
.bind(workspace)
.bind(payload)
.bind(now_ms)
.bind(labels_jsonb)
.fetch_one(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1);
Ok(WriteResult {
resource_version: resource_version_i64.max(1).cast_unsigned(),
created_at_ms: row.get("created_at_ms"),
updated_at_ms: row.get("updated_at_ms"),
})
}
}
}
pub async fn delete_if(
&self,
object_type: &str,
id: &str,
expected_resource_version: u64,
) -> PersistenceResult<bool> {
let result = sqlx::query(
r"
DELETE FROM objects
WHERE object_type = $1 AND id = $2 AND resource_version = $3
",
)
.bind(object_type)
.bind(id)
.bind(i64::try_from(expected_resource_version).unwrap_or(i64::MAX))
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
if result.rows_affected() > 0 {
Ok(true)
} else {
// Check if object exists to distinguish NotFound from Conflict
let existing = self.get(object_type, id).await?;
if let Some(record) = existing {
Err(PersistenceError::Conflict {
current_resource_version: Some(record.resource_version),
})
} else {
Ok(false)
}
}
}
#[allow(clippy::too_many_arguments)]
pub async fn put_scoped(
&self,
object_type: &str,
id: &str,
name: &str,
workspace: &str,
scope: &str,
payload: &[u8],
labels: Option<&str>,
) -> PersistenceResult<()> {
let now_ms = current_time_ms();
let labels_jsonb: Option<serde_json::Value> = labels
.map(serde_json::from_str)
.transpose()
.map_err(|e| PersistenceError::Encode(format!("invalid labels JSON: {e}")))?;
sqlx::query(
r"
INSERT INTO objects (object_type, id, name, workspace, scope, payload, created_at_ms, updated_at_ms, labels, resource_version)
VALUES ($1, $2, $3, $4, $5, $6, $7, $7, COALESCE($8, '{}'::jsonb), 1)
ON CONFLICT (object_type, workspace, name) WHERE name IS NOT NULL DO UPDATE SET
scope = EXCLUDED.scope,
payload = EXCLUDED.payload,
updated_at_ms = EXCLUDED.updated_at_ms,
labels = EXCLUDED.labels,
resource_version = objects.resource_version + 1
",
)
.bind(object_type)
.bind(id)
.bind(name)
.bind(workspace)
.bind(scope)
.bind(payload)
.bind(now_ms)
.bind(labels_jsonb)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(())
}
pub async fn get(
&self,
object_type: &str,
id: &str,
) -> PersistenceResult<Option<ObjectRecord>> {
let row = sqlx::query(
r"
SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version
FROM objects
WHERE object_type = $1 AND id = $2
",
)
.bind(object_type)
.bind(id)
.fetch_optional(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(row.map(row_to_object_record))
}
pub async fn get_by_name(
&self,
object_type: &str,
workspace: &str,
name: &str,
) -> PersistenceResult<Option<ObjectRecord>> {
let row = sqlx::query(
r"
SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version
FROM objects
WHERE object_type = $1 AND workspace = $2 AND name = $3
",
)
.bind(object_type)
.bind(workspace)
.bind(name)
.fetch_optional(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(row.map(row_to_object_record))
}
pub async fn delete(&self, object_type: &str, id: &str) -> PersistenceResult<bool> {
let result = sqlx::query("DELETE FROM objects WHERE object_type = $1 AND id = $2")
.bind(object_type)
.bind(id)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected() > 0)
}
pub async fn count_in_workspace(
&self,
object_type: &str,
workspace: &str,
) -> PersistenceResult<u64> {
let row: (i64,) = sqlx::query_as(
"SELECT COUNT(*) FROM objects WHERE object_type = $1 AND workspace = $2",
)
.bind(object_type)
.bind(workspace)
.fetch_one(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(u64::try_from(row.0).unwrap_or(0))
}
pub async fn delete_all_in_workspace(
&self,
object_type: &str,
workspace: &str,
) -> PersistenceResult<u64> {
let result = sqlx::query("DELETE FROM objects WHERE object_type = $1 AND workspace = $2")
.bind(object_type)
.bind(workspace)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected())
}
pub async fn delete_by_scope(&self, object_type: &str, scope: &str) -> PersistenceResult<u64> {
let result = sqlx::query("DELETE FROM objects WHERE object_type = $1 AND scope = $2")
.bind(object_type)
.bind(scope)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected())
}
pub async fn delete_by_name(
&self,
object_type: &str,
workspace: &str,
name: &str,
) -> PersistenceResult<bool> {
let result = sqlx::query(
"DELETE FROM objects WHERE object_type = $1 AND workspace = $2 AND name = $3",
)
.bind(object_type)
.bind(workspace)
.bind(name)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected() > 0)
}
pub async fn list(
&self,
object_type: &str,
workspace: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<ObjectRecord>> {
let rows = sqlx::query(
r"
SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version
FROM objects
WHERE object_type = $1 AND workspace = $2
ORDER BY created_at_ms ASC, name ASC
LIMIT $3 OFFSET $4
",
)
.bind(object_type)
.bind(workspace)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(rows.into_iter().map(row_to_object_record).collect())
}
pub async fn list_by_type(
&self,
object_type: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<ObjectRecord>> {
let rows = sqlx::query(
r"
SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version
FROM objects
WHERE object_type = $1
ORDER BY created_at_ms ASC, name ASC, workspace ASC, id ASC
LIMIT $2 OFFSET $3
",
)
.bind(object_type)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(rows.into_iter().map(row_to_object_record).collect())
}
pub async fn list_with_membership(
&self,
object_type: &str,
member_type: &str,
member_name: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<ObjectRecord>> {
let rows = sqlx::query(
r"
SELECT w.object_type, w.id, w.name, w.workspace, w.payload,
w.created_at_ms, w.updated_at_ms, w.labels, w.resource_version
FROM objects w
WHERE w.object_type = $1 AND w.workspace = ''
AND EXISTS (
SELECT 1 FROM objects m
WHERE m.object_type = $2
AND m.workspace = w.name
AND m.name = $3
)
ORDER BY w.created_at_ms ASC, w.name ASC
LIMIT $4 OFFSET $5
",
)
.bind(object_type)
.bind(member_type)
.bind(member_name)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(rows.into_iter().map(row_to_object_record).collect())
}
pub async fn list_with_membership_and_selector(
&self,
object_type: &str,
member_type: &str,
member_name: &str,
label_selector: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<ObjectRecord>> {
use super::parse_label_selector;
let required_labels = parse_label_selector(label_selector)?;
let labels_jsonb = serde_json::to_value(&required_labels)
.map_err(|e| PersistenceError::Encode(format!("failed to serialize labels: {e}")))?;
let rows = sqlx::query(
r"
SELECT w.object_type, w.id, w.name, w.workspace, w.payload,
w.created_at_ms, w.updated_at_ms, w.labels, w.resource_version
FROM objects w
WHERE w.object_type = $1 AND w.workspace = ''
AND EXISTS (
SELECT 1 FROM objects m
WHERE m.object_type = $2
AND m.workspace = w.name
AND m.name = $3
)
AND w.labels @> $4
ORDER BY w.created_at_ms ASC, w.name ASC
LIMIT $5 OFFSET $6
",
)
.bind(object_type)
.bind(member_type)
.bind(member_name)
.bind(&labels_jsonb)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(rows.into_iter().map(row_to_object_record).collect())
}
pub async fn list_by_scope(
&self,
object_type: &str,
scope: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<ObjectRecord>> {
let rows = sqlx::query(
r"
SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version
FROM objects
WHERE object_type = $1 AND scope = $2
ORDER BY created_at_ms ASC, name ASC
LIMIT $3 OFFSET $4
",
)
.bind(object_type)
.bind(scope)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(rows.into_iter().map(row_to_object_record).collect())
}
pub async fn list_with_selector(
&self,
object_type: &str,
workspace: &str,
label_selector: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<ObjectRecord>> {
use super::parse_label_selector;
let required_labels = parse_label_selector(label_selector)?;
let labels_jsonb = serde_json::to_value(&required_labels)
.map_err(|e| PersistenceError::Encode(format!("failed to serialize labels: {e}")))?;
let rows = sqlx::query(
r"
SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version
FROM objects
WHERE object_type = $1 AND workspace = $2 AND labels @> $3
ORDER BY created_at_ms ASC, name ASC
LIMIT $4 OFFSET $5
",
)
.bind(object_type)
.bind(workspace)
.bind(&labels_jsonb)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(rows.into_iter().map(row_to_object_record).collect())
}
pub async fn list_all_with_selector(
&self,
object_type: &str,
label_selector: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<ObjectRecord>> {
use super::parse_label_selector;
let required_labels = parse_label_selector(label_selector)?;
let labels_jsonb = serde_json::to_value(&required_labels)
.map_err(|e| PersistenceError::Encode(format!("failed to serialize labels: {e}")))?;
let rows = sqlx::query(
r"
SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version
FROM objects
WHERE object_type = $1 AND labels @> $2
ORDER BY created_at_ms ASC, name ASC, workspace ASC, id ASC
LIMIT $3 OFFSET $4
",
)
.bind(object_type)
.bind(&labels_jsonb)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(rows.into_iter().map(row_to_object_record).collect())
}
pub async fn put_policy_revision(
&self,
id: &str,
sandbox_id: &str,
workspace: &str,
version: i64,
payload: &[u8],
hash: &str,
) -> PersistenceResult<()> {
let now_ms = current_time_ms();
let record = PolicyRecord {
id: id.to_string(),
sandbox_id: sandbox_id.to_string(),
version,
policy_payload: payload.to_vec(),
policy_hash: hash.to_string(),
status: "pending".to_string(),
load_error: None,
created_at_ms: now_ms,
loaded_at_ms: None,
provenance: std::collections::HashMap::default(),
};
let wrapped_payload = policy_payload_from_record(&record)?;
sqlx::query(
r"
INSERT INTO objects (
object_type, id, scope, version, status, payload, created_at_ms, updated_at_ms, workspace
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $7, $8)
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(id)
.bind(sandbox_id)
.bind(version)
.bind("pending")
.bind(wrapped_payload)
.bind(now_ms)
.bind(workspace)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(())
}
pub async fn put_policy_revision_atomic(
&self,
write: &AtomicPolicyRevisionWrite,
) -> PersistenceResult<Sandbox> {
let now_ms = current_time_ms();
let record = policy_record_for_atomic_write(write, now_ms);
let wrapped_payload = policy_payload_from_record(&record)?;
let mut tx = self.pool.begin().await.map_err(|e| map_db_error(&e))?;
let row = sqlx::query(
r"
SELECT payload, resource_version
FROM objects
WHERE object_type = 'sandbox' AND id = $1
FOR UPDATE
",
)
.bind(&write.sandbox_id)
.fetch_optional(&mut *tx)
.await
.map_err(|e| map_db_error(&e))?
.ok_or_else(|| {
PersistenceError::Database(format!("sandbox object {} not found", write.sandbox_id))
})?;
let sandbox_payload: Vec<u8> = row.get("payload");
let current_version: i64 = row.try_get("resource_version").unwrap_or(1);
let current_version = current_version.max(1).cast_unsigned();
let (mut sandbox, sandbox_changed) =
project_policy_revision_onto_sandbox(write, &sandbox_payload, current_version)?;
let resulting_version = if sandbox_changed {
let result = sqlx::query(
r"
UPDATE objects
SET payload = $2, updated_at_ms = $3, resource_version = resource_version + 1
WHERE object_type = 'sandbox' AND id = $1 AND resource_version = $4
",
)
.bind(&write.sandbox_id)
.bind(sandbox.encode_to_vec())
.bind(now_ms)
.bind(i64::try_from(current_version).unwrap_or(i64::MAX))
.execute(&mut *tx)
.await
.map_err(|e| map_db_error(&e))?;
if result.rows_affected() != 1 {
return Err(PersistenceError::Conflict {
current_resource_version: Some(current_version),
});
}
current_version.saturating_add(1)
} else {
current_version
};
sqlx::query(
r"
INSERT INTO objects (
object_type, id, scope, version, status, payload, created_at_ms, updated_at_ms, workspace
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $7, $8)
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(&write.id)
.bind(&write.sandbox_id)
.bind(write.version)
.bind("pending")
.bind(wrapped_payload)
.bind(now_ms)
.bind(&write.workspace)
.execute(&mut *tx)
.await
.map_err(|e| map_db_error(&e))?;
sqlx::query(
r"
UPDATE objects
SET status = 'superseded', updated_at_ms = $4
WHERE object_type = $1
AND scope = $2
AND version < $3
AND status IN ('pending', 'loaded')
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(&write.sandbox_id)
.bind(write.version)
.bind(now_ms)
.execute(&mut *tx)
.await
.map_err(|e| map_db_error(&e))?;
tx.commit().await.map_err(|e| map_db_error(&e))?;
sandbox.set_resource_version(resulting_version);
Ok(sandbox)
}
pub async fn get_latest_policy(
&self,
sandbox_id: &str,
) -> PersistenceResult<Option<PolicyRecord>> {
let row = sqlx::query(
r"
SELECT id, scope, version, status, payload, created_at_ms
FROM objects
WHERE object_type = $1 AND scope = $2
ORDER BY version DESC, created_at_ms DESC
LIMIT 1
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(sandbox_id)
.fetch_optional(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
row.map(row_to_policy_record).transpose()
}
pub async fn get_latest_loaded_policy(
&self,
sandbox_id: &str,
) -> PersistenceResult<Option<PolicyRecord>> {
let row = sqlx::query(
r"
SELECT id, scope, version, status, payload, created_at_ms
FROM objects
WHERE object_type = $1 AND scope = $2 AND status = 'loaded'
ORDER BY version DESC, created_at_ms DESC
LIMIT 1
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(sandbox_id)
.fetch_optional(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
row.map(row_to_policy_record).transpose()
}
pub async fn get_policy_by_version(
&self,
sandbox_id: &str,
version: i64,
) -> PersistenceResult<Option<PolicyRecord>> {
let row = sqlx::query(
r"
SELECT id, scope, version, status, payload, created_at_ms
FROM objects
WHERE object_type = $1 AND scope = $2 AND version = $3
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(sandbox_id)
.bind(version)
.fetch_optional(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
row.map(row_to_policy_record).transpose()
}
pub async fn list_policies(
&self,
sandbox_id: &str,
limit: u32,
offset: u32,
) -> PersistenceResult<Vec<PolicyRecord>> {
let rows = sqlx::query(
r"
SELECT id, scope, version, status, payload, created_at_ms
FROM objects
WHERE object_type = $1 AND scope = $2
ORDER BY version DESC, created_at_ms DESC
LIMIT $3 OFFSET $4
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(sandbox_id)
.bind(i64::from(limit))
.bind(i64::from(offset))
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
rows.into_iter().map(row_to_policy_record).collect()
}
pub async fn update_policy_status(
&self,
sandbox_id: &str,
version: i64,
status: &str,
load_error: Option<&str>,
loaded_at_ms: Option<i64>,
) -> PersistenceResult<bool> {
let Some(mut record) = self.get_policy_by_version(sandbox_id, version).await? else {
return Ok(false);
};
record.status = status.to_string();
record.load_error = load_error.map(ToOwned::to_owned);
record.loaded_at_ms = loaded_at_ms;
let payload = policy_payload_from_record(&record)?;
let now_ms = current_time_ms();
let result = sqlx::query(
r"
UPDATE objects
SET status = $4, payload = $5, updated_at_ms = $6
WHERE object_type = $1 AND scope = $2 AND version = $3
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(sandbox_id)
.bind(version)
.bind(status)
.bind(payload)
.bind(now_ms)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected() > 0)
}
pub async fn supersede_older_policies(
&self,
sandbox_id: &str,
before_version: i64,
) -> PersistenceResult<u64> {
let now_ms = current_time_ms();
let result = sqlx::query(
r"
UPDATE objects
SET status = 'superseded', updated_at_ms = $4
WHERE object_type = $1
AND scope = $2
AND version < $3
AND status IN ('pending', 'loaded')
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(sandbox_id)
.bind(before_version)
.bind(now_ms)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected())
}
pub async fn put_draft_chunk(
&self,
chunk: &DraftChunkRecord,
dedup_key: Option<&str>,
workspace: &str,
) -> PersistenceResult<String> {
let payload = draft_chunk_payload_from_record(chunk)?;
// RETURNING id gives the row's effective id whether INSERT inserted
// a fresh row or ON CONFLICT updated an existing one. See the
// matching sqlite path for the rationale.
let row = sqlx::query(
r"
INSERT INTO objects (
object_type, id, scope, status, dedup_key, hit_count, payload, created_at_ms, updated_at_ms, workspace
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
ON CONFLICT (object_type, scope, dedup_key) WHERE dedup_key IS NOT NULL DO UPDATE SET
hit_count = objects.hit_count + EXCLUDED.hit_count,
updated_at_ms = EXCLUDED.updated_at_ms
RETURNING id
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(&chunk.id)
.bind(&chunk.sandbox_id)
.bind(&chunk.status)
.bind(dedup_key)
.bind(i64::from(chunk.hit_count))
.bind(payload)
.bind(chunk.first_seen_ms)
.bind(chunk.last_seen_ms)
.bind(workspace)
.fetch_one(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(row.get::<String, _>("id"))
}
pub async fn get_draft_chunk(&self, id: &str) -> PersistenceResult<Option<DraftChunkRecord>> {
let row = sqlx::query(
r"
SELECT id, scope, status, hit_count, payload, created_at_ms, updated_at_ms
FROM objects
WHERE object_type = $1 AND id = $2
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(id)
.fetch_optional(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
row.map(row_to_draft_chunk_record).transpose()
}
pub async fn list_draft_chunks(
&self,
sandbox_id: &str,
status_filter: Option<&str>,
) -> PersistenceResult<Vec<DraftChunkRecord>> {
let rows = if let Some(status) = status_filter {
sqlx::query(
r"
SELECT id, scope, status, hit_count, payload, created_at_ms, updated_at_ms
FROM objects
WHERE object_type = $1 AND scope = $2 AND status = $3
ORDER BY created_at_ms DESC
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(sandbox_id)
.bind(status)
.fetch_all(&self.pool)
.await
} else {
sqlx::query(
r"
SELECT id, scope, status, hit_count, payload, created_at_ms, updated_at_ms
FROM objects
WHERE object_type = $1 AND scope = $2
ORDER BY created_at_ms DESC
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(sandbox_id)
.fetch_all(&self.pool)
.await
}
.map_err(|e| map_db_error(&e))?;
rows.into_iter().map(row_to_draft_chunk_record).collect()
}
pub async fn update_draft_chunk_status(
&self,
id: &str,
status: &str,
decided_at_ms: Option<i64>,
rejection_reason: Option<&str>,
) -> PersistenceResult<bool> {
let Some(mut record) = self.get_draft_chunk(id).await? else {
return Ok(false);
};
record.status = status.to_string();
record.decided_at_ms = decided_at_ms;
record.last_seen_ms = current_time_ms();
if let Some(reason) = rejection_reason {
record.rejection_reason = reason.to_string();
}
let payload = draft_chunk_payload_from_record(&record)?;
let result = sqlx::query(
r"
UPDATE objects
SET status = $3, payload = $4, updated_at_ms = $5
WHERE object_type = $1 AND id = $2
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(id)
.bind(status)
.bind(payload)
.bind(record.last_seen_ms)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected() > 0)
}
pub async fn conditionally_reject_draft_chunk(
&self,
id: &str,
decided_at_ms: i64,
rejection_reason: &str,
) -> PersistenceResult<bool> {
let Some(mut record) = self.get_draft_chunk(id).await? else {
return Ok(false);
};
if record.status != "pending" {
return Ok(false);
}
record.status = "rejected".to_string();
record.decided_at_ms = Some(decided_at_ms);
record.rejection_reason = rejection_reason.to_string();
record.last_seen_ms = current_time_ms();
let payload = draft_chunk_payload_from_record(&record)?;
let result = sqlx::query(
r"
UPDATE objects
SET status = $3, payload = $4, updated_at_ms = $5
WHERE object_type = $1 AND id = $2 AND status = 'pending'
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(id)
.bind("rejected")
.bind(payload)
.bind(record.last_seen_ms)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected() > 0)
}
pub async fn update_draft_chunk_rule(
&self,
id: &str,
proposed_rule: &[u8],
) -> PersistenceResult<bool> {
let Some(mut record) = self.get_draft_chunk(id).await? else {
return Ok(false);
};
if record.status != "pending" {
return Ok(false);
}
record.proposed_rule = proposed_rule.to_vec();
record.last_seen_ms = current_time_ms();
let payload = draft_chunk_payload_from_record(&record)?;
let result = sqlx::query(
r"
UPDATE objects
SET payload = $3, updated_at_ms = $4
WHERE object_type = $1 AND id = $2 AND status = 'pending'
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(id)
.bind(payload)
.bind(record.last_seen_ms)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected() > 0)
}
pub async fn delete_draft_chunks(
&self,
sandbox_id: &str,
status: &str,
) -> PersistenceResult<u64> {
let result = sqlx::query(
r"
DELETE FROM objects
WHERE object_type = $1 AND scope = $2 AND status = $3
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(sandbox_id)
.bind(status)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(result.rows_affected())
}
pub async fn get_draft_version(&self, sandbox_id: &str) -> PersistenceResult<i64> {
let rows = sqlx::query(
r"
SELECT payload
FROM objects
WHERE object_type = $1 AND scope = $2
",
)
.bind(DRAFT_CHUNK_OBJECT_TYPE)
.bind(sandbox_id)
.fetch_all(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
let mut max_version = 0_i64;
for row in rows {
let payload: Vec<u8> = row.get("payload");
let wrapper = draft_chunk_record_from_parts(
String::new(),
sandbox_id.to_string(),
String::new(),
0,
&payload,
0,
0,
)?;
max_version = max_version.max(wrapper.draft_version);
}
Ok(max_version)
}
}
fn row_to_object_record(row: sqlx::postgres::PgRow) -> ObjectRecord {
let labels_jsonb: Option<serde_json::Value> = row.get("labels");
let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1);
ObjectRecord {
object_type: row.get("object_type"),
id: row.get("id"),
name: row.get("name"),
workspace: row.try_get("workspace").unwrap_or_default(),
payload: row.get("payload"),
created_at_ms: row.get("created_at_ms"),
updated_at_ms: row.get("updated_at_ms"),
labels: labels_jsonb.map(|value| value.to_string()),
resource_version: resource_version_i64.max(1).cast_unsigned(),
}
}
fn row_to_policy_record(row: sqlx::postgres::PgRow) -> PersistenceResult<PolicyRecord> {
let id: String = row.get("id");
let sandbox_id: String = row.get("scope");
let version: i64 = row.get("version");
let status: String = row.get("status");
let payload: Vec<u8> = row.get("payload");
let created_at_ms: i64 = row.get("created_at_ms");
policy_record_from_parts(id, sandbox_id, version, status, &payload, created_at_ms)
}
fn row_to_draft_chunk_record(row: sqlx::postgres::PgRow) -> PersistenceResult<DraftChunkRecord> {
let id: String = row.get("id");
let sandbox_id: String = row.get("scope");
let status: String = row.get("status");
let hit_count: i64 = row.get("hit_count");
let payload: Vec<u8> = row.get("payload");
let created_at_ms: i64 = row.get("created_at_ms");
let updated_at_ms: i64 = row.get("updated_at_ms");
draft_chunk_record_from_parts(
id,
sandbox_id,
status,
hit_count,
&payload,
created_at_ms,
updated_at_ms,
)
}
Reference in New Issue View Git Blame Copy Permalink