airflow-unicore-integration 0.1.11__py3-none-any.whl → 0.1.13__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,22 +29,43 @@ 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()
+        logger.info(f"Using credential {credential} for SITE {base_url}.")
         conn = client.Client(credential, base_url)
         return 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/util/job.py

@@ -25,6 +25,7 @@ class JobDescriptionGenerator:
     EXECUTOR_CONFIG_PROJECT = "Project"  # gets added to the unicore job description
     EXECUTOR_CONFIG_PRE_COMMANDS = "precommands"  # gets added to the unicore job description
     EXECUTOR_CONFIG_POST_COMMANDS = "postcommands"  # gets added to the unicore job descirption
+    EXECUTOR_CONFIG_JOB_TYPE = "job_type"
     EXECUTOR_CONFIG_UNICORE_CONN_KEY = (
         "unicore_connection_id"  # alternative connection id for the Unicore connection to use
     )
@@ -59,6 +60,7 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
         user_added_pre_commands: list[str] = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_PRE_COMMANDS, [])  # type: ignore
         user_defined_python_env: str = workload.ti.executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_PYTHON_ENV_KEY, None)  # type: ignore
         user_added_post_commands: list[str] = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_POST_COMMANDS, [])  # type: ignore
+        user_defined_job_type: str = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_JOB_TYPE, None)  # type: ignore
         # get local dag path from cmd and fix dag path in arguments
         dag_rel_path = str(workload.dag_rel_path)
         if dag_rel_path.startswith("DAG_FOLDER"):
@@ -71,6 +73,10 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
         )
         logger.debug(f"Server is {server}")
 
+        # set job type
+        if user_defined_job_type:
+            job_descr_dict["Job type"] = user_defined_job_type
+
         # check which python virtualenv to use
         if user_defined_python_env:
             python_env = user_defined_python_env

{airflow_unicore_integration-0.1.11.dist-info → airflow_unicore_integration-0.1.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.1.11
+Version: 0.1.13
 Summary: Running Unicore Jobs from airflow DAGs.
 Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
 License-Expression: BSD-3-Clause

{airflow_unicore_integration-0.1.11.dist-info → airflow_unicore_integration-0.1.13.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=NAR5YmnPVzSsqpNVpRHt1sS3_HkEP_e8qoCv7daUZAA,2579
 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=IsVjDiq3Jj-Eky0mET3awC3yBloHR9Duvm1KFiAPKpU,10264
+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.11.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
-airflow_unicore_integration-0.1.11.dist-info/METADATA,sha256=a6vLdzCg5UYaxOIsKCPfxQPDZ-sBg6rwJQWCVRoupWM,11189
-airflow_unicore_integration-0.1.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-airflow_unicore_integration-0.1.11.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
-airflow_unicore_integration-0.1.11.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
-airflow_unicore_integration-0.1.11.dist-info/RECORD,,
+airflow_unicore_integration-0.1.13.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
+airflow_unicore_integration-0.1.13.dist-info/METADATA,sha256=KXZXAPP9dv8EN2UheZDJ9TeXiIb5jL-t6AggYOJFDPU,11189
+airflow_unicore_integration-0.1.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airflow_unicore_integration-0.1.13.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
+airflow_unicore_integration-0.1.13.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
+airflow_unicore_integration-0.1.13.dist-info/RECORD,,