feat: Batch OPA column mask calls

deploy/helm/trino-operator/crds/crds.yaml

@@ -70,20 +70,25 @@ spec:
                         Authorization options for Trino.
                         Learn more in the [Trino authorization usage guide](https://docs.stackable.tech/home/nightly/trino/usage-guide/security#authorization).
                       nullable: true
+                      oneOf:
+                        - required:
+                            - opa
                       properties:
                         opa:
                           description: |-
                             Configure the OPA stacklet [discovery ConfigMap](https://docs.stackable.tech/home/nightly/concepts/service_discovery)
                             and the name of the Rego package containing your authorization rules.
                             Consult the [OPA authorization documentation](https://docs.stackable.tech/home/nightly/concepts/opa)
                             to learn how to deploy Rego authorization rules with OPA.
-                          nullable: true
                           properties:
                             configMapName:
                               description: |-
                                 The [discovery ConfigMap](https://docs.stackable.tech/home/nightly/concepts/service_discovery)
                                 for the OPA stacklet that should be used for authorization requests.
                               type: string
+                            enableColumnMasking:
+                              default: true
+                              type: boolean
                             package:
                               description: The name of the Rego package containing the Rego rules for the product.
                               nullable: true

rust/operator-binary/src/authorization/opa.rs

@@ -1,13 +1,11 @@
 use std::collections::BTreeMap;
 
 use stackable_operator::{
-    client::Client,
-    commons::opa::{OpaApiVersion, OpaConfig},
-    k8s_openapi::api::core::v1::ConfigMap,
+    client::Client, commons::opa::OpaApiVersion, k8s_openapi::api::core::v1::ConfigMap,
     kube::ResourceExt,
 };
 
-use crate::crd::v1alpha1::TrinoCluster;
+use crate::crd::v1alpha1;
 
 pub const OPA_TLS_VOLUME_NAME: &str = "opa-tls";
 
@@ -23,10 +21,10 @@ pub struct TrinoOpaConfig {
     /// `http://localhost:8081/v1/data/trino/rowFilters` - if not set,
     /// no row filtering will be applied
     pub(crate) row_filters_connection_string: Option<String>,
-    /// URI for fetching column masks, e.g.
-    /// `http://localhost:8081/v1/data/trino/columnMask` - if not set,
-    /// no masking will be applied
-    pub(crate) column_masking_connection_string: Option<String>,
+    /// URI for fetching columns masks in batches, e.g.
+    /// `http://localhost:8081/v1/data/trino/batchColumnMasks` - if not set,
+    /// column-masking-uri will be used for getting column masks in parallel
+    pub(crate) batched_column_masking_connection_string: Option<String>,
     /// Whether to allow permission management (GRANT, DENY, ...) and
     /// role management operations - OPA will not be queried for any
     /// such operations, they will be bulk allowed or denied depending
@@ -41,13 +39,15 @@ pub struct TrinoOpaConfig {
 impl TrinoOpaConfig {
     pub async fn from_opa_config(
         client: &Client,
-        trino: &TrinoCluster,
-        opa_config: &OpaConfig,
+        trino: &v1alpha1::TrinoCluster,
+        opa_config: &v1alpha1::TrinoAuthorizationOpaConfig,
     ) -> Result<Self, stackable_operator::commons::opa::Error> {
         let non_batched_connection_string = opa_config
+            .opa
             .full_document_url_from_config_map(client, trino, Some("allow"), OpaApiVersion::V1)
             .await?;
         let batched_connection_string = opa_config
+            .opa
             .full_document_url_from_config_map(
                 client,
                 trino,
@@ -57,6 +57,7 @@ impl TrinoOpaConfig {
             )
             .await?;
         let row_filters_connection_string = opa_config
+            .opa
             .full_document_url_from_config_map(
                 client,
                 trino,
@@ -65,19 +66,27 @@ impl TrinoOpaConfig {
                 OpaApiVersion::V1,
             )
             .await?;
-        let column_masking_connection_string = opa_config
-            .full_document_url_from_config_map(
-                client,
-                trino,
-                // Sticking to https://github.com/trinodb/trino/blob/455/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControlDataFilteringSystem.java#L47
-                Some("columnMask"),
-                OpaApiVersion::V1,
+
+        let batched_column_masking_connection_string = if trino.column_masking_enabled() {
+            Some(
+                opa_config
+                    .opa
+                    .full_document_url_from_config_map(
+                        client,
+                        trino,
+                        // Sticking to https://github.com/trinodb/trino/blob/455/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaAccessControlDataFilteringSystem.java#L48
+                        Some("batchColumnMasks"),
+                        OpaApiVersion::V1,
+                    )
+                    .await?,
             )
-            .await?;
+        } else {
+            None
+        };
 
         let tls_secret_class = client
             .get::<ConfigMap>(
-                &opa_config.config_map_name,
+                &opa_config.opa.config_map_name,
                 trino.namespace().as_deref().unwrap_or("default"),
             )
             .await
@@ -89,7 +98,7 @@ impl TrinoOpaConfig {
             non_batched_connection_string,
             batched_connection_string,
             row_filters_connection_string: Some(row_filters_connection_string),
-            column_masking_connection_string: Some(column_masking_connection_string),
+            batched_column_masking_connection_string,
             allow_permission_management_operations: true,
             tls_secret_class,
         })
@@ -113,10 +122,12 @@ impl TrinoOpaConfig {
                 Some(row_filters_connection_string.clone()),
             );
         }
-        if let Some(column_masking_connection_string) = &self.column_masking_connection_string {
+        if let Some(batched_column_masking_connection_string) =
+            &self.batched_column_masking_connection_string
+        {
             config.insert(
-                "opa.policy.column-masking-uri".to_string(),
-                Some(column_masking_connection_string.clone()),
+                "opa.policy.batch-column-masking-uri".to_string(),
+                Some(batched_column_masking_connection_string.clone()),
             );
         }
         if self.allow_permission_management_operations {

rust/operator-binary/src/controller.rs

@@ -2064,8 +2064,8 @@ mod tests {
                 "http://simple-opa.default.svc.cluster.local:8081/v1/data/my-product/rowFilters"
                     .to_string(),
             ),
-            column_masking_connection_string: Some(
-                "http://simple-opa.default.svc.cluster.local:8081/v1/data/my-product/columnMask"
+            batched_column_masking_connection_string: Some(
+                "http://simple-opa.default.svc.cluster.local:8081/v1/data/my-product/batchColumnMasks"
                     .to_string(),
             ),
             allow_permission_management_operations: true,
@@ -2167,7 +2167,7 @@ mod tests {
         assert!(access_control_config.contains("foo.bar=true"));
         assert!(access_control_config.contains("opa.allow-permission-management-operations=false"));
         assert!(access_control_config.contains(r#"opa.policy.batched-uri=http\://simple-opa.default.svc.cluster.local\:8081/v1/data/my-product/batch-new"#));
-        assert!(access_control_config.contains(r#"opa.policy.column-masking-uri=http\://simple-opa.default.svc.cluster.local\:8081/v1/data/my-product/columnMask"#));
+        assert!(access_control_config.contains(r#"opa.policy.batch-column-masking-uri=http\://simple-opa.default.svc.cluster.local\:8081/v1/data/my-product/batchColumnMasks"#));
         assert!(access_control_config.contains(r#"opa.policy.row-filters-uri=http\://simple-opa.default.svc.cluster.local\:8081/v1/data/my-product/rowFilters"#));
         assert!(access_control_config.contains(r#"opa.policy.uri=http\://simple-opa.default.svc.cluster.local\:8081/v1/data/my-product/allow"#));
     }

rust/operator-binary/src/crd/mod.rs

@@ -312,10 +312,23 @@ pub mod versioned {
 
     #[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
     #[serde(rename_all = "camelCase")]
-    pub struct TrinoAuthorization {
+    pub enum TrinoAuthorization {
+        Opa {
+            // no doc - it's in the struct.
+            #[serde(default, flatten)]
+            config: TrinoAuthorizationOpaConfig,
+        },
+    }
+
+    #[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
+    #[serde(rename_all = "camelCase")]
+    pub struct TrinoAuthorizationOpaConfig {
         // no doc - it's in the struct.
-        #[serde(default, skip_serializing_if = "Option::is_none")]
-        pub opa: Option<OpaConfig>,
+        #[serde(flatten)]
+        pub opa: OpaConfig,
+
+        #[serde(default = "TrinoAuthorization::enabled_column_masking_default")]
+        pub enable_column_masking: bool,
     }
 
     #[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
@@ -365,6 +378,12 @@ pub mod versioned {
     }
 }
 
+impl v1alpha1::TrinoAuthorization {
+    pub fn enabled_column_masking_default() -> bool {
+        true
+    }
+}
+
 impl Default for v1alpha1::TrinoCoordinatorRoleConfig {
     fn default() -> Self {
         v1alpha1::TrinoCoordinatorRoleConfig {
@@ -877,12 +896,21 @@ impl v1alpha1::TrinoCluster {
         !spec.cluster_config.authentication.is_empty()
     }
 
-    pub fn get_opa_config(&self) -> Option<&OpaConfig> {
+    pub fn column_masking_enabled(&self) -> bool {
+        match self.get_opa_config() {
+            Some(a) => a.enable_column_masking,
+            None => v1alpha1::TrinoAuthorization::enabled_column_masking_default(),
+        }
+    }
+
+    pub fn get_opa_config(&self) -> Option<&v1alpha1::TrinoAuthorizationOpaConfig> {
         self.spec
             .cluster_config
             .authorization
             .as_ref()
-            .and_then(|a| a.opa.as_ref())
+            .and_then(|a| match a {
+                v1alpha1::TrinoAuthorization::Opa { config } => Some(config),
+            })
     }
 
     /// Return user provided server TLS settings

rust/operator-binary/src/main.rs

@@ -238,9 +238,8 @@ fn references_config_map(
     };
 
     match &trino.spec.cluster_config.authorization {
-        Some(trino_authorization) => match &trino_authorization.opa {
-            Some(opa_config) => opa_config.config_map_name == config_map.name_any(),
-            None => false,
+        Some(trino_authorization) => match &trino_authorization {
+            v1alpha1::TrinoAuthorization::Opa { config } => config.opa.config_map_name == config_map.name_any(),
         },
         None => false,
     }