airflow-unicore-integration 0.1.12__py3-none-any.whl → 0.2.0__py3-none-any.whl

airflow_unicore_integration/hooks/unicore_hooks.py

@@ -1,8 +1,14 @@
 from __future__ import annotations
 
-from airflow.hooks.base import BaseHook
+import logging
+from typing import Any
+
+from airflow.providers.common.compat.sdk import BaseHook
 from pyunicore import client
 from pyunicore import credentials
+from wtforms import StringField
+
+logger = logging.getLogger(__name__)
 
 
 class UnicoreHook(BaseHook):
@@ -23,25 +29,59 @@ class UnicoreHook(BaseHook):
         super().__init__()
         self.uc_conn_id = uc_conn_id
 
+    @classmethod
+    def get_connection_form_fields(cls):
+        return {"auth_token": StringField("Auth Token")}
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Return custom UI field behaviour for UNICORE connection."""
+        return {
+            "hidden_fields": ["schema", "port", "extra"],
+            "relabeling": {
+                "login": "Username",
+            },
+            "placeholder": {"auth_token": "UNICORE auth token"},
+        }
+
     def get_conn(
         self,
         overwrite_base_url: str | None = None,
         overwrite_credential: credentials.Credential | None = None,
     ) -> client.Client:
         """Return a Unicore Client. base_url and credentials may be overwritten."""
-        self.log.debug(
+        logger.debug(
             f"Gettig connection with id '{self.uc_conn_id}' from secrets backend. Will be modified with user input for UNICORE."
         )
         params = self.get_connection(self.uc_conn_id)
         base_url = params.host
         credential = credentials.UsernamePassword(params.login, params.password)
+        auth_token = params.extra_dejson.get("auth_token", None)
+        if auth_token is not None:
+            credential = credentials.create_credential(token=auth_token)
         if overwrite_base_url is not None:
             base_url = overwrite_base_url
         if overwrite_credential is not None:
             credential = overwrite_credential
+        if not base_url:
+            raise TypeError()
         conn = client.Client(credential, base_url)
         return conn
 
+    def get_credential(self) -> credentials.Credential:
+        """Return the credential part of the connection as a Credential object."""
+        params = self.get_connection(self.uc_conn_id)
+        credential = credentials.UsernamePassword(params.login, params.password)
+        auth_token = params.extra_dejson.get("auth_token", None)
+        if auth_token is not None:
+            credential = credentials.create_credential(token=auth_token)
+        return credential
+
+    def get_base_url(self) -> str:
+        """Return the base url of the connection."""
+        params = self.get_connection(self.uc_conn_id)
+        return params.host
+
     def test_connection(self) -> tuple[bool, str]:
         """Test the connection by sending an access_info request"""
         conn = self.get_conn()

airflow_unicore_integration/operators/container.py

@@ -0,0 +1,22 @@
+from typing import Dict
+from typing import List
+
+from airflow_unicore_integration.operators.unicore_operators import (
+    UnicoreGenericOperator,
+)
+
+
+class UnicoreContainerOperator(UnicoreGenericOperator):
+    def __init__(
+        self, name: str, docker_image_url: str, command: str, options: str | None = None, **kwargs
+    ):
+        params: Dict[str, str | List[str]] = {"COMMAND": command, "IMAGE_URL": docker_image_url}
+        if options is not None:
+            params["OPTIONS"] = options
+        super().__init__(
+            name=name,
+            application_name="CONTAINER",
+            application_version="1.0",
+            parameters=params,
+            **kwargs,
+        )

airflow_unicore_integration/operators/unicore_operators.py

@@ -6,7 +6,7 @@ from typing import Sequence
 
 import pyunicore.client as uc_client
 import pyunicore.credentials as uc_credentials
-from airflow.models.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
 from airflow.utils.context import Context
 
 from airflow_unicore_integration.hooks import unicore_hooks
@@ -62,6 +62,7 @@ class UnicoreGenericOperator(BaseOperator):
         credential_username: str | None = None,
         credential_password: str | None = None,
         credential_token: str | None = None,
+        conn_id: str | None = None,
         **kwargs,
     ):
         """
@@ -114,6 +115,7 @@ class UnicoreGenericOperator(BaseOperator):
         self.credential_username = credential_username
         self.credential_password = credential_password
         self.credential_token = credential_token
+        self.conn_id = conn_id
 
         self.validate_job_description()
         logger.debug("created Unicore Job Task")
@@ -262,7 +264,7 @@ class UnicoreGenericOperator(BaseOperator):
 
     def execute_async(self, context: Context) -> Any:
         """Submits the job and returns the job object without waiting for it to finish."""
-        client = self.get_uc_client()
+        client = self.get_uc_client(self.conn_id)
         job = client.new_job(job_description=self.get_job_description(), inputs=[])
         return job
 

{airflow_unicore_integration-0.1.12.dist-info → airflow_unicore_integration-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.1.12
+Version: 0.2.0
 Summary: Running Unicore Jobs from airflow DAGs.
 Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
 License-Expression: BSD-3-Clause
@@ -24,7 +24,7 @@ Unicore Airflow Integration
 
 |Generic badge|
 
-.. |Generic badge| image:: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml/badge.svg 
+.. |Generic badge| image:: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml/badge.svg
    :target: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml
 
 This project integrates `UNICORE <https://github.com/UNICORE-EU>`_ and `Apache Airflow <https://airflow.apache.org/>`_.
@@ -34,6 +34,31 @@ Airflow is a platform to programmatically author, schedule and monitor workflows
 In the current state, this projects provides a set of airflow `operators <https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/operators.html>`_, which can be used as part of airflow workflows to submit jobs to Unicore.
 The UnicoreExecutor only offers experimental support for airflow 3 so far. Further support is currently being worked on.
 
+-------------------------
+Using the UnicoreExecutor
+-------------------------
+
+To use the UnicoreExecutor, this library needs to be installed in your airflow environment and then some configuration work needs to be done.
+
+How to configure these settigns is up to your deployment, as it uses the standard airflow configuration mechanism.
+In case of a helm deployemnt via the official helm chart, you will need to use environment variables, as all unicore related options are not present in the chart and will cause schema-validation to fail.
+
+All options fall under the [unicore.executor] section in airflow.cfg, or have the ``AIRFLOW__UNICORE_EXECUTOR__`` prefix as an environment variable.
+
+========================= ============================================ ===========================================================================================
+Option name               default                                      description
+========================= ============================================ ===========================================================================================
+EXECUTION_API_SERVER_URL  <The default from the airflow config>        The url to reach the airflow API server from the execution environment (e.g. compute nodes)
+AUTH_TOKEN                mandatory                                    The unicore auth token to use for job submission
+DEFAULT_URL               http://localhost:8080/DEMO-SITE/rest/core    The default unicore site to submit jobs to
+DEFAULT_ENV               mandatory                                    The default activation script for a functional airflow environment on the execution machine
+TMP_DIR                   /tmp                                         A temporary directory to store data such as GitDagBundles
+========================= ============================================ ===========================================================================================
+
+The default env is loaded via ``. default_env.sh``, and must enable an environment, where python is available in a suitable version, and the ``apache-airflow-task-sdk`` and ``apache-airflow-providers-git`` packages are available. All other dependencies depend on the dags to be run, but must already be included in the environment.
+
+A simple solution for this may be the "activate" script for a python venv. If the target systems requires additional commands to enable python (e.g. ``module load``), these may be added to the top of the activate script.
+
 ---------------------------
 Using the Unicore Operators
 ---------------------------

{airflow_unicore_integration-0.1.12.dist-info → airflow_unicore_integration-0.2.0.dist-info}/RECORD

@@ -3,15 +3,16 @@ airflow_unicore_integration/executors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCe
 airflow_unicore_integration/executors/run_task_via_supervisor.py,sha256=3ErgPf-Oy3B4Di5yNXhhPkaojIJykvCxMZ9MlKSYPI8,2756
 airflow_unicore_integration/executors/unicore_executor.py,sha256=E1nOskWSBmC-ReLRvA8E3bY-G0lpxP403tazlBNhgFQ,5919
 airflow_unicore_integration/hooks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-airflow_unicore_integration/hooks/unicore_hooks.py,sha256=JjcjogWtN1xveagpkraQuYOdXjkp2lSnEdQc0waqhU4,1662
+airflow_unicore_integration/hooks/unicore_hooks.py,sha256=NEJ6sgQCD9AyZO9TA3vOR8RdijqaIX_FqY-Abz74Yys,3148
 airflow_unicore_integration/operators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-airflow_unicore_integration/operators/unicore_operators.py,sha256=xYX1t_QzjpqHZDoEu6jJNNXVmBEnjIlApFWvgYoYoB0,16790
+airflow_unicore_integration/operators/container.py,sha256=D-2La59XZazXvBDVn-_wFIZ0IrCYeGy18miK8S5Usmc,693
+airflow_unicore_integration/operators/unicore_operators.py,sha256=PrcB-riakRUt8DHR9Hkk0ccOe-xk6NjJmNAmDpNgpN4,16868
 airflow_unicore_integration/policies/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 airflow_unicore_integration/util/job.py,sha256=vYyaNzGw2EQWQlDTGO6q4T1gZTHoB9-RIb5KYYd6pAc,10556
 airflow_unicore_integration/util/launch_script_content.py,sha256=42_aFpaCMmvFmmUxQDGcudkleX1YSK_yYWE8T41NOy0,2915
-airflow_unicore_integration-0.1.12.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
-airflow_unicore_integration-0.1.12.dist-info/METADATA,sha256=aK_J520LHA3xOf4vqqpo71zcJ6WDrVTiMf120Ph00R8,11189
-airflow_unicore_integration-0.1.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-airflow_unicore_integration-0.1.12.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
-airflow_unicore_integration-0.1.12.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
-airflow_unicore_integration-0.1.12.dist-info/RECORD,,
+airflow_unicore_integration-0.2.0.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
+airflow_unicore_integration-0.2.0.dist-info/METADATA,sha256=-ifD7f7RX30qQs3OMzsXKeAvo_acGHLigthKu93-rU0,13696
+airflow_unicore_integration-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airflow_unicore_integration-0.2.0.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
+airflow_unicore_integration-0.2.0.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
+airflow_unicore_integration-0.2.0.dist-info/RECORD,,