Merge pull request #32 from NodeDB-Lab/bug/pgwire-tenant-scoping-29-30

fix(pgwire): scope catalog resolution and trust auth to the connecting tenant

Author: farhan-syah

nodedb/src/control/event_trigger.rs

@@ -124,7 +124,7 @@ async fn execute_then_action(
         .replace("$collection", &event.collection)
         .replace("$operation", event.operation.as_str());
 
-    let query_ctx = QueryContext::for_state(shared, 1);
+    let query_ctx = QueryContext::for_state(shared);
 
     match query_ctx.plan_sql(&sql, event.tenant_id).await {
         Ok(tasks) => {

nodedb/src/control/planner/context.rs

@@ -50,29 +50,33 @@ pub struct QueryContext {
 }
 
 /// Inputs needed to construct an `OriginCatalog` per plan call.
+///
+/// Tenant is intentionally **not** stored here: every plan call passes the
+/// effective tenant to `build_adapter`, so a single `QueryContext` shared
+/// across a pgwire handler can serve queries from connections belonging to
+/// different tenants without cross-tenant catalog resolution.
 #[derive(Clone)]
 struct CatalogInputs {
     credentials: Arc<CredentialStore>,
     shared: Option<std::sync::Weak<crate::control::state::SharedState>>,
-    tenant_id: u32,
     retention_policy_registry:
         Option<Arc<crate::engine::timeseries::retention_policy::RetentionPolicyRegistry>>,
 }
 
 impl CatalogInputs {
-    fn build_adapter(&self) -> super::catalog_adapter::OriginCatalog {
+    fn build_adapter(&self, tenant_id: u32) -> super::catalog_adapter::OriginCatalog {
         if let Some(weak) = &self.shared
             && let Some(shared) = weak.upgrade()
         {
             super::catalog_adapter::OriginCatalog::new_with_lease(
                 &shared,
-                self.tenant_id,
+                tenant_id,
                 self.retention_policy_registry.clone(),
             )
         } else {
             super::catalog_adapter::OriginCatalog::new(
                 Arc::clone(&self.credentials),
-                self.tenant_id,
+                tenant_id,
                 self.retention_policy_registry.clone(),
             )
         }
@@ -97,10 +101,9 @@ impl QueryContext {
     /// path would return instantly anyway, but going through the
     /// sub-planner without a direct `Arc<SharedState>` reference
     /// would require threading one through every call site.
-    pub fn for_state(state: &crate::control::state::SharedState, tenant_id: u32) -> Self {
+    pub fn for_state(state: &crate::control::state::SharedState) -> Self {
         Self::with_catalog(
             Arc::clone(&state.credentials),
-            tenant_id,
             Some(Arc::clone(&state.retention_policy_registry)),
         )
     }
@@ -110,16 +113,12 @@ impl QueryContext {
     /// query's plan acquires descriptor leases before execution.
     /// Callers must hold an `Arc<SharedState>` — the adapter
     /// downgrades to `Weak` internally.
-    pub fn for_state_with_lease(
-        state: &Arc<crate::control::state::SharedState>,
-        tenant_id: u32,
-    ) -> Self {
+    pub fn for_state_with_lease(state: &Arc<crate::control::state::SharedState>) -> Self {
         let retention = Some(Arc::clone(&state.retention_policy_registry));
         Self {
             catalog_inputs: Some(CatalogInputs {
                 credentials: Arc::clone(&state.credentials),
                 shared: Some(Arc::downgrade(state)),
-                tenant_id,
                 retention_policy_registry: retention.clone(),
             }),
             retention_registry: retention,
@@ -136,15 +135,13 @@ impl QueryContext {
     /// that construct a context without an `Arc<SharedState>`.
     pub fn with_catalog(
         credentials: Arc<CredentialStore>,
-        tenant_id: u32,
         retention_policy_registry: Option<
             Arc<crate::engine::timeseries::retention_policy::RetentionPolicyRegistry>,
         >,
     ) -> Self {
         let catalog_inputs = Some(CatalogInputs {
             credentials,
             shared: None,
-            tenant_id,
             retention_policy_registry: retention_policy_registry.clone(),
         });
 
@@ -193,7 +190,7 @@ impl QueryContext {
         // `recorded_versions` field is per-plan state, and
         // two concurrent plans through a shared QueryContext
         // would otherwise interleave their recorded sets.
-        let catalog = inputs.build_adapter();
+        let catalog = inputs.build_adapter(tenant_id.as_u32());
         let plans = nodedb_sql::plan_sql(sql, &catalog).map_err(|e| match e {
             nodedb_sql::SqlError::RetryableSchemaChanged { descriptor } => {
                 crate::Error::RetryableSchemaChanged { descriptor }
@@ -292,7 +289,7 @@ impl QueryContext {
         // through a different cache key), but constructing the
         // adapter fresh keeps the adapter's state per-plan and
         // allows future extension.
-        let catalog = inputs.build_adapter();
+        let catalog = inputs.build_adapter(tenant_id.as_u32());
         let plans = nodedb_sql::plan_sql_with_params(sql, params, &catalog).map_err(|e| {
             crate::Error::PlanError {
                 detail: format!("{e}"),

nodedb/src/control/planner/procedural/executor/core/dispatch.rs

@@ -44,10 +44,7 @@ impl<'a> StatementExecutor<'a> {
     pub(super) async fn execute_dml(&self, sql: &str, bindings: &RowBindings) -> crate::Result<()> {
         let bound_sql = fold_literal_string_concat(&bindings.substitute(sql));
 
-        let ctx = crate::control::planner::context::QueryContext::for_state(
-            self.state,
-            self.tenant_id.as_u32(),
-        );
+        let ctx = crate::control::planner::context::QueryContext::for_state(self.state);
         let tasks = ctx.plan_sql(&bound_sql, self.tenant_id).await?;
 
         if let Some(ref tx_ctx) = self.tx_ctx {

nodedb/src/control/scatter_gather.rs

@@ -343,7 +343,6 @@ pub async fn coordinate_cross_shard_hop(
                 // (same pattern as QueryContext::for_state but without &SharedState).
                 let plan_ctx = crate::control::planner::context::QueryContext::with_catalog(
                     std::sync::Arc::clone(&credentials_clone),
-                    tenant_id_u32,
                     Some(std::sync::Arc::clone(&retention_clone)),
                 );
 

nodedb/src/control/server/http/server.rs

@@ -159,7 +159,7 @@ pub async fn run_with_listener(
     let mut shutdown_rx = bus.handle().flat_watch().raw_receiver();
 
     let query_ctx = Arc::new(crate::control::planner::context::QueryContext::for_state(
-        &shared, 1,
+        &shared,
     ));
     let state = AppState {
         shared,
@@ -198,7 +198,7 @@ pub async fn run(
     let mut shutdown_rx = bus.handle().flat_watch().raw_receiver();
 
     let query_ctx = Arc::new(crate::control::planner::context::QueryContext::for_state(
-        &shared, 1,
+        &shared,
     ));
     let state = AppState {
         shared,

nodedb/src/control/server/native/session.rs

@@ -48,7 +48,7 @@ impl NativeSession {
         state: Arc<SharedState>,
         auth_mode: AuthMode,
     ) -> Self {
-        let query_ctx = QueryContext::for_state(&state, 1); // default tenant
+        let query_ctx = QueryContext::for_state(&state);
         Self {
             stream,
             peer_addr,

nodedb/src/control/server/pgwire/ddl/collection/check_constraint.rs

@@ -102,8 +102,7 @@ async fn enforce_subquery_check(
     // General fallback: wrap in subselect.
     let restructured = restructure_subquery_check(&substituted);
 
-    let query_ctx =
-        crate::control::planner::context::QueryContext::for_state(state, tenant_id.as_u32());
+    let query_ctx = crate::control::planner::context::QueryContext::for_state(state);
 
     let tasks = match query_ctx.plan_sql(&restructured.sql, tenant_id).await {
         Ok(t) => t,

nodedb/src/control/server/pgwire/ddl/collection/insert_parse.rs

@@ -304,8 +304,7 @@ pub(super) async fn plan_and_dispatch(
     tenant_id: nodedb_types::TenantId,
     sql: &str,
 ) -> PgWireResult<()> {
-    let query_ctx =
-        crate::control::planner::context::QueryContext::for_state(state, tenant_id.as_u32());
+    let query_ctx = crate::control::planner::context::QueryContext::for_state(state);
     let tasks = query_ctx
         .plan_sql(sql, tenant_id)
         .await

nodedb/src/control/server/pgwire/ddl/maintenance/analyze.rs

@@ -69,7 +69,7 @@ pub async fn handle_analyze(
 
     // Dispatch a scan to the Data Plane to collect all rows.
     let scan_sql = format!("SELECT * FROM {collection}");
-    let query_ctx = crate::control::planner::context::QueryContext::for_state(state, tenant_id);
+    let query_ctx = crate::control::planner::context::QueryContext::for_state(state);
     let rows = match query_ctx.plan_sql(&scan_sql, identity.tenant_id).await {
         Ok(tasks) => {
             let mut json_rows = Vec::new();

nodedb/src/control/server/pgwire/ddl/typeguard/validate.rs

@@ -57,8 +57,7 @@ pub async fn validate_typeguard(
 
     // Scan all documents.
     let scan_sql = format!("SELECT * FROM {coll_name}");
-    let query_ctx =
-        crate::control::planner::context::QueryContext::for_state(state, tenant_id.as_u32());
+    let query_ctx = crate::control::planner::context::QueryContext::for_state(state);
     let tasks = query_ctx
         .plan_sql(&scan_sql, tenant_id)
         .await