airflow-unicore-integration 0.0.4__tar.gz → 0.1.0__tar.gz

airflow_unicore_integration-0.1.0/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) Forschungszentrum Juelich GmbH
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the names of the copyright holders nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

airflow_unicore_integration-0.1.0/PKG-INFO

@@ -0,0 +1,169 @@
+Metadata-Version: 2.4
+Name: airflow-unicore-integration
+Version: 0.1.0
+Summary: Running Unicore Jobs from airflow DAGs.
+Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/UNICORE-EU/airflow-unicore-integration
+Project-URL: Issues, https://github.com/UNICORE-EU/airflow-unicore-integration/issues
+Classifier: Development Status :: 4 - Beta
+Classifier: Framework :: Apache Airflow :: Provider
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/x-rst
+License-File: LICENSE
+Requires-Dist: pyunicore>=1.0.0
+Requires-Dist: apache-airflow>=3.0.0
+Dynamic: license-file
+
+===========================
+Unicore Airflow Integration
+===========================
+
+
+|Generic badge|
+
+.. |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/>`_.
+UNICORE is a software suite that, among other functions, provides seamless access to high-performance compute and data resources.
+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 Unicore Operators
+---------------------------
+
+There are multiple Unicore operators provided by this package. The most versatile one is the ``UnicoreGenericOperator``, which supports a lot of job parameters.
+All other operators are intended to offer a slightly less complex constructor, and therefore simpler usage, but all generic parameters are still available to be used.
+
+All operators support all possible parameters of the `Unicore job description <https://unicore-docs.readthedocs.io/en/latest/user-docs/rest-api/job-description/index.html#overview>`_. Here is an excerpt containing some commonly used parameters:
+
+======================= ======================= =========================================== ====================
+parameter name          type                    default                                     description
+======================= ======================= =========================================== ====================
+application_name        str                     None                                        Application Name
+application_version     str                     None                                        Application Version
+executable              str                     None                                        Command line executable
+arguments               List(str)               None                                        Command line arguments
+environment             Map(str,str)            None                                        environment arguments
+parameters              Map                     None                                        Application Parameters
+project                 str                     None                                        Accounting Project
+imports                 List(imports)           None                                        Stage-in/data import - see Unicore docs
+exports                 List(exports)           None                                        Stage-out/data export - see Unicore docs
+======================= ======================= =========================================== ====================
+
+For imports and exports go `here <https://unicore-docs.readthedocs.io/en/latest/user-docs/rest-api/job-description/index.html#importing-files-into-the-job-workspace>`_ for details.
+
+
+The ``UnicoreGenericOperator`` supports the following additional parameters:
+
+======================= ======================= =========================================== ====================
+parameter name          type                    default                                     description
+======================= ======================= =========================================== ====================
+name                    str                     None                                        name for the airflow task and the Unicore job
+xcom_output_files       List(str)               ["stdout","stderr"]                         list of files of which the content should be put into xcoms
+base_url                str                     configured in airflow connections or None   The base URL of the UNICOREX server to be used for the Unicore client
+credential              pyunicore credential    configured in airflow connections or None   A Unicore Credential to be used for the Unicore client
+credential_username     str                     configured in airflow connections or None   Username for the Unicore client credentials
+credential_password     str                     configured in airflow connections or None   Password the the Unicore client credentials
+credential_token        str                     configured in airflow connections or None   An OIDC token to be used by the Unicore client
+======================= ======================= =========================================== ====================
+
+
+The ``UnicoreScriptOperator`` offers a way to more easily submit a script as a job, where the script content can be provided as a string.
+
+======================= ======================= =========================================== ====================
+parameter name          type                    default                                     description
+======================= ======================= =========================================== ====================
+script_content          str                     None                                        The content of the script file
+======================= ======================= =========================================== ====================
+
+
+The ``UnicoreBSSOperator`` offers a way to directly submit batch-scripts from their content-strings.
+
+======================= ======================= =========================================== ====================
+parameter name          type                    default                                     description
+======================= ======================= =========================================== ====================
+bss_file_content        str                     None                                        The content of the batch script file
+======================= ======================= =========================================== ====================
+
+
+The ``UnicoreExecutableOperator`` offers a reduced constructor that only requires an executable.
+
+======================= ======================= =========================================== ====================
+parameter name          type                    default                                     description
+======================= ======================= =========================================== ====================
+executable              str                     None                                        The executable to run for this job
+xcom_output_files       List(str)               ["stdout","stderr"]                         list of files of which the content should be put into xcoms
+======================= ======================= =========================================== ====================
+
+The ``UnicoreDateOperator`` is more of a testing operator, since it will only run the ``date`` executable.
+
+-------------------------------
+Behaviour on Errors and Success
+-------------------------------
+
+The Unicore Operators do not do a lot of error and exception handling, and mostly just forward any problems to be handled by airflow.
+All of the Unicore logic is handled by the `pyunicore library <https://github.com/HumanBrainProject/pyunicore>`_.
+
+While some validation of the resulting Unicore job description is done automatically, it may still be possible to build an invalid job description with the operators.
+This may lead to a submission failure with Unicore. In this case, an exception is thrown to be handled by airflow.
+
+
+For a successful job submission, the job exit code is returned as the task return value, so that airflow can handle non-zero exit codes.
+All operators will also append the content of the job-log-file from Unicore to the airflow task log.
+Also, some job results and values will be uploaded via airflow-x-coms as well:
+
+======================= ========================================
+xcom name               description
+======================= ========================================
+Unicore Job ID          the Unicore ID for the job
+Unicore Job             the TSI script that was submitted by Unicore
+BSS_SUBMIT              the bss_script submitted by Unicore
+status_message          the status message for the Unicore job
+log                     the Unicore job log
+workdir_content         content of the job workdir upon completion
+[xcom_output_files]     content of each file in their own xcom, by default stdout and stderr
+======================= ========================================
+
+------------
+Example DAGs
+------------
+
+There are some example DAGs in this repository under ``project-dir/dags``.
+
+- ``unicore-test-1.py`` just shows basic date and executable usage.
+- ``unicore-test-2.py`` has some basic examples for the generic operator.
+- ``unicore-test-3.py`` also includes script-operator examples.
+- ``unicore-test-4.py`` has some examples with more arguments.
+- ``unicore-test-bss.py`` shows how bss submission can be done (very simple example).
+- ``unicore-test-credentials.py`` demonstrates that not only the credentials from the airflow connections backend can be used, but they can also be provided in the constructor of the operator.
+- ``unicore-test-import-export.py`` gives short examples for the imports and exports usage.
+
+
+-----------------
+Setup testing env
+-----------------
+
+Ensure a current version of docker is installed.
+
+Run ``python3 -m build`` to build the python package.
+
+Run the ``testing-env/build-image.sh`` script to create the customized airflow image, which will contain the newly build python package.
+
+Run ``testing-env/run-testing-env.sh init`` to initialize the airflow containers, database etc. This only needs to be done once.
+
+Run ``testing-env/run-testing-env.sh up`` to start the local airflow and Unicore deployment. Airflow will be available on port 8080, Unicore on port 8081.
+
+The ``run-testing-env.sh`` script supports the commands up, down, start, stop, ps and init for matching docker compose functions.
+
+-----------------------
+Install package via pip
+-----------------------
+
+``pip install airflow-unicore-integration``

{airflow_unicore_integration-0.0.4 → airflow_unicore_integration-0.1.0}/README.rst

@@ -3,6 +3,18 @@ Unicore Airflow Integration
 ===========================
 
 
+|Generic badge|
+
+.. |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/>`_.
+UNICORE is a software suite that, among other functions, provides seamless access to high-performance compute and data resources.
+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 Unicore Operators
 ---------------------------
@@ -29,7 +41,7 @@ exports                 List(exports)           None
 For imports and exports go `here <https://unicore-docs.readthedocs.io/en/latest/user-docs/rest-api/job-description/index.html#importing-files-into-the-job-workspace>`_ for details.
 
 
-The ``UnicoreGenericOperator`` supports the following additional parameters: 
+The ``UnicoreGenericOperator`` supports the following additional parameters:
 
 ======================= ======================= =========================================== ====================
 parameter name          type                    default                                     description
@@ -80,7 +92,7 @@ Behaviour on Errors and Success
 The Unicore Operators do not do a lot of error and exception handling, and mostly just forward any problems to be handled by airflow.
 All of the Unicore logic is handled by the `pyunicore library <https://github.com/HumanBrainProject/pyunicore>`_.
 
-While some validation of the resulting Unicore job description is done automatically, it may still be possible to build an invalid job description with the operators. 
+While some validation of the resulting Unicore job description is done automatically, it may still be possible to build an invalid job description with the operators.
 This may lead to a submission failure with Unicore. In this case, an exception is thrown to be handled by airflow.
 
 
@@ -110,10 +122,10 @@ There are some example DAGs in this repository under ``project-dir/dags``.
 - ``unicore-test-2.py`` has some basic examples for the generic operator.
 - ``unicore-test-3.py`` also includes script-operator examples.
 - ``unicore-test-4.py`` has some examples with more arguments.
-- ``unicore-test-bss.py`` shows how bss submission can be done (very simple example). 
-- ``unicore-test-credentials.py`` demonstrates that not only the credentials from the airflow connections backend can be used, but they can also be provided in the constructor of the o`perator.
-- ``unicore-test-import-export.py`` gives shprt examples for the imports and exports usage.
- 
+- ``unicore-test-bss.py`` shows how bss submission can be done (very simple example).
+- ``unicore-test-credentials.py`` demonstrates that not only the credentials from the airflow connections backend can be used, but they can also be provided in the constructor of the operator.
+- ``unicore-test-import-export.py`` gives short examples for the imports and exports usage.
+
 
 -----------------
 Setup testing env
@@ -135,4 +147,4 @@ The ``run-testing-env.sh`` script supports the commands up, down, start, stop, p
 Install package via pip
 -----------------------
 
-``pip install airflow-unicore-integration --index-url https://gitlab.jsc.fz-juelich.de/api/v4/projects/6269/packages/pypi/simple``
+``pip install airflow-unicore-integration``

airflow_unicore_integration-0.1.0/pyproject.toml

@@ -0,0 +1,63 @@
+[build-system]
+requires = [
+  "setuptools>=61.0",
+  "wheel"
+  ]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "airflow-unicore-integration"
+version = "0.1.0"
+authors = [
+  { name="Christian Böttcher", email="c.boettcher@fz-juelich.de" },
+]
+description = "Running Unicore Jobs from airflow DAGs."
+readme = "README.rst"
+requires-python = ">=3.9"
+license = "BSD-3-Clause"
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Framework :: Apache Airflow :: Provider",
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+
+
+
+dependencies = [
+    "pyunicore>=1.0.0",
+    "apache-airflow>=3.0.0"
+]
+
+[project.urls]
+Homepage = "https://github.com/UNICORE-EU/airflow-unicore-integration"
+Issues = "https://github.com/UNICORE-EU/airflow-unicore-integration/issues"
+
+[project.entry-points."apache_airflow_provider"]
+provider_info = "airflow_unicore_integration:get_provider_info"
+
+[project.entry-points.'airflow.policy']
+_ = 'airflow_unicore_integration.policies'
+
+[tool.pytest.ini_options]
+pythonpath = [
+  "src"
+]
+
+[tool.black]
+line-length = 100
+
+[tool.flake8]
+max-line-length = 100
+ignore = [
+ "N999",
+ "E501",
+ "W503"
+]
+per-file-ignores = [
+    "__init__.py:F401,E501",
+    "_version.py:E203",
+]
+
+[tool.mypy]
+ignore_missing_imports=true

{airflow_unicore_integration-0.0.4 → airflow_unicore_integration-0.1.0}/src/airflow_unicore_integration/__init__.py

@@ -3,8 +3,10 @@ def get_provider_info():
         "package-name": "airflow-unicore-integration",
         "name": "Unicore",
         "description": "Apache Airflow Unicore provider containing Operators and hooks.",
-        "connection-types": [{
-            "connection-type": "unicore",
-            "hook-class-name": "airflow_unicore_integration.hooks.unicore_hooks.UnicoreHook",
-        }],
+        "connection-types": [
+            {
+                "connection-type": "unicore",
+                "hook-class-name": "airflow_unicore_integration.hooks.unicore_hooks.UnicoreHook",
+            }
+        ],
     }

airflow_unicore_integration-0.1.0/src/airflow_unicore_integration/executors/run_task_via_supervisor.py

@@ -0,0 +1,85 @@
+"""
+Usage:
+
+python run_task_via_supervisor.py [--json-string <workload string> | --json-file <workload filepath>]
+
+"""
+
+import argparse
+import sys
+
+import structlog
+from airflow.configuration import conf
+from airflow.executors import workloads
+from airflow.sdk.execution_time.supervisor import supervise
+from pydantic import TypeAdapter
+from pydantic_core._pydantic_core import ValidationError
+
+log = structlog.get_logger(logger_name=__name__)
+
+
+def execute_workload_locally(workload: workloads.All):
+    if not isinstance(workload, workloads.ExecuteTask):
+        raise ValueError(f"Executor does not know how to handle {type(workload)}")
+
+    base_url = conf.get("api", "base_url", fallback="/")
+    default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+    server = conf.get("core", "execution_api_server_url", fallback=default_execution_api_server)
+    log.info("Connecting to server:", server=server)
+
+    supervise(
+        # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this.
+        ti=workload.ti,  # type: ignore[arg-type]
+        dag_rel_path=workload.dag_rel_path,
+        bundle_info=workload.bundle_info,
+        token=workload.token,
+        server=server,
+        log_path=workload.log_path,
+        # Include the output of the task to stdout too, so that in process logs can be read from via the
+        # kubeapi as pod logs.
+        subprocess_logs_to_stdout=True,
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Execute a workload in a Containerised executor using the task SDK."
+    )
+
+    # Create a mutually exclusive group to ensure that only one of the flags is set
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--json-path",
+        help="Path to the input JSON file containing the execution workload payload.",
+        type=str,
+    )
+    group.add_argument(
+        "--json-string",
+        help="The JSON string itself containing the execution workload payload.",
+        type=str,
+    )
+    args = parser.parse_args()
+
+    decoder = TypeAdapter[workloads.All](workloads.All)
+
+    if args.json_path:
+        try:
+            with open(args.json_path) as file:
+                input_data = file.read()
+                workload = decoder.validate_json(input_data)
+        except OSError as e:
+            log.error("Failed to read file", error=str(e))
+            sys.exit(1)
+
+    elif args.json_string:
+        try:
+            workload = decoder.validate_json(args.json_string)
+        except ValidationError as e:
+            log.error("Failed to parse input JSON string", error=str(e))
+            sys.exit(1)
+
+    execute_workload_locally(workload)
+
+
+if __name__ == "__main__":
+    main()

airflow_unicore_integration-0.1.0/src/airflow_unicore_integration/executors/unicore_executor.py

@@ -0,0 +1,123 @@
+"""
+to configure for executor:
+- Connection details for unicore: conn_id AIRFLOW__UNICORE_EXECUTOR__UNICORE_CONN_ID | should be defined, can be skipped if every task provides one
+- location (path) of python virtualenv prepared on hpc system | AIRFLOW__UNICORE_EXECUTOR__DEFAULT_ENV | should be defined, can be skipped if every task provides one
+
+tasks should be allowed to overwrite SITE, CREDENTIALS_*, UNICORE_CONN_ID and DEFAULT_ENV - i.e. everything but the database connection - credentials should be given as a uc_credential object via executor_config
+
+"""
+
+import time
+from typing import Any
+from typing import Dict
+from typing import List
+
+import pyunicore.client as uc_client
+from airflow.configuration import conf
+from airflow.executors.base_executor import BaseExecutor
+from airflow.executors.workloads import All
+from airflow.executors.workloads import ExecuteTask
+from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.utils.state import TaskInstanceState
+
+from airflow_unicore_integration.hooks import unicore_hooks
+
+from ..util.job import JobDescriptionGenerator
+from ..util.job import NaiveJobDescriptionGenerator
+
+STATE_MAPPINGS: Dict[uc_client.JobStatus, TaskInstanceState] = {
+    uc_client.JobStatus.UNDEFINED: TaskInstanceState.FAILED,
+    uc_client.JobStatus.READY: TaskInstanceState.QUEUED,
+    uc_client.JobStatus.STAGINGIN: TaskInstanceState.QUEUED,
+    uc_client.JobStatus.QUEUED: TaskInstanceState.QUEUED,
+    uc_client.JobStatus.RUNNING: TaskInstanceState.RUNNING,
+    uc_client.JobStatus.STAGINGOUT: TaskInstanceState.RUNNING,
+    uc_client.JobStatus.SUCCESSFUL: TaskInstanceState.SUCCESS,
+    uc_client.JobStatus.FAILED: TaskInstanceState.FAILED,
+}
+
+
+class UnicoreExecutor(BaseExecutor):
+
+    def start(self):
+        self.active_jobs: Dict[TaskInstanceKey, uc_client.Job] = {}
+        self.uc_conn = unicore_hooks.UnicoreHook().get_conn()
+        # TODO get job description generator class and init params from config
+        self.job_descr_generator: JobDescriptionGenerator = NaiveJobDescriptionGenerator()
+
+    def sync(self) -> None:
+        # iterate through task collection and update task/ job status - delete if needed
+        for task, job in list(self.active_jobs.items()):
+            state = STATE_MAPPINGS[job.status]
+            if state == TaskInstanceState.FAILED:
+                self.fail(task)
+                self._forward_unicore_log(task, job)
+                self.active_jobs.pop(task)
+            elif state == TaskInstanceState.SUCCESS:
+                self.success(task)
+                self._forward_unicore_log(task, job)
+                self.active_jobs.pop(task)
+            elif state == TaskInstanceState.RUNNING:
+                self.running_state(task, state)
+
+        return super().sync()
+
+    def _forward_unicore_log(self, task: TaskInstanceKey, job: uc_client.Job) -> List[str]:
+        # TODO retrieve unicore logs from job directory and return
+        return []
+
+    def _get_unicore_client(self, executor_config: dict | None = {}):
+        # TODO fix this only temporary solution
+        return self.uc_conn
+        # END TODO fix this
+        # include client desires from executor_config
+        unicore_conn_id = executor_config.get(  # type: ignore
+            UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_CONN_KEY,
+            conf.get("unicore.executor", "UNICORE_CONN_ID"),
+        )  # task can provide a different unicore connection to use, else airflow-wide default is used
+        self.log.info(f"Using base unicore connection with id '{unicore_conn_id}'")
+        hook = unicore_hooks.UnicoreHook(uc_conn_id=unicore_conn_id)
+        unicore_site = executor_config.get(  # type: ignore
+            UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_SITE_KEY, None
+        )  # task can provide a different site to run at, else default from connetion is used
+        unicore_credential = executor_config.get(  # type: ignore
+            UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_CREDENTIAL_KEY, None
+        )  # task can provide a different credential to use, else default from connection is used
+        return hook.get_conn(
+            overwrite_base_url=unicore_site, overwrite_credential=unicore_credential
+        )
+
+    def _submit_job(self, workload: ExecuteTask):
+        uc_client = self._get_unicore_client(executor_config=workload.ti.executor_config)
+        job_descr = self._create_job_description(workload)
+        self.log.info("Generated job description")
+        self.log.debug(str(job_descr))
+        job = uc_client.new_job(job_descr)
+        self.log.info("Submitted unicore job")
+        self.active_jobs[workload.ti.key] = job
+        return job
+
+    def _create_job_description(self, workload: ExecuteTask) -> Dict[str, Any]:
+        return self.job_descr_generator.create_job_description(workload)
+
+    def queue_workload(self, workload: ExecuteTask | All, session):
+        if not isinstance(workload, ExecuteTask):
+            raise TypeError(f"Don't know how to queue workload of type {type(workload).__name__}")
+
+        # submit job to unicore and add to active_jobs dict for task state management
+        job = self._submit_job(workload)
+        self.active_jobs[workload.ti.key] = job
+
+    def end(self, heartbeat_interval=10) -> None:
+        # wait for current jobs to finish, dont start any new ones
+        while True:
+            self.sync()
+            if not self.active_jobs:
+                break
+            time.sleep(heartbeat_interval)
+
+    def terminate(self):
+        # terminate all jobs
+        for task, job in list(self.active_jobs.items()):
+            job.abort()
+        self.end()

airflow_unicore_integration-0.1.0/src/airflow_unicore_integration/hooks/unicore_hooks.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from airflow.hooks.base import BaseHook
+from pyunicore import client
+from pyunicore import credentials
+
+
+class UnicoreHook(BaseHook):
+    """
+    Interact with Unicore.
+
+    Creates Unicore Clients from airflow connections.
+
+    :param uc_conn_id: The unicore connection id - default: uc_default
+    """
+
+    conn_name_attr = "uc_conn_id"
+    default_conn_name = "uc_default"
+    conn_type = "unicore"
+    hook_name = "Unicore"
+
+    def __init__(self, uc_conn_id: str = default_conn_name) -> None:
+        super().__init__()
+        self.uc_conn_id = uc_conn_id
+
+    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(
+            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)
+        if overwrite_base_url is not None:
+            base_url = overwrite_base_url
+        if overwrite_credential is not None:
+            credential = overwrite_credential
+        conn = client.Client(credential, base_url)
+        return conn
+
+    def test_connection(self) -> tuple[bool, str]:
+        """Test the connection by sending an access_info request"""
+        conn = self.get_conn()
+        conn.access_info()
+        return True, "Connection successfully tested"