refactor(server): move policy persistence into policy_store

Author: johntmyers

crates/openshell-server/src/grpc/policy.rs

@@ -12,6 +12,7 @@
 
 use crate::ServerState;
 use crate::persistence::{DraftChunkRecord, ObjectId, ObjectName, PolicyRecord, Store};
+use crate::policy_store::PolicyStoreExt;
 use openshell_core::proto::policy_merge_operation;
 use openshell_core::proto::setting_value;
 use openshell_core::proto::{

crates/openshell-server/src/lib.rs

@@ -27,6 +27,7 @@ mod http;
 mod inference;
 mod multiplex;
 mod persistence;
+pub(crate) mod policy_store;
 mod sandbox_index;
 mod sandbox_watch;
 mod ssh_tunnel;

crates/openshell-server/src/persistence/mod.rs

@@ -10,13 +10,7 @@ pub use openshell_core::proto::{
     StoredDraftChunk as DraftChunkRecord, StoredPolicyRevision as PolicyRecord,
 };
 
-use openshell_core::{
-    Error as CoreError, Result as CoreResult,
-    proto::{
-        DraftChunkPayload, NetworkPolicyRule, PolicyRevisionPayload,
-        SandboxPolicy as ProtoSandboxPolicy,
-    },
-};
+use openshell_core::{Error as CoreError, Result as CoreResult};
 use prost::Message;
 use rand::Rng;
 use std::time::{SystemTime, UNIX_EPOCH};
@@ -228,122 +222,6 @@ impl Store {
         }
     }
 
-    // -----------------------------------------------------------------------
-    // Policy revision operations
-    // -----------------------------------------------------------------------
-    /// Insert a new policy revision.
-    pub async fn put_policy_revision(
-        &self,
-        id: &str,
-        sandbox_id: &str,
-        version: i64,
-        payload: &[u8],
-        hash: &str,
-    ) -> PersistenceResult<()> {
-        match self {
-            Self::Postgres(store) => {
-                store
-                    .put_policy_revision(id, sandbox_id, version, payload, hash)
-                    .await
-            }
-            Self::Sqlite(store) => {
-                store
-                    .put_policy_revision(id, sandbox_id, version, payload, hash)
-                    .await
-            }
-        }
-    }
-
-    /// Get the latest policy revision for a sandbox (by highest version, any status).
-    pub async fn get_latest_policy(
-        &self,
-        sandbox_id: &str,
-    ) -> PersistenceResult<Option<PolicyRecord>> {
-        match self {
-            Self::Postgres(store) => store.get_latest_policy(sandbox_id).await,
-            Self::Sqlite(store) => store.get_latest_policy(sandbox_id).await,
-        }
-    }
-
-    /// Get the latest loaded policy revision for a sandbox.
-    pub async fn get_latest_loaded_policy(
-        &self,
-        sandbox_id: &str,
-    ) -> PersistenceResult<Option<PolicyRecord>> {
-        match self {
-            Self::Postgres(store) => store.get_latest_loaded_policy(sandbox_id).await,
-            Self::Sqlite(store) => store.get_latest_loaded_policy(sandbox_id).await,
-        }
-    }
-
-    /// Get a specific policy revision by sandbox id and version.
-    pub async fn get_policy_by_version(
-        &self,
-        sandbox_id: &str,
-        version: i64,
-    ) -> PersistenceResult<Option<PolicyRecord>> {
-        match self {
-            Self::Postgres(store) => store.get_policy_by_version(sandbox_id, version).await,
-            Self::Sqlite(store) => store.get_policy_by_version(sandbox_id, version).await,
-        }
-    }
-
-    /// List policy revisions for a sandbox, ordered by version descending.
-    pub async fn list_policies(
-        &self,
-        sandbox_id: &str,
-        limit: u32,
-        offset: u32,
-    ) -> PersistenceResult<Vec<PolicyRecord>> {
-        match self {
-            Self::Postgres(store) => store.list_policies(sandbox_id, limit, offset).await,
-            Self::Sqlite(store) => store.list_policies(sandbox_id, limit, offset).await,
-        }
-    }
-
-    /// Update the status of a policy revision.
-    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> {
-        match self {
-            Self::Postgres(store) => {
-                store
-                    .update_policy_status(sandbox_id, version, status, load_error, loaded_at_ms)
-                    .await
-            }
-            Self::Sqlite(store) => {
-                store
-                    .update_policy_status(sandbox_id, version, status, load_error, loaded_at_ms)
-                    .await
-            }
-        }
-    }
-
-    /// Mark all pending and loaded policy revisions older than `before_version` as superseded.
-    pub async fn supersede_older_policies(
-        &self,
-        sandbox_id: &str,
-        before_version: i64,
-    ) -> PersistenceResult<u64> {
-        match self {
-            Self::Postgres(store) => {
-                store
-                    .supersede_older_policies(sandbox_id, before_version)
-                    .await
-            }
-            Self::Sqlite(store) => {
-                store
-                    .supersede_older_policies(sandbox_id, before_version)
-                    .await
-            }
-        }
-    }
-
     // -----------------------------------------------------------------------
     // Generic protobuf message helpers
     // -----------------------------------------------------------------------
@@ -402,191 +280,6 @@ impl Store {
             .map(Some)
             .map_err(|e| PersistenceError::Decode(format!("protobuf decode error: {e}")))
     }
-
-    /// Insert or merge a new draft policy chunk.
-    pub async fn put_draft_chunk(&self, chunk: &DraftChunkRecord) -> PersistenceResult<()> {
-        match self {
-            Self::Postgres(store) => store.put_draft_chunk(chunk).await,
-            Self::Sqlite(store) => store.put_draft_chunk(chunk).await,
-        }
-    }
-
-    /// Fetch a single draft chunk by id.
-    pub async fn get_draft_chunk(&self, id: &str) -> PersistenceResult<Option<DraftChunkRecord>> {
-        match self {
-            Self::Postgres(store) => store.get_draft_chunk(id).await,
-            Self::Sqlite(store) => store.get_draft_chunk(id).await,
-        }
-    }
-
-    /// List draft chunks for a sandbox, optionally filtered by status.
-    pub async fn list_draft_chunks(
-        &self,
-        sandbox_id: &str,
-        status_filter: Option<&str>,
-    ) -> PersistenceResult<Vec<DraftChunkRecord>> {
-        match self {
-            Self::Postgres(store) => store.list_draft_chunks(sandbox_id, status_filter).await,
-            Self::Sqlite(store) => store.list_draft_chunks(sandbox_id, status_filter).await,
-        }
-    }
-
-    /// Update the status of a draft chunk.
-    pub async fn update_draft_chunk_status(
-        &self,
-        id: &str,
-        status: &str,
-        decided_at_ms: Option<i64>,
-    ) -> PersistenceResult<bool> {
-        match self {
-            Self::Postgres(store) => {
-                store
-                    .update_draft_chunk_status(id, status, decided_at_ms)
-                    .await
-            }
-            Self::Sqlite(store) => {
-                store
-                    .update_draft_chunk_status(id, status, decided_at_ms)
-                    .await
-            }
-        }
-    }
-
-    /// Update the proposed rule on a pending draft chunk.
-    pub async fn update_draft_chunk_rule(
-        &self,
-        id: &str,
-        proposed_rule: &[u8],
-    ) -> PersistenceResult<bool> {
-        match self {
-            Self::Postgres(store) => store.update_draft_chunk_rule(id, proposed_rule).await,
-            Self::Sqlite(store) => store.update_draft_chunk_rule(id, proposed_rule).await,
-        }
-    }
-
-    /// Delete all draft chunks for a sandbox with a given status.
-    pub async fn delete_draft_chunks(
-        &self,
-        sandbox_id: &str,
-        status: &str,
-    ) -> PersistenceResult<u64> {
-        match self {
-            Self::Postgres(store) => store.delete_draft_chunks(sandbox_id, status).await,
-            Self::Sqlite(store) => store.delete_draft_chunks(sandbox_id, status).await,
-        }
-    }
-
-    /// Get the current maximum draft version for a sandbox.
-    pub async fn get_draft_version(&self, sandbox_id: &str) -> PersistenceResult<i64> {
-        match self {
-            Self::Postgres(store) => store.get_draft_version(sandbox_id).await,
-            Self::Sqlite(store) => store.get_draft_version(sandbox_id).await,
-        }
-    }
-}
-
-pub(crate) fn policy_payload_from_record(record: &PolicyRecord) -> PersistenceResult<Vec<u8>> {
-    let policy = ProtoSandboxPolicy::decode(record.policy_payload.as_slice())
-        .map_err(|e| PersistenceError::Decode(format!("decode policy payload failed: {e}")))?;
-    Ok(PolicyRevisionPayload {
-        policy: Some(policy),
-        hash: record.policy_hash.clone(),
-        load_error: record.load_error.clone().unwrap_or_default(),
-        loaded_at_ms: record.loaded_at_ms.unwrap_or(0),
-    }
-    .encode_to_vec())
-}
-
-pub(crate) fn policy_record_from_parts(
-    id: String,
-    sandbox_id: String,
-    version: i64,
-    status: String,
-    payload: &[u8],
-    created_at_ms: i64,
-) -> PersistenceResult<PolicyRecord> {
-    let wrapper = PolicyRevisionPayload::decode(payload)
-        .map_err(|e| PersistenceError::Decode(format!("decode policy wrapper failed: {e}")))?;
-    let policy = wrapper
-        .policy
-        .ok_or_else(|| PersistenceError::Decode("policy wrapper missing policy".to_string()))?;
-    Ok(PolicyRecord {
-        id,
-        sandbox_id,
-        version,
-        policy_payload: policy.encode_to_vec(),
-        policy_hash: wrapper.hash,
-        status,
-        load_error: if wrapper.load_error.is_empty() {
-            None
-        } else {
-            Some(wrapper.load_error)
-        },
-        created_at_ms,
-        loaded_at_ms: (wrapper.loaded_at_ms > 0).then_some(wrapper.loaded_at_ms),
-    })
-}
-
-pub(crate) fn draft_chunk_payload_from_record(
-    chunk: &DraftChunkRecord,
-) -> PersistenceResult<Vec<u8>> {
-    let proposed_rule = if chunk.proposed_rule.is_empty() {
-        None
-    } else {
-        Some(
-            NetworkPolicyRule::decode(chunk.proposed_rule.as_slice())
-                .map_err(|e| PersistenceError::Decode(format!("decode draft rule failed: {e}")))?,
-        )
-    };
-    Ok(DraftChunkPayload {
-        rule_name: chunk.rule_name.clone(),
-        proposed_rule,
-        rationale: chunk.rationale.clone(),
-        security_notes: chunk.security_notes.clone(),
-        confidence: chunk.confidence as f32,
-        decided_at_ms: chunk.decided_at_ms.unwrap_or(0),
-        host: chunk.host.clone(),
-        port: chunk.port,
-        binary: chunk.binary.clone(),
-        draft_version: chunk.draft_version,
-    }
-    .encode_to_vec())
-}
-
-pub(crate) fn draft_chunk_record_from_parts(
-    id: String,
-    sandbox_id: String,
-    status: String,
-    hit_count: i64,
-    payload: &[u8],
-    created_at_ms: i64,
-    updated_at_ms: i64,
-) -> PersistenceResult<DraftChunkRecord> {
-    let wrapper = DraftChunkPayload::decode(payload)
-        .map_err(|e| PersistenceError::Decode(format!("decode draft chunk wrapper failed: {e}")))?;
-    let proposed_rule = wrapper
-        .proposed_rule
-        .map(|rule| rule.encode_to_vec())
-        .unwrap_or_default();
-    Ok(DraftChunkRecord {
-        id,
-        sandbox_id,
-        draft_version: wrapper.draft_version,
-        status,
-        rule_name: wrapper.rule_name,
-        proposed_rule,
-        rationale: wrapper.rationale,
-        security_notes: wrapper.security_notes,
-        confidence: f64::from(wrapper.confidence),
-        created_at_ms,
-        decided_at_ms: (wrapper.decided_at_ms > 0).then_some(wrapper.decided_at_ms),
-        host: wrapper.host,
-        port: wrapper.port,
-        binary: wrapper.binary,
-        hit_count: i32::try_from(hit_count).unwrap_or(i32::MAX),
-        first_seen_ms: created_at_ms,
-        last_seen_ms: updated_at_ms,
-    })
 }
 
 pub fn current_time_ms() -> PersistenceResult<i64> {

crates/openshell-server/src/persistence/postgres.rs

@@ -2,9 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::{
-    DraftChunkRecord, ObjectRecord, PersistenceResult, PolicyRecord, current_time_ms,
-    draft_chunk_payload_from_record, draft_chunk_record_from_parts, map_db_error,
-    map_migrate_error, policy_payload_from_record, policy_record_from_parts,
+    DraftChunkRecord, ObjectRecord, PersistenceResult, PolicyRecord, current_time_ms, map_db_error,
+    map_migrate_error,
+};
+use crate::policy_store::{
+    draft_chunk_payload_from_record, draft_chunk_record_from_parts, policy_payload_from_record,
+    policy_record_from_parts,
 };
 use sqlx::postgres::PgPoolOptions;
 use sqlx::{PgPool, Row};

crates/openshell-server/src/persistence/sqlite.rs

@@ -2,9 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::{
-    DraftChunkRecord, ObjectRecord, PersistenceResult, PolicyRecord, current_time_ms,
-    draft_chunk_payload_from_record, draft_chunk_record_from_parts, map_db_error,
-    map_migrate_error, policy_payload_from_record, policy_record_from_parts,
+    DraftChunkRecord, ObjectRecord, PersistenceResult, PolicyRecord, current_time_ms, map_db_error,
+    map_migrate_error,
+};
+use crate::policy_store::{
+    draft_chunk_payload_from_record, draft_chunk_record_from_parts, policy_payload_from_record,
+    policy_record_from_parts,
 };
 use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
 use sqlx::{Row, SqlitePool};

crates/openshell-server/src/persistence/tests.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::{ObjectType, Store, generate_name};
+use crate::policy_store::PolicyStoreExt;
 use openshell_core::proto::{ObjectForTest, SandboxPolicy};
 use prost::Message;