Adopt new fields in API, change constructor arguments format (#10)

* feat: add new fields for run operator, change payloads to dict

* fix: make runtime field string

airflow_providers_wherobots/hooks/rest_api.py

@@ -15,7 +15,6 @@
 from airflow_providers_wherobots.hooks.base import DEFAULT_CONN_ID
 from airflow_providers_wherobots.wherobots.models import (
  Run,
- CreateRunPayload,
  LogsResponse,
 )
 
@@ -81,11 +80,11 @@ def get_run(self, run_id: str) -> Run:
  resp_json = self._api_call("GET", f"/runs/{run_id}").json()
  return Run.model_validate(resp_json)
 
- def create_run(self, payload: CreateRunPayload) -> Run:
+ def create_run(self, payload: dict[str, Any]) -> Run:
  resp_json = self._api_call(
  "POST",
  "/runs",
- payload=payload.model_dump(mode="json"),
+ payload=payload,
  ).json()
  return Run.model_validate(resp_json)
 

airflow_providers_wherobots/operators/run.py

@@ -8,14 +8,10 @@
 
 from airflow.models import BaseOperator
 from strenum import StrEnum
-from wherobots.db import Runtime
 
 from airflow_providers_wherobots.hooks.base import DEFAULT_CONN_ID
 from airflow_providers_wherobots.hooks.rest_api import WherobotsRestAPIHook
 from airflow_providers_wherobots.wherobots.models import (
- PythonRunPayload,
- JavaRunPayload,
- CreateRunPayload,
  RUN_NAME_ALPHABET,
  RunStatus,
  Run,
@@ -36,9 +32,10 @@ class WherobotsRunOperator(BaseOperator):
  def __init__(
  self,
  name: Optional[str] = None,
- runtime: Optional[Runtime] = Runtime.SEDONA,
- python: Optional[PythonRunPayload] = None,
- java: Optional[JavaRunPayload] = None,
+ runtime: str = "TINY",
+ run_python: Optional[dict[str, Any]] = None,
+ run_jar: Optional[dict[str, Any]] = None,
+ environment: Optional[dict[str, Any]] = None,
  polling_interval: int = 20,
  wherobots_conn_id: str = DEFAULT_CONN_ID,
  poll_logs: bool = False,
@@ -47,12 +44,16 @@ def __init__(
  ):
  super().__init__(**kwargs)
  # If the user specifies the name, we will use it and rely on the server to validate the name
- self.run_payload = CreateRunPayload(
- runtime=runtime,
- name=name or self.default_run_name,
- python=python,
- java=java,
- )
+ self.run_payload: dict[str, Any] = {
+ "runtime": runtime,
+ "name": name or self.default_run_name,
+ }
+ if run_python:
+ self.run_payload["runPython"] = run_python
+ if run_jar:
+ self.run_payload["runJar"] = run_jar
+ if environment:
+ self.run_payload["environment"] = environment
  self._polling_interval = polling_interval
  self.wherobots_conn_id = wherobots_conn_id
  self.xcom_push = xcom_push

airflow_providers_wherobots/wherobots/models.py

@@ -5,11 +5,10 @@
 import string
 from datetime import datetime
 from enum import auto
-from typing import Optional, Sequence, List
+from typing import Optional, List
 
-from pydantic import BaseModel, Field, ConfigDict, computed_field
+from pydantic import BaseModel, Field, ConfigDict
 from strenum import StrEnum
-from wherobots.db import Runtime
 
 RUN_NAME_ALPHABET = string.ascii_letters + string.digits + "-_."
 
@@ -39,81 +38,6 @@ class Run(WherobotsModel):
  end_time: Optional[datetime] = Field(default=None, alias="completeTime")
 
 
-class PythonRunPayload(BaseModel):
- """
- Model for the payload of Run with type == "python"
- """
-
- # For airflow to render the template fields
- template_fields: Sequence[str] = Field(
- ("uri", "args", "entrypoint"), exclude=True, init=False
- )
-
- uri: str
- args: list[str] = []
- entrypoint: Optional[str] = None
-
- @classmethod
- def create(cls, uri: str, args: list[str], entrypoint: Optional[str] = None):
- return cls(uri=uri, args=args, entrypoint=entrypoint)
-
-
-class JavaRunPayload(BaseModel):
- """
- Model for the payload of Run with type == "python"
- """
-
- # For airflow to render the template fields
- template_fields: Sequence[str] = Field(("uri", "args", "main_class"), exclude=True)
-
- uri: str
- args: list[str] = []
- main_class: Optional[str] = Field(None, alias="mainClass")
-
- @classmethod
- def create(cls, uri: str, args: list[str], main_class: Optional[str] = None):
- return cls(uri=uri, args=args, mainClass=main_class)
-
-
-class RunType(StrEnum):
- python = auto()
- java = auto()
-
-
-class CreateRunPayload(BaseModel):
- # For airflow to render the template fields
- template_fields: Sequence[str] = Field(("name", "python", "java"), exclude=True)
-
- runtime: Runtime
- name: Optional[str] = None
- python: Optional[PythonRunPayload] = None
- java: Optional[JavaRunPayload] = None
- timeout_seconds: int = Field(3600, alias="timeoutSeconds")
-
- @computed_field
- def type(self) -> RunType:
- run_type = RunType.python if self.python else RunType.java
- assert isinstance(run_type, RunType)
- return run_type
-
- @classmethod
- def create(
- cls,
- runtime: Runtime,
- name: str,
- python: Optional[PythonRunPayload] = None,
- java: Optional[JavaRunPayload] = None,
- timeout_seconds: int = 3600,
- ):
- return cls(
- runtime=runtime,
- name=name,
- python=python,
- java=java,
- timeoutSeconds=timeout_seconds,
- )
-
-
 class LogItem(BaseModel):
  timestamp: int
  raw: str

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "airflow-providers-wherobots"
-version = "0.1.9"
+version = "0.1.10"
 description = "Airflow extension for communicating with Wherobots Cloud"
 authors = ["zongsi.zhang <[email protected]>"]
 readme = "README.md"

tests/integration_tests/operator/test_run.py

@@ -10,9 +10,6 @@
 from airflow.models import Connection
 
 from airflow_providers_wherobots.operators.run import WherobotsRunOperator
-from airflow_providers_wherobots.wherobots.models import (
- PythonRunPayload,
-)
 from tests.unit_tests.operator.test_run import execute_dag
 
 DEFAULT_START = pendulum.datetime(2021, 9, 13, tz="UTC")
@@ -28,9 +25,9 @@ def test_staging_run_success(staging_conn: Connection, dag: DAG) -> None:
  wherobots_conn_id=staging_conn.conn_id,
  task_id="test_run_smoke",
  name="airflow_operator_test_run_{{ ts_nodash }}",
- python=PythonRunPayload(
- uri="s3://wbts-wbc-rcv7vl73oy/hao9o6y8ci/data/customer-z4asgjn7clrcbz/very_simple_job.py"
- ),
+ run_python={
+ "uri": "s3://wbts-wbc-rcv7vl73oy/hao9o6y8ci/data/customer-z4asgjn7clrcbz/very_simple_job.py"
+ },
  dag=dag,
  )
  execute_dag(dag, task_id=operator.task_id)

tests/unit_tests/hooks/test_rest_api.py

@@ -18,8 +18,6 @@
 )
 from airflow_providers_wherobots.wherobots.models import (
  Run,
- CreateRunPayload,
- PythonRunPayload,
  LogsResponse,
 )
 from tests.unit_tests import helpers
@@ -108,22 +106,20 @@ def test_create_run(self, test_default_conn) -> None:
  """
  test_run: Run = helpers.run_factory.build()
  url = f"https://{test_default_conn.host}/runs"
- create_payload = CreateRunPayload.create(
- name=test_run.name,
- runtime=Runtime.SEDONA,
- python=PythonRunPayload(
- uri="s3://bucket/test.py",
- args=["arg1", "arg2"],
- entrypoint="src.main",
- ),
- )
+ create_payload = {
+ "name": test_run.name,
+ "runtime": Runtime.SEDONA.value,
+ "python": {
+ "uri": "s3://bucket/test.py",
+ "args": ["arg1", "arg2"],
+ "entrypoint": "src.main",
+ },
+ }
  responses.add(
  responses.POST,
  url,
  json=test_run.model_dump(mode="json"),
- match=[
- matchers.json_params_matcher(create_payload.model_dump(mode="json"))
- ],
+ match=[matchers.json_params_matcher(create_payload)],
  status=HTTPStatus.OK,
  )
  with WherobotsRestAPIHook() as hook:

tests/unit_tests/operator/test_run.py

@@ -18,9 +18,7 @@
 
 from airflow_providers_wherobots.operators.run import WherobotsRunOperator
 from airflow_providers_wherobots.wherobots.models import (
- PythonRunPayload,
  RunStatus,
- CreateRunPayload,
  LogsResponse,
  Run,
  LogItem,
@@ -74,22 +72,22 @@ def test_render_template(self, mocker: MockerFixture, dag: DAG):
  operator = WherobotsRunOperator(
  task_id="test_render_template_python",
  name="test_run_{{ ds }}",
- python=PythonRunPayload(
- uri="s3://bucket/test-{{ ds }}.py",
- args=["{{ ds }}"],
- entrypoint="src.main_{{ ds }}",
- ),
+ run_python={
+ "uri": "s3://bucket/test-{{ ds }}.py",
+ "args": ["{{ ds }}"],
+ },
  dag=dag,
  )
  execute_dag(dag, task_id=operator.task_id)
  assert create_run.call_count == 1
  rendered_payload = create_run.call_args.args[0]
- assert isinstance(rendered_payload, CreateRunPayload)
+ assert isinstance(rendered_payload, dict)
  expected_ds = data_interval_start.format("YYYY-MM-DD")
- assert rendered_payload.name == f"test_run_{expected_ds}"
- assert rendered_payload.python.uri == f"s3://bucket/test-{expected_ds}.py"
- assert rendered_payload.python.args == [expected_ds]
- assert rendered_payload.python.entrypoint == f"src.main_{expected_ds}"
+ assert rendered_payload["name"] == f"test_run_{expected_ds}"
+ assert (
+ rendered_payload["runPython"]["uri"] == f"s3://bucket/test-{expected_ds}.py"
+ )
+ assert rendered_payload["runPython"]["args"] == [expected_ds]
 
  @pytest.mark.usefixtures("clean_airflow_db")
  def test_default_name(self, mocker: MockerFixture, dag: DAG):
@@ -100,13 +98,13 @@ def test_default_name(self, mocker: MockerFixture, dag: DAG):
  )
  operator = WherobotsRunOperator(
  task_id="test_default_name",
- python=PythonRunPayload(uri=""),
+ run_python={"uri": ""},
  dag=dag,
  )
  execute_dag(dag, task_id=operator.task_id)
  rendered_payload = create_run.call_args.args[0]
- assert isinstance(rendered_payload, CreateRunPayload)
- assert rendered_payload.name == operator.default_run_name.replace(
+ assert isinstance(rendered_payload, dict)
+ assert rendered_payload["name"] == operator.default_run_name.replace(
  "{{ ts_nodash }}", data_interval_start.strftime("%Y%m%dT%H%M%S")
  )
 
@@ -160,7 +158,7 @@ def test_execute_handle_states(
  )
  operator = WherobotsRunOperator(
  task_id=f"test_execute_{uuid.uuid4()}",
- python=PythonRunPayload(uri=""),
+ run_python={"uri": ""},
  dag=dag,
  polling_interval=0,
  poll_logs=poll_logs,
@@ -186,11 +184,11 @@ def test_on_kill(
  operator = WherobotsRunOperator(
  task_id="test_render_template_python",
  name="test_run_{{ ds }}",
- python=PythonRunPayload(
- uri="s3://bucket/test-{{ ds }}.py",
- args=["{{ ds }}"],
- entrypoint="src.main_{{ ds }}",
- ),
+ run_python={
+ "uri": "s3://bucket/test-{{ ds }}.py",
+ "args": ["{{ ds }}"],
+ "entrypoint": "src.main_{{ ds }}",
+ },
  dag=dag,
  )
  operator.on_kill()
@@ -209,7 +207,7 @@ def test_poll_and_display_logs(self, mocker: MockerFixture):
  )
  operator = WherobotsRunOperator(
  task_id="test_poll_and_display_logs",
- python=PythonRunPayload(uri=""),
+ run_python={"uri": ""},
  dag=DAG("test_poll_and_display_logs"),
  )
  assert operator.poll_and_display_logs(hook, test_run, 0) == 2