Fix: Automatically refresh the MWAA CLI token (#1564)

Author: izeigerman

docs/integrations/airflow.md

@@ -84,11 +84,3 @@ default_scheduler:
     type: mwaa
     environment: <The MWAA Environment Name>
 ```
-
-Alternatively, the Airflow Webserver URL and the MWAA CLI token can be provided directly instead of the environment name:
-```yaml linenums="1"
-default_scheduler:
-    type: mwaa
-    airflow_url: https://<Airflow Webserver Host>/
-    auth_token: <The MWAA CLI Token>
-```

sqlmesh/core/config/scheduler.py

@@ -220,8 +220,6 @@ class MWAASchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig):
 
     Args:
         environment: The name of the MWAA environment.
-        airflow_url: The URL of the Airflow Webserver.
-        auth_token: The MWAA authentication token.
         dag_run_poll_interval_secs: Determines how often a running DAG can be polled (in seconds).
         dag_creation_poll_interval_secs: Determines how often SQLMesh should check whether a DAG has been created (in seconds).
         dag_creation_max_retry_attempts: Determines the maximum number of attempts that SQLMesh will make while checking for
@@ -230,9 +228,7 @@ class MWAASchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig):
         ddl_concurrent_tasks: The number of concurrent tasks used for DDL operations (table / view creation, deletion, etc).
     """
 
-    environment: t.Optional[str] = None
-    airflow_url: t.Optional[str] = None
-    auth_token: t.Optional[str] = None
+    environment: str
     dag_run_poll_interval_secs: int = 10
     dag_creation_poll_interval_secs: int = 30
     dag_creation_max_retry_attempts: int = 10
@@ -245,20 +241,10 @@ class MWAASchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig):
     _concurrent_tasks_validator = concurrent_tasks_validator
 
     def create_plan_evaluator(self, context: Context) -> PlanEvaluator:
-        from sqlmesh.schedulers.airflow.mwaa_client import (
-            MWAAClient,
-            url_and_auth_token_for_environment,
-        )
-
-        if self.environment:
-            airflow_url, auth_token = url_and_auth_token_for_environment(self.environment)
-        else:
-            assert self.airflow_url and self.auth_token  # Make mypy happy
-            airflow_url = self.airflow_url
-            auth_token = self.auth_token
+        from sqlmesh.schedulers.airflow.mwaa_client import MWAAClient
 
         return MWAAPlanEvaluator(
-            client=MWAAClient(airflow_url, auth_token, console=context.console),
+            client=MWAAClient(self.environment, console=context.console),
             state_sync=context.state_sync,
             console=context.console,
             dag_run_poll_interval_secs=self.dag_run_poll_interval_secs,
@@ -270,19 +256,6 @@ def create_plan_evaluator(self, context: Context) -> PlanEvaluator:
             users=context.users,
         )
 
-    @model_validator(mode="before")
-    @classmethod
-    def _ensure_environment_or_url_with_auth_token(
-        cls, values: t.Dict[str, t.Any]
-    ) -> t.Dict[str, t.Any]:
-        if not values.get("environment"):
-            if not values.get("airflow_url") or not values.get("auth_token"):
-                raise ValueError(
-                    "Either 'environment' or 'airflow_url' and 'auth_token' must be specified for the MWAA scheduler config."
-                )
-
-        return values
-
 
 SchedulerConfig = Annotated[
     t.Union[

sqlmesh/schedulers/airflow/mwaa_client.py

@@ -2,24 +2,31 @@
 
 import base64
 import json
+import logging
 import typing as t
 from urllib.parse import urljoin
 
 from requests import Session
 
 from sqlmesh.core.console import Console
 from sqlmesh.schedulers.airflow.client import BaseAirflowClient, raise_for_status
+from sqlmesh.utils.date import now_timestamp
 from sqlmesh.utils.errors import NotFoundError
 
+logger = logging.getLogger(__name__)
+
+
+TOKEN_TTL_MS = 30 * 1000
+
 
 class MWAAClient(BaseAirflowClient):
-    def __init__(self, airflow_url: str, auth_token: str, console: t.Optional[Console] = None):
+    def __init__(self, environment: str, console: t.Optional[Console] = None):
+        airflow_url, auth_token = url_and_auth_token_for_environment(environment)
         super().__init__(airflow_url, console)
 
-        self._session = Session()
-        self._session.headers.update(
-            {"Authorization": f"Bearer {auth_token}", "Content-Type": "text/plain"}
-        )
+        self._environment = environment
+        self._last_token_refresh_ts = now_timestamp()
+        self.__session: Session = _create_session(auth_token)
 
     def get_first_dag_run_id(self, dag_id: str) -> t.Optional[str]:
         dag_runs = self._list_dag_runs(dag_id)
@@ -49,10 +56,27 @@ def _post(self, data: str) -> t.Tuple[str, str]:
         cli_stderr = base64.b64decode(response_body["stderr"]).decode("utf8").strip()
         return cli_stdout, cli_stderr
 
+    @property
+    def _session(self) -> Session:
+        current_ts = now_timestamp()
+        if current_ts - self._last_token_refresh_ts > TOKEN_TTL_MS:
+            _, auth_token = url_and_auth_token_for_environment(self._environment)
+            self.__session = _create_session(auth_token)
+            self._last_token_refresh_ts = current_ts
+        return self.__session
+
+
+def _create_session(auth_token: str) -> Session:
+    session = Session()
+    session.headers.update({"Authorization": f"Bearer {auth_token}", "Content-Type": "text/plain"})
+    return session
+
 
 def url_and_auth_token_for_environment(environment_name: str) -> t.Tuple[str, str]:
     import boto3
 
+    logger.info("Fetching the MWAA CLI token")
+
     client = boto3.client("mwaa")
     cli_token = client.create_cli_token(Name=environment_name)
 

tests/schedulers/airflow/test_mwaa_client.py

@@ -16,14 +16,20 @@ def test_get_first_dag_run_id(mocker: MockerFixture):
     list_runs_mock = mocker.patch("requests.Session.post")
     list_runs_mock.return_value = list_runs_response_mock
 
-    client = MWAAClient("https://test_airflow_host", "test_token")
+    url_and_auth_token_mock = mocker.patch(
+        "sqlmesh.schedulers.airflow.mwaa_client.url_and_auth_token_for_environment"
+    )
+    url_and_auth_token_mock.return_value = ("https://test_airflow_host", "test_token")
+
+    client = MWAAClient("test_environment")
 
     assert client.get_first_dag_run_id("test_dag_id") == "test_run_id"
 
     list_runs_mock.assert_called_once_with(
         "https://test_airflow_host/aws_mwaa/cli",
         data="dags list-runs -o json -d test_dag_id",
     )
+    url_and_auth_token_mock.assert_called_once_with("test_environment")
 
 
 def test_get_dag_run_state(mocker: MockerFixture):
@@ -43,14 +49,56 @@ def test_get_dag_run_state(mocker: MockerFixture):
     list_runs_mock = mocker.patch("requests.Session.post")
     list_runs_mock.return_value = list_runs_response_mock
 
-    client = MWAAClient("https://test_airflow_host", "test_token")
+    url_and_auth_token_mock = mocker.patch(
+        "sqlmesh.schedulers.airflow.mwaa_client.url_and_auth_token_for_environment"
+    )
+    url_and_auth_token_mock.return_value = ("https://test_airflow_host", "test_token")
+
+    client = MWAAClient("test_environment")
 
     assert client.get_dag_run_state("test_dag_id", "test_run_id_b") == "failed"
 
     list_runs_mock.assert_called_once_with(
         "https://test_airflow_host/aws_mwaa/cli",
         data="dags list-runs -o json -d test_dag_id",
     )
+    url_and_auth_token_mock.assert_called_once_with("test_environment")
+
+
+def test_token_refresh(mocker: MockerFixture):
+    list_runs_response_mock = mocker.Mock()
+    list_runs_response_mock.json.return_value = {
+        "stdout": _encode_output(json.dumps([{"run_id": "test_run_id", "state": "success"}])),
+        "stderr": "",
+    }
+    list_runs_response_mock.status_code = 200
+    list_runs_mock = mocker.patch("requests.Session.post")
+    list_runs_mock.return_value = list_runs_response_mock
+
+    url_and_auth_token_mock = mocker.patch(
+        "sqlmesh.schedulers.airflow.mwaa_client.url_and_auth_token_for_environment"
+    )
+    url_and_auth_token_mock.return_value = ("https://test_airflow_host", "test_token")
+
+    now_mock = mocker.patch("sqlmesh.schedulers.airflow.mwaa_client.now_timestamp")
+    now_mock.return_value = 0
+
+    client = MWAAClient("test_environment")
+    client.get_first_dag_run_id("test_dag_id")
+
+    now_mock.return_value = 15000  # 15 seconds later
+    client.get_first_dag_run_id("test_dag_id")
+
+    now_mock.return_value = 31000  # 31 seconds later
+    client.get_first_dag_run_id("test_dag_id")
+
+    now_mock.return_value = 45000  # 45 seconds later
+    client.get_first_dag_run_id("test_dag_id")
+
+    now_mock.return_value = 63000  # 63 seconds later
+    client.get_first_dag_run_id("test_dag_id")
+
+    assert url_and_auth_token_mock.call_count == 3
 
 
 def _encode_output(out: str) -> str: