weftgraph 0.1.0

//! A `GraphServices` a test can drive: the real service, a real
//! `PolicyEnforcer` in front of a stub PDP, the in-memory store, and a
//! one-hop engine over that same store so traversal runs rather than being
//! mocked away.
//!
//! Shared by the service cases and the REST cases, which differ only in
//! where they enter.

#![allow(clippy::expect_used, clippy::unwrap_used, dead_code)]

use super::conformance;

use std::sync::Arc;

use async_trait::async_trait;
use authz_resolver_sdk::api::AuthZResolverApi;
use authz_resolver_sdk::constraints::{Constraint, InPredicate, Predicate};
use authz_resolver_sdk::models::{
    EvaluationRequest, EvaluationResponse, EvaluationResponseContext,
};
use authz_resolver_sdk::pep::PolicyEnforcer;
use graph_storage::config::{GraphStorageConfig, ValidatedConfig};
use graph_storage::domain::service::GraphServices;
use graph_storage::infra::fake_store::FakeGraphStore;
use graph_storage_sdk::models::{Direction, EdgeRef, GraphRevision, NodeId};
use graph_storage_sdk::plugin_api::{
    EngineCursor, ExpandRequest, ExpandResponse, GraphEngineError, GraphEngineV1, GraphStoreV1,
    PathResponse, PatternRequest, PatternResponse, ShortestPathRequest, StoreCtx,
};
use toolkit_canonical_errors::CanonicalError;
use toolkit_security::{PlatformSecurityContext, SecurityContext, pep_properties};
use uuid::Uuid;

// ---------------------------------------------------------------------------
// Harness
// ---------------------------------------------------------------------------

/// A PDP that admits everything, scoped to the caller's own tenant.
///
/// Allow-all *at the action level* and tenant-bounded at the data level,
/// which is the deployment posture the gear is written for: the PDP decides
/// who may call, the compiled scope decides what they see. A denying variant
/// below covers the other branch.
pub struct AllowInOwnTenant;

#[async_trait]
impl AuthZResolverApi for AllowInOwnTenant {
    async fn evaluate(
        &self,
        _ctx: PlatformSecurityContext,
        request: EvaluationRequest,
    ) -> Result<EvaluationResponse, CanonicalError> {
        let tenant = request
            .subject
            .properties
            .get("tenant_id")
            .and_then(|value| value.as_str())
            .and_then(|raw| Uuid::parse_str(raw).ok())
            .unwrap_or_else(Uuid::nil);
        Ok(EvaluationResponse {
            decision: true,
            context: EvaluationResponseContext {
                constraints: vec![Constraint {
                    predicates: vec![Predicate::In(InPredicate::new(
                        pep_properties::OWNER_TENANT_ID,
                        [tenant],
                    ))],
                }],
                deny_reason: None,
            },
        })
    }
}

/// A PDP that grants `read` in the caller's own tenant and refuses every
/// other action -- a producer team's permissions, not an ontology
/// administrator's.
pub struct ReadOnly;

#[async_trait]
impl AuthZResolverApi for ReadOnly {
    async fn evaluate(
        &self,
        ctx: PlatformSecurityContext,
        request: EvaluationRequest,
    ) -> Result<EvaluationResponse, CanonicalError> {
        if request.action.name != "read" {
            return Ok(EvaluationResponse {
                decision: false,
                context: EvaluationResponseContext {
                    constraints: Vec::new(),
                    deny_reason: None,
                },
            });
        }
        AllowInOwnTenant.evaluate(ctx, request).await
    }
}

/// A PDP that denies. The gear must answer `permission_denied` from the
/// service rather than reaching the store at all.
struct DenyEverything;

#[async_trait]
impl AuthZResolverApi for DenyEverything {
    async fn evaluate(
        &self,
        _ctx: PlatformSecurityContext,
        _request: EvaluationRequest,
    ) -> Result<EvaluationResponse, CanonicalError> {
        Ok(EvaluationResponse {
            decision: false,
            context: EvaluationResponseContext {
                constraints: Vec::new(),
                deny_reason: None,
            },
        })
    }
}

/// A PDP that must never be asked.
///
/// `no_new_work_starts_after_the_budget_is_spent` claims the refusal happens
/// before the policy call. Asserting the returned variant cannot show that: a
/// deadline check moved below `authorize` returns the same `Deadline` and the
/// wasted round trip is invisible. This one records every evaluation, so the
/// claim in the test's name becomes something the test can fail on.
#[derive(Default)]
pub struct CountingPdp {
    pub calls: std::sync::atomic::AtomicUsize,
}

impl CountingPdp {
    pub fn calls(&self) -> usize {
        self.calls.load(std::sync::atomic::Ordering::SeqCst)
    }
}

#[async_trait]
impl AuthZResolverApi for CountingPdp {
    async fn evaluate(
        &self,
        ctx: PlatformSecurityContext,
        request: EvaluationRequest,
    ) -> Result<EvaluationResponse, CanonicalError> {
        self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        // Answers exactly as `AllowInOwnTenant` does, so swapping it in
        // changes what is observed and not what the service decides. A decision
        // with no constraints would compile to a different scope and make this
        // case fail for a reason that has nothing to do with deadlines.
        AllowInOwnTenant.evaluate(ctx, request).await
    }
}

/// A one-hop engine over the same store the service uses.
///
/// Enough to drive the traversal service for real -- seeds resolve, hops
/// expand, budgets and filters apply -- without a `PostgreSQL` server. It
/// declares neither shortest path nor pattern matching, so the service's
/// unsupported branches are reached as well.
struct HopOverStore {
    store: Arc<FakeGraphStore>,
}

#[async_trait]
impl GraphEngineV1 for HopOverStore {
    fn capabilities(&self) -> graph_storage_sdk::models::EngineCapabilities {
        graph_storage_sdk::models::EngineCapabilities::default()
    }

    async fn cursor(&self, ctx: &StoreCtx<'_>) -> Result<EngineCursor, GraphEngineError> {
        let revision = self.store.revision(ctx).await.unwrap_or(GraphRevision {
            source_epoch: 1,
            revision: 0,
        });
        Ok(EngineCursor { revision })
    }

    async fn expand(
        &self,
        ctx: &StoreCtx<'_>,
        req: ExpandRequest,
    ) -> Result<ExpandResponse, GraphEngineError> {
        // Hydration answers with the node but not its adjacency, so the
        // frontier is turned back into keys and each one is read: this stub
        // stands in for an engine, and an engine is exactly the thing that
        // knows the edges.
        let frontier = self
            .store
            .hydrate_nodes(ctx, &req.frontier)
            .await
            .map_err(|error| GraphEngineError::Unavailable {
                reason: error.to_string(),
            })?;
        let mut views = Vec::new();
        for node in frontier {
            if let Ok(view) = self.store.get_node(ctx, &node.node_key, 1000).await {
                views.push(view);
            }
        }
        let mut edges: Vec<EdgeRef> = Vec::new();
        let mut neighbours: Vec<String> = Vec::new();
        for view in views {
            for entry in view.adjacency {
                let outgoing = entry.side == graph_storage_sdk::models::AdjacencySide::Outgoing;
                let wanted = match req.direction {
                    Direction::Outgoing => outgoing,
                    Direction::Incoming => !outgoing,
                    Direction::Either => true,
                };
                let admitted = req
                    .edge_types
                    .as_ref()
                    .is_none_or(|set| set.contains(&entry.edge_type_id));
                if !wanted || !admitted {
                    continue;
                }
                let (src, dst) = if outgoing {
                    (view.node_key.clone(), entry.neighbor_key.clone())
                } else {
                    (entry.neighbor_key.clone(), view.node_key.clone())
                };
                edges.push(EdgeRef {
                    edge_key: entry.edge_key,
                    edge_type_id: entry.edge_type_id,
                    src,
                    dst,
                });
                neighbours.push(entry.neighbor_key);
            }
        }
        neighbours.sort();
        neighbours.dedup();
        let resolved = self
            .store
            .resolve_node_ids(ctx, &neighbours)
            .await
            .map_err(|error| GraphEngineError::Unavailable {
                reason: error.to_string(),
            })?;
        let mut reached: Vec<NodeId> = Vec::new();
        let mut degrees: Vec<u32> = Vec::new();
        for (key, id) in resolved {
            reached.push(id);
            if req.with_degrees {
                // The stub reads the neighbour back to count its adjacency,
                // which is this engine's version of the second scoped read
                // the real ones do.
                let degree = self
                    .store
                    .get_node(ctx, &key, 1000)
                    .await
                    .map_or(0, |view| {
                        u32::try_from(view.adjacency.len()).unwrap_or(u32::MAX)
                    });
                degrees.push(degree);
            }
        }
        Ok(ExpandResponse {
            reached,
            degrees,
            edges,
            truncated: None,
            served_by: graph_storage_sdk::plugin_api::HopBackend::TwoQuery,
        })
    }

    async fn shortest_path(
        &self,
        _ctx: &StoreCtx<'_>,
        _req: ShortestPathRequest,
    ) -> Result<PathResponse, GraphEngineError> {
        Err(GraphEngineError::Unsupported {
            what: "shortest_path",
        })
    }

    async fn match_pattern(
        &self,
        _ctx: &StoreCtx<'_>,
        _req: PatternRequest,
    ) -> Result<PatternResponse, GraphEngineError> {
        Err(GraphEngineError::Unsupported {
            what: "match_pattern",
        })
    }
}

pub struct Harness {
    pub services: Arc<GraphServices>,
    pub tenant: Uuid,
}

impl Harness {
    pub fn with(authz: Arc<dyn AuthZResolverApi>) -> Self {
        Self::configured(authz, GraphStorageConfig::default())
    }

    /// The same, with limits a case can make small enough to reach.
    ///
    /// A byte ceiling of tens of megabytes is not something a test should
    /// build its way up to: the case would spend its time allocating rather
    /// than asserting, and would be measuring the machine.
    /// The same, over a store the case supplies -- one that declines
    /// snapshots, for instance.
    pub fn over(store: Arc<FakeGraphStore>, authz: Arc<dyn AuthZResolverApi>) -> Self {
        let engine = Arc::new(HopOverStore {
            store: Arc::clone(&store),
        });
        Self {
            services: Arc::new(GraphServices::new(
                GraphStorageConfig::default()
                    .validated()
                    .expect("the default configuration is valid"),
                store,
                engine,
                PolicyEnforcer::new(authz),
                conformance::coordinator(),
            )),
            tenant: Uuid::now_v7(),
        }
    }

    pub fn configured(authz: Arc<dyn AuthZResolverApi>, config: GraphStorageConfig) -> Self {
        Self::configured_over(Arc::new(FakeGraphStore::new()), authz, config)
    }

    /// Both at once: limits the case sets, over a store the case keeps a
    /// handle to -- for a case that has to look at what the store was asked
    /// for, not only at what the service answered.
    pub fn configured_over(
        store: Arc<FakeGraphStore>,
        authz: Arc<dyn AuthZResolverApi>,
        config: GraphStorageConfig,
    ) -> Self {
        let engine = Arc::new(HopOverStore {
            store: Arc::clone(&store),
        });
        Self {
            services: Arc::new(GraphServices::new(
                // Deliberately unchecked: cases reach a limit by setting it
                // below the hard range, which is the only cheap way to prove
                // what happens there, and a zero deadline has no legal value.
                ValidatedConfig::unchecked(config),
                store,
                engine,
                PolicyEnforcer::new(authz),
                conformance::coordinator(),
            )),
            tenant: Uuid::now_v7(),
        }
    }

    /// The same, over any store, with the fake the test engine walks.
    pub fn configured_over_store(
        store: Arc<dyn GraphStoreV1>,
        fake: Arc<FakeGraphStore>,
        authz: Arc<dyn AuthZResolverApi>,
        config: GraphStorageConfig,
    ) -> Self {
        let engine = Arc::new(HopOverStore { store: fake });
        Self {
            services: Arc::new(GraphServices::new(
                ValidatedConfig::unchecked(config),
                store,
                engine,
                PolicyEnforcer::new(authz),
                conformance::coordinator(),
            )),
            tenant: Uuid::now_v7(),
        }
    }

    /// Over an embedding coordinator the case supplies -- one whose space is
    /// blocked, for instance, which no store write can produce.
    pub fn with_coordinator(
        authz: Arc<dyn AuthZResolverApi>,
        coordinator: graph_storage::domain::embedding::EmbeddingCoordinator,
    ) -> Self {
        let store = Arc::new(FakeGraphStore::new());
        let engine = Arc::new(HopOverStore {
            store: Arc::clone(&store),
        });
        Self {
            services: Arc::new(GraphServices::new(
                GraphStorageConfig::default()
                    .validated()
                    .expect("the default configuration is valid"),
                store,
                engine,
                PolicyEnforcer::new(authz),
                coordinator,
            )),
            tenant: Uuid::now_v7(),
        }
    }

    pub fn allowed() -> Self {
        Self::with(Arc::new(AllowInOwnTenant))
    }

    pub fn denied() -> Self {
        Self::with(Arc::new(DenyEverything))
    }

    pub fn ctx(&self) -> SecurityContext {
        SecurityContext::builder()
            .subject_id(Uuid::now_v7())
            .subject_tenant_id(self.tenant)
            .build()
            .expect("a valid security context")
    }

    /// The base ontology plus the suite's producer types, through the service.
    pub async fn seed_ontology(&self, ctx: &SecurityContext) {
        self.services
            .register_types(ctx, conformance::ontology_batch())
            .await
            .expect("the ontology registers");
    }
}

/// A store that does **not** implement `node_types`, so the trait's default
/// answers -- the thing an external or older store would rely on, and the
/// thing `FakeGraphStore::without_node_types` only imitates. Every required
/// method delegates to the fake it wraps; the optional ones keep their
/// defaults. Generated from the trait's required signatures; a new required
/// method is a compile error here, which is the point.
pub mod without_node_types {
    use std::sync::Arc;

    use async_trait::async_trait;
    use graph_storage::infra::fake_store::FakeGraphStore;
    use graph_storage_sdk::models::*;
    use graph_storage_sdk::plugin_api::*;

    pub struct StoreWithoutNodeTypes(pub Arc<FakeGraphStore>);

    #[async_trait]
    impl GraphStoreV1 for StoreWithoutNodeTypes {
        fn capabilities(&self) -> StoreCapabilities {
            self.0.capabilities()
        }

        async fn register_types_with(
            &self,
            ctx: &StoreCtx<'_>,
            batch: Vec<TypeRegistration>,
            options: TypeRegistrationOptions,
        ) -> Result<Vec<RegisteredType>, GraphStoreError> {
            self.0.register_types_with(ctx, batch, options).await
        }

        async fn get_type(
            &self,
            ctx: &StoreCtx<'_>,
            id: &GtsTypeId,
        ) -> Result<TypeRecord, GraphStoreError> {
            self.0.get_type(ctx, id).await
        }

        async fn list_types(
            &self,
            ctx: &StoreCtx<'_>,
            query: TypeQuery,
        ) -> Result<Page<TypeRecord>, GraphStoreError> {
            self.0.list_types(ctx, query).await
        }

        async fn probe_readiness(&self) -> Vec<ComponentReadiness> {
            self.0.probe_readiness().await
        }

        async fn list_source_namespaces(
            &self,
            ctx: &StoreCtx<'_>,
        ) -> Result<Vec<SourceNamespaceOwner>, GraphStoreError> {
            self.0.list_source_namespaces(ctx).await
        }

        async fn transfer_source_namespace(
            &self,
            ctx: &StoreCtx<'_>,
            namespace: &str,
            owner_principal: &str,
        ) -> Result<SourceNamespaceOwner, GraphStoreError> {
            self.0
                .transfer_source_namespace(ctx, namespace, owner_principal)
                .await
        }

        async fn resolve_type_set(
            &self,
            ctx: &StoreCtx<'_>,
            patterns: &[String],
        ) -> Result<TypeIdSet, GraphStoreError> {
            self.0.resolve_type_set(ctx, patterns).await
        }

        async fn ingest(
            &self,
            ctx: &StoreCtx<'_>,
            req: IngestRequest,
            embedding: EmbeddingPlan,
        ) -> Result<IngestOutcome, GraphStoreError> {
            self.0.ingest(ctx, req, embedding).await
        }

        async fn soft_delete(
            &self,
            ctx: &StoreCtx<'_>,
            req: DeleteRequest,
        ) -> Result<DeleteOutcome, GraphStoreError> {
            self.0.soft_delete(ctx, req).await
        }

        async fn begin_read(&self, ctx: &StoreCtx<'_>) -> Result<ReadSnapshot, GraphStoreError> {
            self.0.begin_read(ctx).await
        }

        async fn end_read(&self, snapshot: ReadSnapshot) -> Result<(), GraphStoreError> {
            self.0.end_read(snapshot).await
        }

        async fn revision(&self, ctx: &StoreCtx<'_>) -> Result<GraphRevision, GraphStoreError> {
            self.0.revision(ctx).await
        }

        async fn get_node(
            &self,
            ctx: &StoreCtx<'_>,
            key: &NodeKey,
            adjacency_limit: u32,
        ) -> Result<NodeView, GraphStoreError> {
            self.0.get_node(ctx, key, adjacency_limit).await
        }

        async fn hydrate_nodes(
            &self,
            ctx: &StoreCtx<'_>,
            ids: &[NodeId],
        ) -> Result<Vec<NodeView>, GraphStoreError> {
            self.0.hydrate_nodes(ctx, ids).await
        }

        async fn get_edge(
            &self,
            ctx: &StoreCtx<'_>,
            key: &EdgeKey,
        ) -> Result<EdgeView, GraphStoreError> {
            self.0.get_edge(ctx, key).await
        }

        async fn search(
            &self,
            ctx: &StoreCtx<'_>,
            req: SearchRequest,
            vector: Option<VectorArm>,
        ) -> Result<SearchResponse, GraphStoreError> {
            self.0.search(ctx, req, vector).await
        }

        async fn project_table(
            &self,
            ctx: &StoreCtx<'_>,
            req: ProjectionRequest,
        ) -> Result<toolkit_odata::Page<NodeRow>, GraphStoreError> {
            self.0.project_table(ctx, req).await
        }

        async fn load_topology(
            &self,
            ctx: &StoreCtx<'_>,
            req: TopologyRequest,
        ) -> Result<TopologyPage, GraphStoreError> {
            self.0.load_topology(ctx, req).await
        }

        async fn resolve_node_ids(
            &self,
            ctx: &StoreCtx<'_>,
            keys: &[NodeKey],
        ) -> Result<Vec<(NodeKey, NodeId)>, GraphStoreError> {
            self.0.resolve_node_ids(ctx, keys).await
        }

        async fn embedding_state(
            &self,
            ctx: &StoreCtx<'_>,
            keys: &[NodeKey],
        ) -> Result<Vec<Option<EmbeddingState>>, GraphStoreError> {
            self.0.embedding_state(ctx, keys).await
        }
    }
}