apache-airflow-providers-edge3 1.1.2rc1__py3-none-any.whl → 1.1.2rc2__py3-none-any.whl

airflow/providers/edge3/cli/edge_command.py

@@ -351,6 +351,48 @@ def remote_worker_request_shutdown(args) -> None:
     logger.info("Requested shutdown of Edge Worker host %s by %s.", args.edge_hostname, getuser())
 
 
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def add_worker_queues(args) -> None:
+    """Add queues to an edge worker."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import add_worker_queues
+
+    queues = args.queues.split(",") if args.queues else []
+    if not queues:
+        raise SystemExit("Error: No queues specified to add.")
+
+    try:
+        add_worker_queues(args.edge_hostname, queues)
+        logger.info("Added queues %s to Edge Worker host %s by %s.", queues, args.edge_hostname, getuser())
+    except TypeError as e:
+        logger.error(str(e))
+        raise SystemExit
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def remove_worker_queues(args) -> None:
+    """Remove queues from an edge worker."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import remove_worker_queues
+
+    queues = args.queues.split(",") if args.queues else []
+    if not queues:
+        raise SystemExit("Error: No queues specified to remove.")
+
+    try:
+        remove_worker_queues(args.edge_hostname, queues)
+        logger.info(
+            "Removed queues %s from Edge Worker host %s by %s.", queues, args.edge_hostname, getuser()
+        )
+    except TypeError as e:
+        logger.error(str(e))
+        raise SystemExit
+
+
 ARG_CONCURRENCY = Arg(
     ("-c", "--concurrency"),
     type=int,
@@ -380,6 +422,11 @@ ARG_REQUIRED_MAINTENANCE_COMMENT = Arg(
     help="Maintenance comments to report reason. Required if enabling maintenance",
     required=True,
 )
+ARG_QUEUES_MANAGE = Arg(
+    ("-q", "--queues"),
+    help="Comma delimited list of queues to add or remove.",
+    required=True,
+)
 ARG_WAIT_MAINT = Arg(
     ("-w", "--wait"),
     default=False,
@@ -516,4 +563,22 @@ EDGE_COMMANDS: list[ActionCommand] = [
         func=remote_worker_request_shutdown,
         args=(ARG_REQUIRED_EDGE_HOSTNAME,),
     ),
+    ActionCommand(
+        name="add-worker-queues",
+        help=add_worker_queues.__doc__,
+        func=add_worker_queues,
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_QUEUES_MANAGE,
+        ),
+    ),
+    ActionCommand(
+        name="remove-worker-queues",
+        help=remove_worker_queues.__doc__,
+        func=remove_worker_queues,
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_QUEUES_MANAGE,
+        ),
+    ),
 ]

airflow/providers/edge3/cli/worker.py

@@ -188,8 +188,8 @@ class EdgeWorker:
 
             try:
                 api_url = conf.get("edge", "api_url")
-                execution_api_server_url = conf.get("core", "execution_api_server_url", fallback=...)
-                if execution_api_server_url is ...:
+                execution_api_server_url = conf.get("core", "execution_api_server_url", fallback="")
+                if not execution_api_server_url:
                     parsed = urlparse(api_url)
                     execution_api_server_url = f"{parsed.scheme}://{parsed.netloc}/execution/"
 

airflow/providers/edge3/example_dags/integration_test.py

@@ -41,13 +41,13 @@ try:
     from airflow.sdk import DAG, Param, Variable, task, task_group
 except ImportError:
     # Airflow 2.10 compat
-    from airflow.decorators import task, task_group  # type: ignore[no-redef,attr-defined]
-    from airflow.models.dag import DAG  # type: ignore[no-redef,attr-defined,assignment]
-    from airflow.models.param import Param  # type: ignore[no-redef,attr-defined]
-    from airflow.models.variable import Variable  # type: ignore[no-redef,attr-defined]
-    from airflow.operators.bash import BashOperator  # type: ignore[no-redef,attr-defined]
-    from airflow.operators.empty import EmptyOperator  # type: ignore[no-redef,attr-defined]
-    from airflow.operators.python import PythonOperator  # type: ignore[no-redef,attr-defined]
+    from airflow.decorators import task, task_group  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[assignment]
+    from airflow.models.param import Param  # type: ignore[no-redef]
+    from airflow.models.variable import Variable
+    from airflow.operators.bash import BashOperator  # type: ignore[no-redef]
+    from airflow.operators.empty import EmptyOperator  # type: ignore[no-redef]
+    from airflow.operators.python import PythonOperator  # type: ignore[no-redef]
 
 with DAG(
     dag_id="integration_test",

airflow/providers/edge3/example_dags/win_test.py

@@ -32,7 +32,11 @@ from subprocess import STDOUT, Popen
 from time import sleep
 from typing import TYPE_CHECKING, Any
 
-from airflow.decorators import task, task_group
+try:
+    from airflow.sdk import task, task_group
+except ImportError:
+    # Airflow 2 path
+    from airflow.decorators import task, task_group  # type: ignore[attr-defined,no-redef]
 from airflow.exceptions import AirflowException, AirflowNotFoundException, AirflowSkipException
 from airflow.models import BaseOperator
 from airflow.models.dag import DAG
@@ -52,11 +56,11 @@ if TYPE_CHECKING:
     try:
         from airflow.sdk.types import RuntimeTaskInstanceProtocol as TaskInstance
     except ImportError:
-        from airflow.models import TaskInstance  # type: ignore[assignment, no-redef]
+        from airflow.models import TaskInstance  # type: ignore[assignment]
     from airflow.utils.context import Context
 
 try:
-    from airflow.operators.python import PythonOperator  # type: ignore
+    from airflow.operators.python import PythonOperator
 except ImportError:
     from airflow.providers.common.compat.standard.operators import PythonOperator
 
@@ -277,7 +281,7 @@ with DAG(
 
         @task.virtualenv(requirements="numpy")
         def virtualenv():
-            import numpy  # type: ignore
+            import numpy
 
             print(f"Welcome to virtualenv with numpy version {numpy.__version__}.")
 

airflow/providers/edge3/executors/edge_executor.py

@@ -30,7 +30,6 @@ from sqlalchemy.orm import Session
 from airflow.cli.cli_config import GroupCommand
 from airflow.configuration import conf
 from airflow.executors.base_executor import BaseExecutor
-from airflow.models.abstractoperator import DEFAULT_QUEUE
 from airflow.models.taskinstance import TaskInstance, TaskInstanceState
 from airflow.providers.edge3.cli.edge_command import EDGE_COMMANDS
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
@@ -55,6 +54,7 @@ if TYPE_CHECKING:
     TaskTuple = tuple[TaskInstanceKey, CommandType, str | None, Any | None]
 
 PARALLELISM: int = conf.getint("core", "PARALLELISM")
+DEFAULT_QUEUE: str = conf.get_mandatory_value("operators", "default_queue")
 
 
 class EdgeExecutor(BaseExecutor):

airflow/providers/edge3/models/edge_worker.py

@@ -109,7 +109,7 @@ class EdgeWorkerModel(Base, LoggingMixin):
         super().__init__()
 
     @property
-    def sysinfo_json(self) -> dict:
+    def sysinfo_json(self) -> dict | None:
         return json.loads(self.sysinfo) if self.sysinfo else None
 
     @property
@@ -283,3 +283,37 @@ def request_shutdown(worker_name: str, session: Session = NEW_SESSION) -> None:
         EdgeWorkerState.UNKNOWN,
     ):
         worker.state = EdgeWorkerState.SHUTDOWN_REQUEST
+
+
+@provide_session
+def add_worker_queues(worker_name: str, queues: list[str], session: Session = NEW_SESSION) -> None:
+    """Add queues to an edge worker."""
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if worker.state in (
+        EdgeWorkerState.OFFLINE,
+        EdgeWorkerState.OFFLINE_MAINTENANCE,
+        EdgeWorkerState.UNKNOWN,
+    ):
+        error_message = f"Cannot add queues to edge worker {worker_name} as it is in {worker.state} state!"
+        logger.error(error_message)
+        raise TypeError(error_message)
+    worker.add_queues(queues)
+
+
+@provide_session
+def remove_worker_queues(worker_name: str, queues: list[str], session: Session = NEW_SESSION) -> None:
+    """Remove queues from an edge worker."""
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if worker.state in (
+        EdgeWorkerState.OFFLINE,
+        EdgeWorkerState.OFFLINE_MAINTENANCE,
+        EdgeWorkerState.UNKNOWN,
+    ):
+        error_message = (
+            f"Cannot remove queues from edge worker {worker_name} as it is in {worker.state} state!"
+        )
+        logger.error(error_message)
+        raise TypeError(error_message)
+    worker.remove_queues(queues)

airflow/providers/edge3/plugins/edge_executor_plugin.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import sys
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
@@ -213,12 +214,24 @@ try:
 except AirflowConfigException:
     EDGE_EXECUTOR_ACTIVE = False
 
+# Load the API endpoint only on api-server (Airflow 3.x) or webserver (Airflow 2.x)
+# todo(jscheffl): Remove this check when the discussion in
+#                 https://lists.apache.org/thread/w170czq6r7bslkqp1tk6bjjjo0789wgl
+#                 resulted in a proper API to selective initialize. Maybe backcompat-shim
+#                 is also needed to support Airflow-versions prior the rework.
+if AIRFLOW_V_3_0_PLUS:
+    RUNNING_ON_APISERVER = (len(sys.argv) > 1 and sys.argv[1] in ["api-server"]) or (
+        len(sys.argv) > 2 and sys.argv[2] == "airflow-core/src/airflow/api_fastapi/main.py"
+    )
+else:
+    RUNNING_ON_APISERVER = "gunicorn" in sys.argv[0] and "airflow-webserver" in sys.argv
+
 
 class EdgeExecutorPlugin(AirflowPlugin):
     """EdgeExecutor Plugin - provides API endpoints for Edge Workers in Webserver."""
 
     name = "edge_executor"
-    if EDGE_EXECUTOR_ACTIVE:
+    if EDGE_EXECUTOR_ACTIVE and RUNNING_ON_APISERVER:
         if AIRFLOW_V_3_0_PLUS:
             fastapi_apps = [_get_api_endpoint()]
         else:

airflow/providers/edge3/worker_api/auth.py

@@ -58,7 +58,7 @@ if AIRFLOW_V_3_0_PLUS:
 
 else:
     # Airflow 2.10 compatibility
-    from airflow.utils.jwt_signer import JWTSigner  # type: ignore
+    from airflow.utils.jwt_signer import JWTSigner
 
     @cache
     def jwt_signer() -> JWTSigner:

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

@@ -127,7 +127,7 @@ else:
 
     # In Airflow 3 with AIP-72 we get workload addressed by ExecuteTask
     # But in Airflow 2.10 it is a command line array
-    ExecuteTask = list[str]  # type: ignore[no-redef,assignment,misc]
+    ExecuteTask = list[str]  # type: ignore[assignment,misc]
 
     def parse_command(command: str) -> ExecuteTask:
         from ast import literal_eval

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

@@ -66,7 +66,7 @@ def rpcapi_v2(body: dict[str, Any]) -> APIResponse:
     # Note: Except the method map this _was_ a 100% copy of internal API module
     #       airflow.api_internal.endpoints.rpc_api_endpoint.internal_airflow_api()
     # As of rework for FastAPI in Airflow 3.0, this is updated and to be removed in the future.
-    from airflow.api_internal.endpoints.rpc_api_endpoint import (  # type: ignore[attr-defined]
+    from airflow.api_internal.endpoints.rpc_api_endpoint import (
         # Note: This is just for compatibility with Airflow 2.10, not working for Airflow 3 / main as removed
         initialize_method_map,
     )

{apache_airflow_providers_edge3-1.1.2rc1.dist-info → apache_airflow_providers_edge3-1.1.2rc2.dist-info}/METADATA

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-edge3
-Version: 1.1.2rc1
+Version: 1.1.2rc2
 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.10
+Requires-Python: >=3.10
 Description-Content-Type: text/x-rst
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: Console
@@ -18,6 +18,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: apache-airflow>=2.10.0rc1
 Requires-Dist: pydantic>=2.11.0
@@ -57,6 +58,7 @@ Package ``apache-airflow-providers-edge3``
 
 Release: ``1.1.2``
 
+Release Date: ``|PypiReleaseDate|``
 
 Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites.
 
@@ -89,7 +91,7 @@ You can install this package on top of an existing Airflow 2 installation (see `
 for the minimum Airflow version supported) via
 ``pip install apache-airflow-providers-edge3``
 
-The package supports the following python versions: 3.10,3.11,3.12
+The package supports the following python versions: 3.10,3.11,3.12,3.13
 
 Requirements
 ------------

{apache_airflow_providers_edge3-1.1.2rc1.dist-info → apache_airflow_providers_edge3-1.1.2rc2.dist-info}/RECORD

@@ -5,37 +5,37 @@ airflow/providers/edge3/version_compat.py,sha256=j5PCtXvZ71aBjixu-EFTNtVDPsngzzs
 airflow/providers/edge3/cli/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/edge3/cli/api_client.py,sha256=334KHVB4eMSzRpQ5emS56o-RTUJQprxf5Q3xQldCHDQ,7440
 airflow/providers/edge3/cli/dataclasses.py,sha256=JUuvvmzSVWvG9uOEfzLIiXrTZ-HbESvu50jkPpVIYVw,2895
-airflow/providers/edge3/cli/edge_command.py,sha256=v5HFkPnSYQO2Iggvr2pqfZrkoU3WwZO3-m0jntZA2Jg,18262
+airflow/providers/edge3/cli/edge_command.py,sha256=qfHjVxA8QJ6zD8Eb8Ipieso4wo51_3MBaTtuSSIEOI4,20339
 airflow/providers/edge3/cli/signalling.py,sha256=sf4S6j6OoP0bLkda3UlCmlZabjv5wsMypy3kAvx56Z0,3220
-airflow/providers/edge3/cli/worker.py,sha256=nsHjMgEFd71ABSSwX51NMag4I2ZDSrXd2bIv9UjFBdA,17263
+airflow/providers/edge3/cli/worker.py,sha256=x1SX5J2mFgfUTERKv5m8lDxNnoss3BV4F7Jgv6u8rPg,17259
 airflow/providers/edge3/example_dags/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/edge3/example_dags/integration_test.py,sha256=Ba55JwDK6j3x-L-C6GWKJ4NW28nJlQtm6I9OXE28vh8,6233
+airflow/providers/edge3/example_dags/integration_test.py,sha256=4iXvv379nFPxF1_gfbqb5NW1z-24MscYWLDxDwTNbPo,6120
 airflow/providers/edge3/example_dags/win_notepad.py,sha256=zYcrKqODN4KLZQ-5wNnZQQskrDd5LA-nKJNgKQDntSE,2832
-airflow/providers/edge3/example_dags/win_test.py,sha256=i7P21nG8ND0NAgKHWsTZg6IlxtDlmspxqYdQJr30Z9w,13326
+airflow/providers/edge3/example_dags/win_test.py,sha256=wJCeij-iPQOnWZlDU7Oqcv6mkE8cVe03rOFiB8We_oQ,13418
 airflow/providers/edge3/executors/__init__.py,sha256=y830gGSKCvjOcLwLuCDp84NCrHWWB9RSSH1qvJpFhyY,923
-airflow/providers/edge3/executors/edge_executor.py,sha256=izmLzGrtSky1FPqpifuy63EoSRfB9UJys-R_AoduJjk,16494
+airflow/providers/edge3/executors/edge_executor.py,sha256=arcq959o4l5x8yjt6f4ck1pDSDFzkGiAh4jGQ_jmqso,16512
 airflow/providers/edge3/models/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/edge3/models/edge_job.py,sha256=3D5HAzcVkyI2bxl3pVbbRxjIz--Tnr_eNFiw2oI6gEQ,3167
 airflow/providers/edge3/models/edge_logs.py,sha256=bNstp7gR54O2vbxzz4NTL0erbifFbGUjZ-YOM0I4sqk,2768
-airflow/providers/edge3/models/edge_worker.py,sha256=4qr5K-QU5yTd2p3AtdVhWlbvToxFFBWk2qIRFpJToWo,10749
+airflow/providers/edge3/models/edge_worker.py,sha256=z3Rru2iFqaFVyvJ6vBQ4-W0T9MenDHip6bJHfLqI6K4,12102
 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=PWQQxvskPGj2EuxpQo_mI6KCgCqvP3gWvY3pa3fdcrA,10258
+airflow/providers/edge3/plugins/edge_executor_plugin.py,sha256=cJcjmk-4OYBISpeGFXGvi78GPZ5EFgymzBm6FWvMMFA,10992
 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=XVTfL-c0JYUhpVkKdqqaxZACZXqvnPp_3W6q2TK0Bjc,4883
+airflow/providers/edge3/worker_api/auth.py,sha256=nmwfUz-nokUKyQp-UKwlMn-i2U5AXzq7c8Xfrt_bGeU,4867
 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=U5VuxoW6sUt1fCWfz-OwbqjUFn2omAewxv9K5Bqzdic,4552
-airflow/providers/edge3/worker_api/routes/_v2_routes.py,sha256=-WAofvXJpOYpTyh983cxMarA5FkPlH7LMhZuSHE-qPg,10830
+airflow/providers/edge3/worker_api/routes/_v2_compat.py,sha256=PuzSL9dMuji_MsluCoPdCvYboqJD0h4zERcIoB5kJvI,4543
+airflow/providers/edge3/worker_api/routes/_v2_routes.py,sha256=xcbf6RdOHx5zOl9JlIAW9nQhiy3ju-EIyq1tbXGJSYc,10800
 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=BGARu1RZ74lW9X-ltuMYbbVXczm_MZdqHaai2MhDWtY,8969
-apache_airflow_providers_edge3-1.1.2rc1.dist-info/entry_points.txt,sha256=7WUIGfd3o9NvvbK5trbZxNXTgYGc6pqg74wZPigbx5o,206
-apache_airflow_providers_edge3-1.1.2rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-apache_airflow_providers_edge3-1.1.2rc1.dist-info/METADATA,sha256=y4M-dHQwRqbExuYM_Ctw-BbQry7x0FC2gS-cPiAZwek,4707
-apache_airflow_providers_edge3-1.1.2rc1.dist-info/RECORD,,
+apache_airflow_providers_edge3-1.1.2rc2.dist-info/entry_points.txt,sha256=7WUIGfd3o9NvvbK5trbZxNXTgYGc6pqg74wZPigbx5o,206
+apache_airflow_providers_edge3-1.1.2rc2.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_edge3-1.1.2rc2.dist-info/METADATA,sha256=khGdhmJivDTg4Jpo-QLB4KGFCN0W62SKy6eUTTn42SU,4799
+apache_airflow_providers_edge3-1.1.2rc2.dist-info/RECORD,,