From 34378a19dc2e851708017d7dabdb8e1d1ca28068 Mon Sep 17 00:00:00 2001
From: razumeiko <2330426+razumeiko@users.noreply.github.com>
Date: Fri, 27 Aug 2021 15:25:48 +0300
Subject: [PATCH] Added AWS Lambda Python 3.9 runtime support

---
 .craft.yml                                |  1 +
 sentry_sdk/integrations/aws_lambda.py     | 18 ++++++++++++++----
 tests/integrations/aws_lambda/test_aws.py |  4 +++-
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/.craft.yml b/.craft.yml
index c6d13cfc2c..864d689271 100644
--- a/.craft.yml
+++ b/.craft.yml
@@ -21,6 +21,7 @@ targets:
           - python3.6
           - python3.7
           - python3.8
+          - python3.9
     license: MIT
 changelog: CHANGELOG.md
 changelogPolicy: simple
diff --git a/sentry_sdk/integrations/aws_lambda.py b/sentry_sdk/integrations/aws_lambda.py
index 533250efaa..0eae710bff 100644
--- a/sentry_sdk/integrations/aws_lambda.py
+++ b/sentry_sdk/integrations/aws_lambda.py
@@ -284,12 +284,14 @@ def get_lambda_bootstrap():
     # Python 3.7: If the bootstrap module is *already imported*, it is the
     # one we actually want to use (no idea what's in __main__)
     #
-    # On Python 3.8 bootstrap is also importable, but will be the same file
+    # Python 3.8: bootstrap is also importable, but will be the same file
     # as __main__ imported under a different name:
     #
     #     sys.modules['__main__'].__file__ == sys.modules['bootstrap'].__file__
     #     sys.modules['__main__'] is not sys.modules['bootstrap']
     #
+    # Python 3.9: bootstrap is in __main__.awslambdaricmain
+    #
     # On container builds using the `aws-lambda-python-runtime-interface-client`
     # (awslamdaric) module, bootstrap is located in sys.modules['__main__'].bootstrap
     #
@@ -297,10 +299,18 @@ def get_lambda_bootstrap():
     if "bootstrap" in sys.modules:
         return sys.modules["bootstrap"]
     elif "__main__" in sys.modules:
-        if hasattr(sys.modules["__main__"], "bootstrap"):
+        module = sys.modules["__main__"]
+        # python3.9 runtime
+        if hasattr(module, "awslambdaricmain") and hasattr(
+            module.awslambdaricmain, "bootstrap"  # type: ignore
+        ):
+            return module.awslambdaricmain.bootstrap  # type: ignore
+        elif hasattr(module, "bootstrap"):
             # awslambdaric python module in container builds
-            return sys.modules["__main__"].bootstrap  # type: ignore
-        return sys.modules["__main__"]
+            return module.bootstrap  # type: ignore
+
+        # python3.8 runtime
+        return module
     else:
         return None
 
diff --git a/tests/integrations/aws_lambda/test_aws.py b/tests/integrations/aws_lambda/test_aws.py
index 0f50753be7..c9084beb14 100644
--- a/tests/integrations/aws_lambda/test_aws.py
+++ b/tests/integrations/aws_lambda/test_aws.py
@@ -105,7 +105,9 @@ def lambda_client():
     return get_boto_client()
 
 
-@pytest.fixture(params=["python3.6", "python3.7", "python3.8", "python2.7"])
+@pytest.fixture(
+    params=["python3.6", "python3.7", "python3.8", "python3.9", "python2.7"]
+)
 def lambda_runtime(request):
     return request.param