Allow arbitrary python code (#493)

* Bumping operator-rs to 0.74.0

* Adding docs

* Allow arbitrary python code with Footer and Header keys

* Update Changelog

* Adding some tests for new file vars

* Reworking tests, green now

* Code Review

Co-authored-by: Sebastian Bernauer <[email protected]>

---------

Co-authored-by: Sebastian Bernauer <[email protected]>

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Allowing arbitrary python code as `EXPERIMENTAL_FILE_HEADER` and `EXPERIMENTAL_FILE_FOOTER` in `webserver_config.py` ([#493]).
+
 ### Changed
 
 - Reduce CRD size from `1.7MB` to `111KB` by accepting arbitrary YAML input instead of the underlying schema for the following fields ([#488]):
@@ -16,6 +20,7 @@
 
 [#488]: https://github.com/stackabletech/airflow-operator/pull/488
 [#489]: https://github.com/stackabletech/airflow-operator/pull/489
+[#493]: https://github.com/stackabletech/airflow-operator/pull/493
 
 ## [24.7.0] - 2024-07-24
 

docs/modules/airflow/pages/usage-guide/overrides.adoc

@@ -10,6 +10,33 @@ that each component has the same configuration: not all components use each sett
 
 Airflow exposes an environment variable for every Airflow configuration setting, a list of which can be found in the https://airflow.apache.org/docs/apache-airflow/stable/configurations-ref.html[Configuration Reference].
 
+As Airflow can be configured with python code too, arbitrary code can be added to the `webserver_config.py`.
+You can use either `EXPERIMENTAL_FILE_HEADER` to add code to the top or `EXPERIMENTAL_FILE_FOOTER` to add to the bottom.
+
+IMPORTANT: This is an experimental feature
+
+[source,yaml]
+----
+webservers:
+  configOverrides:
+    webserver_config.py:
+      CSV_EXPORT: "{'encoding': 'utf-8'}"
+      EXPERIMENTAL_FILE_HEADER: |
+        from modules.my_module import my_class
+      EXPERIMENTAL_FILE_FOOTER: |
+        import logging
+        from airflow.security import AirflowSecurityManager
+
+        class myCustomSecurityManger(AirflowSecurityManager):
+          def __init__():
+            init()
+
+        CUSTOM_SECURITY_MANAGER = myCustomSecurityManger
+  roleGroups:
+    default:
+      config: {}
+----
+
 Although Kubernetes can override these settings in one of two ways (Configuration overrides, or Environment Variable overrides), the affect is the same
 and currently only the latter is implemented. This is described in the following section.
 

rust/operator-binary/src/airflow_controller.rs

@@ -50,7 +50,10 @@ use stackable_operator::{
     },
     kvp::{Label, LabelError, Labels},
     logging::controller::ReconcilerError,
-    product_config_utils::{transform_all_roles_to_config, validate_all_roles_and_groups_config},
+    product_config_utils::{
+        transform_all_roles_to_config, validate_all_roles_and_groups_config,
+        CONFIG_OVERRIDE_FILE_FOOTER_KEY, CONFIG_OVERRIDE_FILE_HEADER_KEY,
+    },
     product_logging::{
         self,
         spec::{ContainerLogConfig, Logging},
@@ -65,6 +68,7 @@ use stackable_operator::{
 };
 use std::{
     collections::{BTreeMap, HashMap},
+    io::Write,
     str::FromStr,
     sync::Arc,
 };
@@ -280,6 +284,11 @@ pub enum Error {
 
     #[snafu(display("failed to construct config"))]
     ConstructConfig { source: config::Error },
+
+    #[snafu(display(
+        "failed to write to String (Vec<u8> to be precise) containing Airflow config"
+    ))]
+    WriteToConfigFileString { source: std::io::Error },
 }
 
 type Result<T, E = Error> = std::result::Result<T, E>;
@@ -639,6 +648,15 @@ fn build_rolegroup_config_map(
     config::add_airflow_config(&mut config, authentication_config).context(ConstructConfigSnafu)?;
 
     let mut config_file = Vec::new();
+
+    // By removing the keys from `config_properties`, we avoid pasting the Python code into a Python variable as well
+    // (which would be bad)
+    if let Some(header) = config.remove(CONFIG_OVERRIDE_FILE_HEADER_KEY) {
+        writeln!(config_file, "{}", header).context(WriteToConfigFileStringSnafu)?;
+    }
+
+    let temp_file_footer: Option<String> = config.remove(CONFIG_OVERRIDE_FILE_FOOTER_KEY);
+
     flask_app_config_writer::write::<AirflowConfigOptions, _, _>(
         &mut config_file,
         config.iter(),
@@ -648,6 +666,10 @@ fn build_rolegroup_config_map(
         rolegroup: rolegroup.clone(),
     })?;
 
+    if let Some(footer) = temp_file_footer {
+        writeln!(config_file, "{}", footer).context(WriteToConfigFileStringSnafu)?;
+    }
+
     let mut cm_builder = ConfigMapBuilder::new();
 
     cm_builder

tests/templates/kuttl/smoke/40-install-airflow-cluster.yaml.j2

@@ -45,9 +45,20 @@ spec:
     config:
       logging:
         enableVectorAgent: {{ lookup('env', 'VECTOR_AGGREGATOR') | length > 0 }}
+    configOverrides:
+      webserver_config.py:
+        EXPERIMENTAL_FILE_HEADER: |
+          COMMON_HEADER_VAR = "role-value"
+          ROLE_HEADER_VAR = "role-value"
+        EXPERIMENTAL_FILE_FOOTER: |
+          ROLE_FOOTER_VAR = "role-value"
     roleGroups:
       default:
         replicas: 1
+        configOverrides:
+          webserver_config.py:
+            EXPERIMENTAL_FILE_HEADER: |
+              COMMON_HEADER_VAR = "group-value"
 {% if test_scenario['values']['executor'] == 'celery' %}
   celeryExecutors:
     config:

tests/templates/kuttl/smoke/41-assert.yaml

@@ -0,0 +1,11 @@
+---
+apiVersion: kuttl.dev/v1beta1
+kind: TestAssert
+timeout: 600
+commands:
+  #
+  # Test envOverrides
+  #
+  - script: |
+      kubectl -n $NAMESPACE get cm airflow-webserver-default -o yaml | yq -e '.data."webserver_config.py"' | grep "COMMON_HEADER_VAR = \"group-value\""
+      kubectl -n $NAMESPACE get cm airflow-webserver-default -o yaml | yq -e '.data."webserver_config.py"' | grep "ROLE_FOOTER_VAR = \"role-value\""