From b25a75b67048ebcc3757bbf676d2ac2bf4479d59 Mon Sep 17 00:00:00 2001
From: Christopher DeCarolis <cdguitar817@gmail.com>
Date: Wed, 20 Nov 2024 17:29:04 -0800
Subject: [PATCH] [dagster-airlift] mwaa_session -> mwaa_client (#26066)

In the 1.9 rush I incorrectly updated MwaaSessionAuthBackend to take a
mwaa_session instead of a mwaa_client. This has been fixed.
---
 .../dagster-airlift/dagster_airlift/mwaa/auth.py       | 10 +++++-----
 .../unit_tests/mwaa_tests/test_mwaa_auth.py            |  2 +-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/python_modules/libraries/dagster-airlift/dagster_airlift/mwaa/auth.py b/python_modules/libraries/dagster-airlift/dagster_airlift/mwaa/auth.py
index 5b80ab8d7276b..ea07f4aaac26d 100644
--- a/python_modules/libraries/dagster-airlift/dagster_airlift/mwaa/auth.py
+++ b/python_modules/libraries/dagster-airlift/dagster_airlift/mwaa/auth.py
@@ -48,19 +48,19 @@ class MwaaSessionAuthBackend(AirflowAuthBackend):
             from dagster_airlift.mwaa import MwaaSessionAuthBackend
             from dagster_airlift.core import AirflowInstance
 
-            boto_session = boto3.Session(profile_name="my_profile", region_name="us-west-2")
+            boto_client = boto3.client("mwaa")
             af_instance = AirflowInstance(
                 name="my-mwaa-instance",
                 auth_backend=MwaaSessionAuthBackend(
-                    mwaa_session=boto_session,
+                    mwaa_client=boto_client,
                     env_name="my-mwaa-env"
                 )
             )
 
     """
 
-    def __init__(self, mwaa_session: boto3.Session, env_name: str) -> None:
-        self.mwaa_client = mwaa_session
+    def __init__(self, mwaa_client: Any, env_name: str) -> None:
+        self.mwaa_client = mwaa_client
         self.env_name = env_name
         # Session info is generated when we either try to retrieve a session or retrieve the web server url
         self._session_info: Optional[Tuple[str, str]] = None
@@ -69,7 +69,7 @@ def __init__(self, mwaa_session: boto3.Session, env_name: str) -> None:
     def from_profile(region: str, env_name: str, profile_name: Optional[str] = None):
         boto_session = boto3.Session(profile_name=profile_name, region_name=region)
         mwaa = boto_session.client("mwaa")
-        return MwaaSessionAuthBackend(mwaa_session=mwaa, env_name=env_name)
+        return MwaaSessionAuthBackend(mwaa_client=mwaa, env_name=env_name)
 
     def get_session(self) -> requests.Session:
         # Get the session info
diff --git a/python_modules/libraries/dagster-airlift/dagster_airlift_tests/unit_tests/mwaa_tests/test_mwaa_auth.py b/python_modules/libraries/dagster-airlift/dagster_airlift_tests/unit_tests/mwaa_tests/test_mwaa_auth.py
index f05ce1a444622..a876f065d19b9 100644
--- a/python_modules/libraries/dagster-airlift/dagster_airlift_tests/unit_tests/mwaa_tests/test_mwaa_auth.py
+++ b/python_modules/libraries/dagster-airlift/dagster_airlift_tests/unit_tests/mwaa_tests/test_mwaa_auth.py
@@ -20,7 +20,7 @@ def test_mwaa_session_auth_direct_mwaa_client_creation() -> None:
         mock_get_session_info.return_value = ("my-webserver-hostname", "my-session-cookie")
         boto_session = boto3.Session(region_name="us-west-2")
         mwaa = boto_session.client("mwaa")
-        auth_backend = MwaaSessionAuthBackend(mwaa_session=mwaa, env_name="my-env")
+        auth_backend = MwaaSessionAuthBackend(mwaa_client=mwaa, env_name="my-env")
         session = auth_backend.get_session()
         assert session.cookies["session"] == "my-session-cookie"
         assert auth_backend.get_webserver_url() == "https://my-webserver-hostname"