From 2c0dc75016d93c02fd4806157eb7ac966c57dedc Mon Sep 17 00:00:00 2001
From: Frank Hamand <frankhamand@gmail.com>
Date: Wed, 17 Apr 2024 11:39:39 +0100
Subject: [PATCH] feat: Switch from wsgi to asgi (attempt 3) (#21579)

* Try switching from wsgi to asgi

By default I don't believe this will change anything (everything runs
sync by default), but it will allow us to wrap e.g. clickhouse calls in
sync_to_async wrappers to stop them blocking for ages

* Trigger e2e tests

* Set env var with ASGI/WSGI info

* Fix sentry tracing sampling in ASGI

This was relying on wsgi_environ which is replaced by asgi_scope

* Remove bad headers typecast

* fix linting errors
---
 .github/workflows/ci-e2e.yml | 1 -
 posthog/asgi.py              | 8 ++++++++
 posthog/settings/sentry.py   | 6 ++++++
 posthog/wsgi.py              | 1 +
 unit.json                    | 3 ++-
 5 files changed, 17 insertions(+), 2 deletions(-)
 create mode 100644 posthog/asgi.py

diff --git a/.github/workflows/ci-e2e.yml b/.github/workflows/ci-e2e.yml
index 9322c51df37c5..79799757a7be9 100644
--- a/.github/workflows/ci-e2e.yml
+++ b/.github/workflows/ci-e2e.yml
@@ -59,7 +59,6 @@ jobs:
         timeout-minutes: 5
         outputs:
             chunks: ${{ steps.chunk.outputs.chunks }}
-
         steps:
             - name: Check out
               uses: actions/checkout@v3
diff --git a/posthog/asgi.py b/posthog/asgi.py
new file mode 100644
index 0000000000000..22912a0c7b76e
--- /dev/null
+++ b/posthog/asgi.py
@@ -0,0 +1,8 @@
+import os
+
+from django.core.asgi import get_asgi_application
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "posthog.settings")
+os.environ.setdefault("SERVER_GATEWAY_INTERFACE", "ASGI")
+
+application = get_asgi_application()
diff --git a/posthog/settings/sentry.py b/posthog/settings/sentry.py
index 225e6dc61f51a..d61d99077ed52 100644
--- a/posthog/settings/sentry.py
+++ b/posthog/settings/sentry.py
@@ -74,6 +74,12 @@ def traces_sampler(sampling_context: dict) -> float:
     if op == "http.server":
         path = sampling_context.get("wsgi_environ", {}).get("PATH_INFO")
         force_sample = bool(sampling_context.get("wsgi_environ", {}).get("HTTP_FORCE_SAMPLE"))
+        if os.environ.get("SERVER_GATEWAY_INTERFACE") == "ASGI":
+            path = sampling_context.get("asgi_scope", {}).get("path")
+            headers = sampling_context.get("asgi_scope", {}).get("headers", [])
+            for name, value in headers:
+                if name.lower().replace(b"_", b"-") == "force-sample":
+                    force_sample = bool(value)
 
         # HTTP header to force sampling set
         if force_sample:
diff --git a/posthog/wsgi.py b/posthog/wsgi.py
index 5be06203c3d26..b118b29fbf4bf 100644
--- a/posthog/wsgi.py
+++ b/posthog/wsgi.py
@@ -12,5 +12,6 @@
 from django.core.wsgi import get_wsgi_application
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "posthog.settings")
+os.environ.setdefault("SERVER_GATEWAY_INTERFACE", "WSGI")
 
 application = get_wsgi_application()
diff --git a/unit.json b/unit.json
index 0b8de8774edf1..5a5d5fa312708 100644
--- a/unit.json
+++ b/unit.json
@@ -43,7 +43,8 @@
             "processes": 4,
             "working_directory": "/code",
             "path": ".",
-            "module": "posthog.wsgi",
+            "module": "posthog.asgi",
+            "protocol": "asgi",
             "user": "nobody",
             "limits": {
                 "requests": 50000