apache-airflow-providers-edge3 1.0.0__py3-none-any.whl

airflow/providers/edge3/worker_api/routes/worker.py

@@ -0,0 +1,224 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import json
+from typing import Annotated
+
+from sqlalchemy import select
+
+from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, set_metrics
+from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
+from airflow.providers.edge3.worker_api.datamodels import (
+    WorkerQueueUpdateBody,
+    WorkerRegistrationReturn,
+    WorkerSetStateReturn,
+    WorkerStateBody,
+)
+from airflow.providers.edge3.worker_api.routes._v2_compat import (
+    AirflowRouter,
+    Body,
+    Depends,
+    HTTPException,
+    Path,
+    SessionDep,
+    create_openapi_http_exception_doc,
+    status,
+)
+from airflow.stats import Stats
+from airflow.utils import timezone
+
+worker_router = AirflowRouter(
+    tags=["Worker"],
+    prefix="/worker",
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_400_BAD_REQUEST,
+            status.HTTP_403_FORBIDDEN,
+        ]
+    ),
+)
+
+
+def _assert_version(sysinfo: dict[str, str | int]) -> None:
+    """Check if the Edge Worker version matches the central API site."""
+    from airflow import __version__ as airflow_version
+    from airflow.providers.edge3 import __version__ as edge_provider_version
+
+    # Note: In future, more stable versions we might be more liberate, for the
+    #       moment we require exact version match for Edge Worker and core version
+    if "airflow_version" in sysinfo:
+        airflow_on_worker = sysinfo["airflow_version"]
+        if airflow_on_worker != airflow_version:
+            raise HTTPException(
+                status.HTTP_400_BAD_REQUEST,
+                f"Edge Worker runs on Airflow {airflow_on_worker} "
+                f"and the core runs on {airflow_version}. Rejecting access due to difference.",
+            )
+    else:
+        raise HTTPException(
+            status.HTTP_400_BAD_REQUEST, "Edge Worker does not specify the version it is running on."
+        )
+
+    if "edge_provider_version" in sysinfo:
+        provider_on_worker = sysinfo["edge_provider_version"]
+        if provider_on_worker != edge_provider_version:
+            raise HTTPException(
+                status.HTTP_400_BAD_REQUEST,
+                f"Edge Worker runs on Edge Provider {provider_on_worker} "
+                f"and the core runs on {edge_provider_version}. Rejecting access due to difference.",
+            )
+    else:
+        raise HTTPException(
+            status.HTTP_400_BAD_REQUEST, "Edge Worker does not specify the provider version it is running on."
+        )
+
+
+_worker_name_doc = Path(title="Worker Name", description="Hostname or instance name of the worker")
+_worker_state_doc = Body(
+    title="Worker State",
+    description="State of the worker with details",
+    examples=[
+        {
+            "state": "running",
+            "jobs_active": 3,
+            "queues": ["large_node", "wisconsin_site"],
+            "sysinfo": {
+                "concurrency": 4,
+                "airflow_version": "2.10.0",
+                "edge_provider_version": "1.0.0",
+            },
+        }
+    ],
+)
+_worker_queue_doc = Body(
+    title="Changes in worker queues",
+    description="Changes to be applied to current queues of worker",
+    examples=[{"new_queues": ["new_queue"], "remove_queues": ["old_queue"]}],
+)
+
+
+def redefine_state(worker_state: EdgeWorkerState, body_state: EdgeWorkerState) -> EdgeWorkerState:
+    """Redefine the state of the worker based on maintenance request."""
+    if (
+        worker_state == EdgeWorkerState.MAINTENANCE_REQUEST
+        and body_state
+        not in (
+            EdgeWorkerState.MAINTENANCE_PENDING,
+            EdgeWorkerState.MAINTENANCE_MODE,
+        )
+        or worker_state == EdgeWorkerState.OFFLINE_MAINTENANCE
+        and body_state == EdgeWorkerState.STARTING
+    ):
+        return EdgeWorkerState.MAINTENANCE_REQUEST
+
+    if worker_state == EdgeWorkerState.MAINTENANCE_EXIT:
+        if body_state == EdgeWorkerState.MAINTENANCE_PENDING:
+            return EdgeWorkerState.RUNNING
+        if body_state == EdgeWorkerState.MAINTENANCE_MODE:
+            return EdgeWorkerState.IDLE
+
+    return body_state
+
+
+def redefine_maintenance_comments(
+    worker_maintenance_comment: str | None, body_maintenance_comments: str | None
+) -> str | None:
+    """Add new maintenance comments or overwrite the old ones if it is too long."""
+    if body_maintenance_comments:
+        if (
+            worker_maintenance_comment
+            and len(body_maintenance_comments) + len(worker_maintenance_comment) < 1020
+        ):
+            return f"{worker_maintenance_comment}\n\n{body_maintenance_comments}"
+        return body_maintenance_comments
+    return worker_maintenance_comment
+
+
+@worker_router.post("/{worker_name}", dependencies=[Depends(jwt_token_authorization_rest)])
+def register(
+    worker_name: Annotated[str, _worker_name_doc],
+    body: Annotated[WorkerStateBody, _worker_state_doc],
+    session: SessionDep,
+) -> WorkerRegistrationReturn:
+    """Register a new worker to the backend."""
+    _assert_version(body.sysinfo)
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if not worker:
+        worker = EdgeWorkerModel(worker_name=worker_name, state=body.state, queues=body.queues)
+    worker.state = redefine_state(worker.state, body.state)
+    worker.maintenance_comment = redefine_maintenance_comments(
+        worker.maintenance_comment, body.maintenance_comments
+    )
+    worker.queues = body.queues
+    worker.sysinfo = json.dumps(body.sysinfo)
+    worker.last_update = timezone.utcnow()
+    session.add(worker)
+    return WorkerRegistrationReturn(last_update=worker.last_update)
+
+
+@worker_router.patch("/{worker_name}", dependencies=[Depends(jwt_token_authorization_rest)])
+def set_state(
+    worker_name: Annotated[str, _worker_name_doc],
+    body: Annotated[WorkerStateBody, _worker_state_doc],
+    session: SessionDep,
+) -> WorkerSetStateReturn:
+    """Set state of worker and returns the current assigned queues."""
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    worker.state = redefine_state(worker.state, body.state)
+    worker.maintenance_comment = redefine_maintenance_comments(
+        worker.maintenance_comment, body.maintenance_comments
+    )
+    worker.jobs_active = body.jobs_active
+    worker.sysinfo = json.dumps(body.sysinfo)
+    worker.last_update = timezone.utcnow()
+    session.commit()
+    Stats.incr(f"edge_worker.heartbeat_count.{worker_name}", 1, 1)
+    Stats.incr("edge_worker.heartbeat_count", 1, 1, tags={"worker_name": worker_name})
+    set_metrics(
+        worker_name=worker_name,
+        state=body.state,
+        jobs_active=body.jobs_active,
+        concurrency=int(body.sysinfo.get("concurrency", -1)),
+        free_concurrency=int(body.sysinfo["free_concurrency"]),
+        queues=worker.queues,
+    )
+    _assert_version(body.sysinfo)  #  Exception only after worker state is in the DB
+    return WorkerSetStateReturn(
+        state=worker.state, queues=worker.queues, maintenance_comments=worker.maintenance_comment
+    )
+
+
+@worker_router.patch(
+    "/queues/{worker_name}",
+    dependencies=[Depends(jwt_token_authorization_rest)],
+)
+def update_queues(
+    worker_name: Annotated[str, _worker_name_doc],
+    body: Annotated[WorkerQueueUpdateBody, _worker_queue_doc],
+    session: SessionDep,
+) -> None:
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if body.new_queues:
+        worker.add_queues(body.new_queues)
+    if body.remove_queues:
+        worker.remove_queues(body.remove_queues)
+    session.add(worker)

apache_airflow_providers_edge3-1.0.0.dist-info/METADATA

@@ -0,0 +1,117 @@
+Metadata-Version: 2.4
+Name: apache-airflow-providers-edge3
+Version: 1.0.0
+Summary: Provider package apache-airflow-providers-edge3 for Apache Airflow
+Keywords: airflow-provider,edge3,airflow,integration
+Author-email: Apache Software Foundation <dev@airflow.apache.org>
+Maintainer-email: Apache Software Foundation <dev@airflow.apache.org>
+Requires-Python: ~=3.9
+Description-Content-Type: text/x-rst
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Environment :: Console
+Classifier: Environment :: Web Environment
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: System Administrators
+Classifier: Framework :: Apache Airflow
+Classifier: Framework :: Apache Airflow :: Provider
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: System :: Monitoring
+Requires-Dist: apache-airflow>=2.10.0
+Requires-Dist: pydantic>=2.11.0
+Requires-Dist: retryhttp>=1.2.0,!=1.3.0
+Requires-Dist: apache-airflow-providers-fab ; extra == "fab"
+Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0
+Project-URL: Mastodon, https://fosstodon.org/@airflow
+Project-URL: Slack Chat, https://s.apache.org/airflow-slack
+Project-URL: Source Code, https://github.com/apache/airflow
+Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
+Provides-Extra: fab
+
+
+.. Licensed to the Apache Software Foundation (ASF) under one
+   or more contributor license agreements.  See the NOTICE file
+   distributed with this work for additional information
+   regarding copyright ownership.  The ASF licenses this file
+   to you under the Apache License, Version 2.0 (the
+   "License"); you may not use this file except in compliance
+   with the License.  You may obtain a copy of the License at
+
+..   http://www.apache.org/licenses/LICENSE-2.0
+
+.. Unless required by applicable law or agreed to in writing,
+   software distributed under the License is distributed on an
+   "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+   KIND, either express or implied.  See the License for the
+   specific language governing permissions and limitations
+   under the License.
+
+.. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN!
+
+.. IF YOU WANT TO MODIFY TEMPLATE FOR THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
+   ``PROVIDER_README_TEMPLATE.rst.jinja2`` IN the ``dev/breeze/src/airflow_breeze/templates`` DIRECTORY
+
+Package ``apache-airflow-providers-edge3``
+
+Release: ``1.0.0``
+
+
+Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites
+
+
+Provider package
+----------------
+
+This is a provider package for ``edge3`` provider. All classes for this provider package
+are in ``airflow.providers.edge3`` python package.
+
+You can find package information and changelog for the provider
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0/>`_.
+
+Installation
+------------
+
+You can install this package on top of an existing Airflow 2 installation (see ``Requirements`` below
+for the minimum Airflow version supported) via
+``pip install apache-airflow-providers-edge3``
+
+The package supports the following python versions: 3.9,3.10,3.11,3.12
+
+Requirements
+------------
+
+==================  ===================
+PIP package         Version required
+==================  ===================
+``apache-airflow``  ``>=2.10.0``
+``pydantic``        ``>=2.11.0``
+``retryhttp``       ``>=1.2.0,!=1.3.0``
+==================  ===================
+
+Cross provider package dependencies
+-----------------------------------
+
+Those are dependencies that might be needed in order to use all the features of the package.
+You need to install the specified providers in order to use them.
+
+You can install such cross-provider dependencies when installing from PyPI. For example:
+
+.. code-block:: bash
+
+    pip install apache-airflow-providers-edge3[fab]
+
+
+==============================================================================================  =======
+Dependent package                                                                               Extra
+==============================================================================================  =======
+`apache-airflow-providers-fab <https://airflow.apache.org/docs/apache-airflow-providers-fab>`_  ``fab``
+==============================================================================================  =======
+
+The changelog for the provider package can be found in the
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0/changelog.html>`_.
+

apache_airflow_providers_edge3-1.0.0.dist-info/RECORD

@@ -0,0 +1,39 @@
+airflow/providers/edge3/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
+airflow/providers/edge3/__init__.py,sha256=hb8OUkltNOLHQS7pnEmS9LOc9NJyajbwwZY_51qn6kk,1494
+airflow/providers/edge3/get_provider_info.py,sha256=T8CMBInzQT7TRJFEFf3tbDK6otJ210gXj4qOaOLyfnE,5548
+airflow/providers/edge3/version_compat.py,sha256=aHg90_DtgoSnQvILFICexMyNlHlALBdaeWqkX3dFDug,1605
+airflow/providers/edge3/cli/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/edge3/cli/api_client.py,sha256=J5l99mwwKXC7Ub_kO6fe_56A8Ue3ag8mIYxD_qPLfuQ,7497
+airflow/providers/edge3/cli/dataclasses.py,sha256=JUuvvmzSVWvG9uOEfzLIiXrTZ-HbESvu50jkPpVIYVw,2895
+airflow/providers/edge3/cli/edge_command.py,sha256=xWNcgvUJyHqT51_Q-6bVPjCd1lrI7KpxSrFfejPvgjA,27466
+airflow/providers/edge3/example_dags/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/edge3/example_dags/integration_test.py,sha256=pXf6rmS4crHKfCBT73LvPmhJS8Rmri8pUlfIvuLja7E,5240
+airflow/providers/edge3/example_dags/win_notepad.py,sha256=2evbqiupi5Ko4tyRkIEC5TPbc2c0wyR2YpsDYsBLMMM,2828
+airflow/providers/edge3/example_dags/win_test.py,sha256=GegWqjvbsSdbsA_f3S9_FRYftVO0pggXwQQggB9Vvz4,13220
+airflow/providers/edge3/executors/__init__.py,sha256=y830gGSKCvjOcLwLuCDp84NCrHWWB9RSSH1qvJpFhyY,923
+airflow/providers/edge3/executors/edge_executor.py,sha256=d05OpRKd82mSkKeybNZPXw1yiRdZAKYLZGh15eIgaiM,15254
+airflow/providers/edge3/models/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/edge3/models/edge_job.py,sha256=rdl9cH1bBrc7id8zkZ7uxsCJNPLG-8o9cnspGwBPBcQ,3167
+airflow/providers/edge3/models/edge_logs.py,sha256=bNstp7gR54O2vbxzz4NTL0erbifFbGUjZ-YOM0I4sqk,2768
+airflow/providers/edge3/models/edge_worker.py,sha256=MsWzvJRTdvCRLZAMDy-nYYkMOl6Ji6AUeNfSQOeO6do,8587
+airflow/providers/edge3/openapi/__init__.py,sha256=0O-WvmDx8GeKSoECpHYrbe0hW-LgjlKny3VqTCpBQeQ,927
+airflow/providers/edge3/openapi/edge_worker_api_v1.yaml,sha256=GAE2IdOXmcUueNy5KFkLBgNpoWnOjnHT9TrW5NZEWpI,24938
+airflow/providers/edge3/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/edge3/plugins/edge_executor_plugin.py,sha256=wvveKcIUP9QkdTFIQmxq4GS15M8KnbHkDfKjlpvdvwk,9051
+airflow/providers/edge3/plugins/templates/edge_worker_hosts.html,sha256=0_P2yfZwpy3Kvqd3GBvu_PgmmKCUbso3ieW8aYa76iU,8997
+airflow/providers/edge3/plugins/templates/edge_worker_jobs.html,sha256=bZ-6ysmIy6j4eR_TPHiqbgb3qpNMKCcEEB-SpxuxNgc,2831
+airflow/providers/edge3/worker_api/__init__.py,sha256=nnPvxWGTEKZ9YyB1Yd7P9IvDOenK01LVHm22Owwxj3g,839
+airflow/providers/edge3/worker_api/app.py,sha256=Dda2VjkzgBtbQbSWSVEAoqd22RlqvBMyiPau65uKkv4,2006
+airflow/providers/edge3/worker_api/auth.py,sha256=4UFKc6wQPt9AqoMXMxt1rXljEy8xofgaIhzIdLR54gc,4936
+airflow/providers/edge3/worker_api/datamodels.py,sha256=FAiXqnrSN8zH4YE2fUMjXfXcH9cHlhRh4uZvvr936Ys,6696
+airflow/providers/edge3/worker_api/routes/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/edge3/worker_api/routes/_v2_compat.py,sha256=Q4b2Io0yoK5V_hbgk6fiFviTeT6CFbFMOGYRZgLEeR4,4543
+airflow/providers/edge3/worker_api/routes/_v2_routes.py,sha256=-WAofvXJpOYpTyh983cxMarA5FkPlH7LMhZuSHE-qPg,10830
+airflow/providers/edge3/worker_api/routes/health.py,sha256=XxqIppnRA138Q6mAHCdyL2JvoeeganUiI-TXyXSPTGo,1075
+airflow/providers/edge3/worker_api/routes/jobs.py,sha256=UK1w6nXEUadOLwE9abZ4jHH4KtbvXcwaAF0EnwSa3y4,5733
+airflow/providers/edge3/worker_api/routes/logs.py,sha256=uk0SZ5hAimj3sAcq1FYCDu0AXYNeTeyjZDGBvw-986E,4945
+airflow/providers/edge3/worker_api/routes/worker.py,sha256=jb6m9eE9J1fyejf-DuEtGPzGrynyUmf5qkZxLeN-DQo,8554
+apache_airflow_providers_edge3-1.0.0.dist-info/entry_points.txt,sha256=7WUIGfd3o9NvvbK5trbZxNXTgYGc6pqg74wZPigbx5o,206
+apache_airflow_providers_edge3-1.0.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_edge3-1.0.0.dist-info/METADATA,sha256=zsqZL7eFQZvX8sSAkU9e99QOvOAGaJXYIpKEenYwScM,4980
+apache_airflow_providers_edge3-1.0.0.dist-info/RECORD,,

apache_airflow_providers_edge3-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: flit 3.12.0
+Root-Is-Purelib: true
+Tag: py3-none-any

apache_airflow_providers_edge3-1.0.0.dist-info/entry_points.txt

@@ -0,0 +1,6 @@
+[airflow.plugins]
+edge_executor=airflow.providers.edge3.plugins.edge_executor_plugin:EdgeExecutorPlugin
+
+[apache_airflow_provider]
+provider_info=airflow.providers.edge3.get_provider_info:get_provider_info
+