potiuk commented on code in PR #44311:
URL: https://github.com/apache/airflow/pull/44311#discussion_r1863140184
##########
providers/src/airflow/providers/edge/worker_api/datamodels.py:
##########
@@ -16,17 +16,54 @@
# under the License.
from __future__ import annotations
-from typing import Any, Optional
+from typing import Any, Dict, List, Optional, Union # noqa: UP035 - prevent
pytest failing in back-compat
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+from airflow.providers.edge.models.edge_worker import EdgeWorkerState # noqa:
TCH001
-class JsonRpcRequest(BaseModel):
+
+class JsonRpcRequestBase(BaseModel):
+ """Base JSON RPC request model to define just the method."""
+
+ method: str = Field(description="Fully qualified python module method name
that is called via JSON RPC.")
+
+
+class JsonRpcRequest(JsonRpcRequestBase):
"""JSON RPC request model."""
- method: str
- """Fully qualified python module method name that is called via JSON
RPC."""
- jsonrpc: str
- """JSON RPC version."""
- params: Optional[dict[str, Any]] = None # noqa: UP007 - prevent pytest
failing in back-compat
- """Parameters passed to the method."""
+ jsonrpc: str = Field(description="JSON RPC Version", examples=["2.0"])
+ params: Optional[Dict[str, Any]] = Field( # noqa: UP006, UP007 - prevent
pytest failing in back-compat
+ description="Dictionary of parameters passed to the method."
+ )
+
+
+class WorkerStateBody(BaseModel):
+ """Details of the worker state sent to the scheduler."""
+
+ state: EdgeWorkerState = Field(description="State of the worker from the
view of the worker.")
+ jobs_active: int = Field(0, description="Number of active jobs the worker
is running.")
+ queues: Optional[List[str]] = Field( # noqa: UP006, UP007 - prevent
pytest failing in back-compat
+ description="List of queues the worker is pulling jobs from. If not
provided, worker pulls from all queues."
+ )
+ sysinfo: Dict[str, Union[str, int]] = Field( # noqa: UP006, UP007 -
prevent pytest failing in back-compat
+ description="System information of the worker.",
+ examples=[
+ {
+ "concurrency": 4,
+ "airflow_version": "2.0.0",
+ "edge_provider_version": "1.0.0",
+ }
+ ],
+ )
Review Comment:
Yep. Annotated is cool.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
