Override python3 kernelspec for notebook operator if uv is used

navikt · Jun 11, 2024 · 520a39c · 520a39c
1 parent 625a0cb
commit 520a39c
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/dataverk_airflow/notebook_operator.py b/dataverk_airflow/notebook_operator.py
@@ -73,7 +73,12 @@ def notebook_operator(
         except KeyError:
             image = os.getenv("KNADA_AIRFLOW_OPERATOR_IMAGE")
 
-    cmds = [f"papermill {Path(nb_path).name} output.ipynb"]
+    # Needs to override python3 kernelspec if using uv
+    if use_uv_pip_install:
+        cmds = [f"python -m ipykernel install --user --name=python3 --display-name 'Python (UV)' && papermill {Path(nb_path).name} output.ipynb"]
+    else:
+        cmds = [f"papermill {Path(nb_path).name} output.ipynb"]
+
     if log_output:
         cmds[-1] += " --log-output"
     if override_notebook_kernelspec: