From d1ccaa5bb5dcdaa215b23fa5d42f6aae68521fa3 Mon Sep 17 00:00:00 2001
From: Masaki Yatsu <yatsu@yatsu.info>
Date: Thu, 11 Sep 2025 03:59:00 +0900
Subject: [PATCH] fix(jupyterhub): airflow integration

---
 .../images/datastack-cuda-notebook/Dockerfile       |  2 +-
 jupyterhub/images/datastack-notebook/Dockerfile     |  2 +-
 jupyterhub/jupyterhub-values.gomplate.yaml          | 13 ++++++++++++-
 jupyterhub/justfile                                 |  2 +-
 4 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/jupyterhub/images/datastack-cuda-notebook/Dockerfile b/jupyterhub/images/datastack-cuda-notebook/Dockerfile
index 95791ce..1ac920c 100644
--- a/jupyterhub/images/datastack-cuda-notebook/Dockerfile
+++ b/jupyterhub/images/datastack-cuda-notebook/Dockerfile
@@ -77,7 +77,6 @@ RUN pip uninstall -y \
 # 2. Find the pandas version in the file `dev/infra/Dockerfile`.
 RUN mamba install --yes \
     'aif360' \
-    'airflow' \
     'chromadb' \
     'clickhouse-connect' \
     'csvkit' \
@@ -142,6 +141,7 @@ RUN mamba install --yes \
 
 RUN pip install \
     agno \
+    apache-airflow \
     apache-airflow-client \
     fastembed \
     feature-engine \
diff --git a/jupyterhub/images/datastack-notebook/Dockerfile b/jupyterhub/images/datastack-notebook/Dockerfile
index 8f05850..5be7f9b 100644
--- a/jupyterhub/images/datastack-notebook/Dockerfile
+++ b/jupyterhub/images/datastack-notebook/Dockerfile
@@ -77,7 +77,6 @@ RUN pip uninstall -y \
 # 2. Find the pandas version in the file `dev/infra/Dockerfile`.
 RUN mamba install --yes \
     'aif360' \
-    'airflow' \
     'chromadb' \
     'clickhouse-connect' \
     'csvkit' \
@@ -142,6 +141,7 @@ RUN mamba install --yes \
 
 RUN pip install \
     agno \
+    apache-airflow \
     apache-airflow-client \
     fastembed \
     feature-engine \
diff --git a/jupyterhub/jupyterhub-values.gomplate.yaml b/jupyterhub/jupyterhub-values.gomplate.yaml
index 515b5a3..84a94bf 100644
--- a/jupyterhub/jupyterhub-values.gomplate.yaml
+++ b/jupyterhub/jupyterhub-values.gomplate.yaml
@@ -178,12 +178,23 @@ singleuser:
           optional: true  # Don't fail if PVC doesn't exist yet
     extraVolumeMounts:
       - name: airflow-dags
-        mountPath: /opt/airflow-dags
+        mountPath: /home/jovyan/airflow-dags
         readOnly: false
     {{- end }}
 
   networkPolicy:
     egress:
+      {{- if eq .Env.JUPYTERHUB_AIRFLOW_DAGS_PERSISTENCE_ENABLED "true" }}
+      # Allow communication with Airflow API server in the same namespace
+      - to:
+          - podSelector:
+              matchLabels:
+                release: airflow
+                component: api-server
+        ports:
+          - port: 8080
+            protocol: TCP
+      {{- end }}
       - to:
           - namespaceSelector:
               matchLabels:
diff --git a/jupyterhub/justfile b/jupyterhub/justfile
index 8db61db..576e9f2 100644
--- a/jupyterhub/justfile
+++ b/jupyterhub/justfile
@@ -9,7 +9,7 @@ export JUPYTERHUB_NFS_PV_ENABLED := env("JUPYTERHUB_NFS_PV_ENABLED", "")
 export JUPYTERHUB_STORAGE_CLASS := env("JUPYTERHUB_STORAGE_CLASS", "")
 export JUPYTERHUB_VAULT_INTEGRATION_ENABLED := env("JUPYTERHUB_VAULT_INTEGRATION_ENABLED", "")
 export JUPYTERHUB_AIRFLOW_DAGS_PERSISTENCE_ENABLED := env("JUPYTERHUB_AIRFLOW_DAGS_PERSISTENCE_ENABLED", "")
-export JUPYTER_PYTHON_KERNEL_TAG := env("JUPYTER_PYTHON_KERNEL_TAG", "python-3.12-36")
+export JUPYTER_PYTHON_KERNEL_TAG := env("JUPYTER_PYTHON_KERNEL_TAG", "python-3.12-37")
 export KERNEL_IMAGE_BUUN_STACK_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_REPOSITORY", "buun-stack-notebook")
 export KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY", "buun-stack-cuda-notebook")
 export JUPYTER_PROFILE_MINIMAL_ENABLED := env("JUPYTER_PROFILE_MINIMAL_ENABLED", "false")