wip: kafka dag

Author: adwk67

demos/demos-v2.yaml

@@ -48,7 +48,7 @@ demos:
     manifests:
       - plainYaml: https://raw.githubusercontent.com/stackabletech/demos/main/demos/airflow-scheduled-job/01-airflow-spark-clusterrole.yaml
       - plainYaml: https://raw.githubusercontent.com/stackabletech/demos/main/demos/airflow-scheduled-job/02-airflow-spark-clusterrolebinding.yaml
-      - plainYaml: https://raw.githubusercontent.com/stackabletech/demos/main/demos/airflow-scheduled-job/03-enable-and-run-spark-dag.yaml
+      #- plainYaml: https://raw.githubusercontent.com/stackabletech/demos/main/demos/airflow-scheduled-job/03-enable-and-run-spark-dag.yaml
       - plainYaml: https://raw.githubusercontent.com/stackabletech/demos/main/demos/airflow-scheduled-job/04-enable-and-run-date-dag.yaml
     supportedNamespaces: []
     resourceRequests:

stacks/airflow/airflow.yaml

@@ -26,44 +26,101 @@ spec:
       - name: airflow-dags
         mountPath: /dags/pyspark_pi.yaml
         subPath: pyspark_pi.yaml
+      - name: airflow-dags
+        mountPath: /dags/kafka.py
+        subPath: kafka.py
   webservers:
     roleConfig:
       listenerClass: external-unstable
     config:
       resources:
         cpu:
-          min: 400m
-          max: "1"
+          min: "2"
+          max: "3"
         memory:
-          limit: 2Gi
-      gracefulShutdownTimeout: 30s
+          limit: 3Gi
+    envOverrides: &envOverrides
+      AIRFLOW__CORE__DAGS_FOLDER: "/dags"
+      PYTHONPATH: "/stackable/app/log_config:/dags"
+      AIRFLOW_CONN_KUBERNETES_IN_CLUSTER: "kubernetes://?__extra__=%7B%22extra__kubernetes__in_cluster%22%3A+true%2C+%22extra__kubernetes__kube_config%22%3A+%22%22%2C+%22extra__kubernetes__kube_config_path%22%3A+%22%22%2C+%22extra__kubernetes__namespace%22%3A+%22%22%7D"
+    podOverrides: &podOverrides
+      spec:
+        containers:
+          - name: airflow
+            env:
+              - name: KAFKA_BOOTSTRAP
+                valueFrom:
+                  configMapKeyRef:
+                    name: kafka
+                    key: KAFKA
+              - name: AIRFLOW_CONN_KAFKA_CONN # $(KAFKA_BOOTSTRAP)
+                value: "{\"conn_type\": \"kafka\", \"extra\": {\"bootstrap.servers\": \"kafka-broker-default-0-listener-broker.demo.svc.cluster.local:9092\", \"group.id\": \"airflow_group\", \"auto.offset.reset\": \"latest\"}}"
     roleGroups:
       default:
-        envOverrides: &envOverrides
-          AIRFLOW__CORE__DAGS_FOLDER: "/dags"
-          AIRFLOW_CONN_KUBERNETES_IN_CLUSTER: "kubernetes://?__extra__=%7B%22extra__kubernetes__in_cluster%22%3A+true%2C+%22extra__kubernetes__kube_config%22%3A+%22%22%2C+%22extra__kubernetes__kube_config_path%22%3A+%22%22%2C+%22extra__kubernetes__namespace%22%3A+%22%22%7D"
         replicas: 1
   kubernetesExecutors:
     envOverrides: *envOverrides
+    podOverrides: *podOverrides
   schedulers:
-    config:
-      gracefulShutdownTimeout: 30s
-      resources:
-        cpu:
-          min: 400m
-          max: "1"
-        memory:
-          limit: 1Gi
+    envOverrides: *envOverrides
+    podOverrides: *podOverrides
+    roleGroups:
+      default:
+        replicas: 1
+  dagProcessors:
+    envOverrides: *envOverrides
+    podOverrides: *podOverrides
+    roleGroups:
+      default:
+        replicas: 1
+  triggerers:
+    envOverrides: *envOverrides
+    podOverrides: *podOverrides
     roleGroups:
       default:
-        envOverrides: *envOverrides
         replicas: 1
 ---
 apiVersion: v1
 kind: ConfigMap
 metadata:
   name: airflow-dags
 data:
+  kafka.py: |
+    from airflow.providers.apache.kafka.triggers.msg_queue import KafkaMessageQueueTrigger
+    from airflow.providers.standard.operators.empty import EmptyOperator
+    from airflow.sdk import DAG, Asset, AssetWatcher
+
+    import logging
+    logger = logging.getLogger(__name__)
+
+    logger.info("✅ kafka.apply_function module imported")
+
+    def apply_function(message):
+      try:
+        logger.info("apply_function called")
+        logger.info("message payload: %r", message.value())
+        return True
+      except Exception:
+        logger.exception("apply_function failed")
+        return False
+
+    # Define a trigger that listens to an Apache Kafka message queue
+    trigger = KafkaMessageQueueTrigger(
+      topics=["test-topic"],
+      apply_function="kafka.apply_function",
+      kafka_config_id="kafka_conn",
+      apply_function_args=None,
+      apply_function_kwargs=None,
+      poll_timeout=1,
+      poll_interval=5,
+    )
+
+    # Define an asset that watches for messages on the queue
+    asset = Asset("kafka_queue_asset", watchers=[AssetWatcher(name="kafka_watcher", trigger=trigger)])
+
+    with DAG(dag_id="example_kafka_watcher", schedule=[asset]) as dag:
+      EmptyOperator(task_id="task")
+
   date_demo.py: |
     """Example DAG returning the current date"""
     from datetime import datetime, timedelta

stacks/airflow/kafka.yaml

@@ -0,0 +1,22 @@
+---
+apiVersion: kafka.stackable.tech/v1alpha1
+kind: KafkaCluster
+metadata:
+  name: kafka
+spec:
+  image:
+    productVersion: 4.1.0
+  clusterConfig:
+    tls:
+      serverSecretClass: null
+  controllers:
+    roleGroups:
+      default:
+        replicas: 1
+  brokers:
+    config:
+      bootstrapListenerClass: cluster-internal
+      brokerListenerClass: cluster-internal
+    roleGroups:
+      default:
+        replicas: 1

stacks/stacks-v2.yaml

@@ -139,12 +139,14 @@ stacks:
       - listener
       - secret
       - airflow
-      - spark-k8s # Some demo does schedule a Spark job
+      - spark-k8s
+      - kafka
     labels:
       - airflow
     manifests:
       - helmChart: https://raw.githubusercontent.com/stackabletech/demos/main/stacks/_templates/postgresql-airflow.yaml
-      - plainYaml: https://raw.githubusercontent.com/stackabletech/demos/main/stacks/airflow/airflow.yaml
+      - plainYaml: stacks/airflow/airflow.yaml
+      - plainYaml: stacks/airflow/kafka.yaml
     supportedNamespaces: []
     resourceRequests:
       cpu: 3400m