recombine oxql_query and oxql_query_project with scope concept

david-crespo · david-crespo · commit d33d83fff34e · 2025-04-02T15:14:14.000-05:00
diff --git a/nexus/src/app/metrics.rs b/nexus/src/app/metrics.rs
@@ -14,7 +14,7 @@ use nexus_db_queries::{
 use nexus_external_api::TimeseriesSchemaPaginationParams;
 use nexus_types::external_api::params::SystemMetricName;
 use omicron_common::api::external::{Error, InternalContext};
-use oximeter_db::{Measurement, TimeseriesSchema};
+use oximeter_db::{Measurement, QueryAuthzScope, TimeseriesSchema};
 use std::num::NonZeroU32;
 
 impl super::Nexus {
@@ -138,7 +138,7 @@ impl super::Nexus {
         // resources they have access to.
         opctx.authorize(authz::Action::Read, &authz::FLEET).await?;
         self.timeseries_client
-            .oxql_query(query)
+            .oxql_query(query, QueryAuthzScope::Fleet)
             .await
             // TODO-observability: The query method returns information
             // about the duration of the OxQL query and the database
@@ -162,7 +162,13 @@ impl super::Nexus {
         let (authz_silo, authz_project) =
             project_lookup.lookup_for(authz::Action::Read).await?;
         self.timeseries_client
-            .oxql_query_project(query, authz_silo.id(), authz_project.id())
+            .oxql_query(
+                query,
+                QueryAuthzScope::Project {
+                    silo_id: authz_silo.id(),
+                    project_id: authz_project.id(),
+                },
+            )
             .await
             .map(|result| result.tables)
             .map_err(map_timeseries_err)
diff --git a/oximeter/db/src/client/oxql.rs b/oximeter/db/src/client/oxql.rs
@@ -20,7 +20,7 @@ use crate::oxql::ast::table_ops::filter;
 use crate::oxql::ast::table_ops::filter::Filter;
 use crate::oxql::ast::table_ops::limit::Limit;
 use crate::oxql::ast::table_ops::limit::LimitKind;
-use crate::oxql::query::uuid_eq_filter;
+use crate::oxql::query::QueryAuthzScope;
 use crate::query::field_table_name;
 use oximeter::Measurement;
 use oximeter::TimeseriesSchema;
@@ -148,72 +148,11 @@ impl Client {
     pub async fn oxql_query(
         &self,
         query: impl AsRef<str>,
+        scope: QueryAuthzScope,
     ) -> Result<OxqlResult, Error> {
-        // TODO-security: Need a way to implement authz checks for things like
-        // viewing resources in another project or silo.
-        //
-        // I think one way to do that is look at the predicates and make sure
-        // they refer to things the user has access to. Another is to add some
-        // implicit predicates here, indicating the subset of fields that the
-        // query should be able to access.
-        //
-        // This probably means we'll need to parse the query in Nexus, so that
-        // we can attach the other filters ourselves.
-        //
-        // See https://github.com/oxidecomputer/omicron/issues/5298.
         let query = query.as_ref();
         let parsed_query = oxql::Query::new(query)?;
-        let plan = self.build_query_plan(&parsed_query).await?;
-        if plan.requires_full_table_scan() {
-            return Err(Error::Oxql(anyhow::anyhow!(
-                "This query requires at least one full table scan. \
-                Please rewrite the query to filter either the fields \
-                or timestamps, in order to reduce the amount of data \
-                fetched from the database."
-            )));
-        }
-        let query_id = Uuid::new_v4();
-        let query_log =
-            self.log.new(slog::o!("query_id" => query_id.to_string()));
-        debug!(
-            query_log,
-            "parsed OxQL query";
-            "query" => query,
-            "parsed_query" => ?parsed_query,
-        );
-        let id = usdt::UniqueId::new();
-        probes::oxql__query__start!(|| (&id, &query_id, query));
-        let mut total_rows_fetched = 0;
-        let result = self
-            .run_oxql_query(
-                &query_log,
-                &mut self.claim_connection().await?,
-                query_id,
-                parsed_query,
-                &mut total_rows_fetched,
-                None,
-                None,
-            )
-            .await;
-        probes::oxql__query__done!(|| (&id, &query_id));
-        result
-    }
-
-    /// Run a OxQL query.
-    pub async fn oxql_query_project(
-        &self,
-        query: impl AsRef<str>,
-        silo_id: Uuid,
-        project_id: Uuid,
-    ) -> Result<OxqlResult, Error> {
-        let query = query.as_ref();
-
-        let parsed_query = oxql::Query::new(query)?;
-        // this has to be used for everything
-        let filtered_query = parsed_query.insert_filters(vec![
-            uuid_eq_filter("silo_id", silo_id),
-            uuid_eq_filter("project_id", project_id),
-        ]);
+        let filtered_query = parsed_query.insert_authz_filters(scope);
 
         let plan = self.build_query_plan(&filtered_query).await?;
         if plan.requires_full_table_scan() {
@@ -231,8 +170,8 @@ impl Client {
             query_log,
             "parsed OxQL query";
             "query" => query,
-            "filtered_query" => ?filtered_query,
             "parsed_query" => ?parsed_query,
+            "filtered_query" => ?filtered_query,
         );
         let id = usdt::UniqueId::new();
         probes::oxql__query__start!(|| (&id, &query_id, query));
@@ -1241,7 +1180,9 @@ fn update_total_rows_and_check(
 #[cfg(test)]
 mod tests {
     use super::ConsistentKeyGroup;
-    use crate::client::oxql::chunk_consistent_key_groups_impl;
+    use crate::client::oxql::{
+        QueryAuthzScope, chunk_consistent_key_groups_impl,
+    };
     use crate::oxql::ast::grammar::query_parser;
     use crate::{Client, DATABASE_TIMESTAMP_FORMAT, DbWrite};
     use crate::{Metric, Target};
@@ -1402,7 +1343,7 @@ mod tests {
             "get some_target:some_metric | filter timestamp > @2020-01-01";
         let result = ctx
             .client
-            .oxql_query(query)
+            .oxql_query(query, QueryAuthzScope::Fleet)
             .await
             .expect("failed to run OxQL query");
         assert_eq!(result.tables.len(), 1, "Should be exactly 1 table");
@@ -1461,7 +1402,7 @@ mod tests {
         );
         let result = ctx
             .client
-            .oxql_query(&query)
+            .oxql_query(&query, QueryAuthzScope::Fleet)
             .await
             .expect("failed to run OxQL query");
         assert_eq!(result.tables.len(), 1, "Should be exactly 1 table");
@@ -1517,7 +1458,7 @@ mod tests {
         );
         let result = ctx
             .client
-            .oxql_query(&query)
+            .oxql_query(&query, QueryAuthzScope::Fleet)
             .await
             .expect("failed to run OxQL query");
         assert_eq!(result.tables.len(), 1, "Should be exactly 1 table");
@@ -1749,7 +1690,7 @@ mod tests {
         );
         let result = ctx
             .client
-            .oxql_query(&query)
+            .oxql_query(&query, QueryAuthzScope::Fleet)
             .await
             .expect("failed to run OxQL query");
         assert_eq!(result.tables.len(), 1, "Should be exactly 1 table");
diff --git a/oximeter/db/src/lib.rs b/oximeter/db/src/lib.rs
@@ -54,6 +54,7 @@ pub use client::TestDbWrite;
 pub use client::oxql::OxqlResult;
 pub use client::query_summary::QuerySummary;
 pub use model::OXIMETER_VERSION;
+pub use oxql::query::QueryAuthzScope;
 
 #[derive(Debug, Error)]
 pub enum Error {
diff --git a/oximeter/db/src/oxql/query/mod.rs b/oximeter/db/src/oxql/query/mod.rs
@@ -37,6 +37,12 @@ pub struct Query {
     pub(super) end_time: DateTime<Utc>,
 }
 
+pub enum QueryAuthzScope {
+    Fleet,
+    Silo { silo_id: Uuid },
+    Project { silo_id: Uuid, project_id: Uuid },
+}
+
 impl Query {
     /// Construct a query written in OxQL.
     pub fn new(query: impl AsRef<str>) -> Result<Self, Error> {
@@ -367,19 +373,27 @@ impl Query {
         &self.parsed
     }
 
-    /// Insert filters after the `get`, or in the case of subqueries, recurse
-    /// down the tree and insert them after each get.
-    pub(crate) fn insert_filters(&self, filters: Vec<Filter>) -> Self {
-        Self {
-            parsed: self.parsed.insert_filters(filters),
-            end_time: self.end_time,
-        }
+    /// Insert silo and project filters after the `get`, or in the case of
+    /// subqueries, recurse down the tree and insert them after each get.
+    pub(crate) fn insert_authz_filters(&self, scope: QueryAuthzScope) -> Self {
+        let filtered_query = match scope {
+            QueryAuthzScope::Fleet => self.parsed.clone(),
+            QueryAuthzScope::Silo { silo_id } => self
+                .parsed
+                .insert_filters(vec![uuid_eq_filter("silo_id", silo_id)]),
+            QueryAuthzScope::Project { silo_id, project_id } => {
+                self.parsed.insert_filters(vec![
+                    uuid_eq_filter("silo_id", silo_id),
+                    uuid_eq_filter("project_id", project_id),
+                ])
+            }
+        };
+        Self { parsed: filtered_query, end_time: self.end_time }
     }
 }
 
-// TODO: move this to a more appropriate spot
 /// Just a helper for creating a UUID filter node concisely
-pub fn uuid_eq_filter(key: impl AsRef<str>, id: Uuid) -> Filter {
+fn uuid_eq_filter(key: impl AsRef<str>, id: Uuid) -> Filter {
     let simple_filter = SimpleFilter {
         ident: Ident(key.as_ref().to_string()),
         cmp: Comparison::Eq,
@@ -432,6 +446,7 @@ mod tests {
     use super::Filter;
     use super::Ident;
     use super::Query;
+    use crate::QueryAuthzScope;
     use crate::oxql::ast::SplitQuery;
     use crate::oxql::ast::cmp::Comparison;
     use crate::oxql::ast::literal::Literal;
@@ -1041,10 +1056,8 @@ mod tests {
         let query = Query::new("get a:b | filter timestamp > @now()").unwrap();
         let silo_id = Uuid::new_v4();
         let project_id = Uuid::new_v4();
-        let new_query = query.insert_filters(vec![
-            uuid_eq_filter("silo_id", silo_id),
-            uuid_eq_filter("project_id", project_id),
-        ]);
+        let scope = QueryAuthzScope::Project { silo_id, project_id };
+        let new_query = query.insert_authz_filters(scope);
 
         assert_eq!(query.parsed.table_ops().len(), 2);
         assert_eq!(new_query.parsed.table_ops().len(), 4);
@@ -1079,7 +1092,8 @@ mod tests {
         let expected_project_op =
             TableOp::Basic(BasicTableOp::Filter(project_filter.clone()));
 
-        let new_query = query.insert_filters(vec![silo_filter, project_filter]);
+        let scope = QueryAuthzScope::Project { silo_id, project_id };
+        let new_query = query.insert_authz_filters(scope);
 
         // Check top-level structure (should remain one grouped op and one filter)
         let orig_ops = query.parsed.table_ops().collect::<Vec<_>>();
@@ -1131,7 +1145,8 @@ mod tests {
         let expected_project_op =
             TableOp::Basic(BasicTableOp::Filter(project_filter.clone()));
 
-        let new_query = query.insert_filters(vec![silo_filter, project_filter]);
+        let scope = QueryAuthzScope::Project { silo_id, project_id };
+        let new_query = query.insert_authz_filters(scope);
 
         // Check top-level structure (should remain a single grouped op)
         assert_eq!(query.parsed.table_ops().len(), 1);
diff --git a/oximeter/db/src/shells/oxql.rs b/oximeter/db/src/shells/oxql.rs
@@ -7,7 +7,7 @@
 // Copyright 2024 Oxide Computer
 
 use super::{list_timeseries, prepare_columns};
-use crate::{Client, OxqlResult, make_client};
+use crate::{Client, OxqlResult, make_client, oxql::query::QueryAuthzScope};
 use clap::Args;
 use crossterm::style::Stylize;
 use oxql_types::Table;
@@ -126,7 +126,10 @@ pub async fn shell(
                             }
                         } else {
                             match client
-                                .oxql_query(cmd.trim().trim_end_matches(';'))
+                                .oxql_query(
+                                    cmd.trim().trim_end_matches(';'),
+                                    QueryAuthzScope::Fleet,
+                                )
                                 .await
                             {
                                 Ok(result) => {
diff --git a/oximeter/db/tests/integration_test.rs b/oximeter/db/tests/integration_test.rs
@@ -7,6 +7,7 @@ use clickward::{BasePorts, Deployment, DeploymentConfig, KeeperId};
 use dropshot::test_util::log_prefix_for_test;
 use omicron_test_utils::dev::poll;
 use omicron_test_utils::dev::test_setup_log;
+use oximeter_db::QueryAuthzScope;
 use oximeter_db::{Client, DbWrite, OxqlResult, Sample, TestDbWrite};
 use oximeter_test_utils::wait_for_keepers;
 use slog::{Logger, info};
@@ -213,6 +214,7 @@ async fn test_cluster() -> anyhow::Result<()> {
     let oxql_res1 = client1
         .oxql_query(
             "get virtual_machine:cpu_busy | filter timestamp > @2000-01-01",
+            QueryAuthzScope::Fleet,
         )
         .await
         .expect("failed to get all samples");
@@ -427,7 +429,9 @@ async fn wait_for_num_points(
     poll::wait_for_condition(
         || async {
             let oxql_res = client
-                .oxql_query("get virtual_machine:cpu_busy | filter timestamp > @2000-01-01")
+                .oxql_query(
+                    "get virtual_machine:cpu_busy | filter timestamp > @2000-01-01",
+                    QueryAuthzScope::Fleet)
                 .await
                 .map_err(|_| {
                     poll::CondCheckError::<oximeter_db::Error>::NotYet
