From 95fec19aafb5724550b4743625d71058a2386705 Mon Sep 17 00:00:00 2001
From: Julian Bez <julian@posthog.com>
Date: Wed, 17 Jan 2024 11:54:12 +0000
Subject: [PATCH] feat(celery): Prepare to run on multiple queues (#19157)

* Add Celery queues env file with default queues

Reasoning:
We need to configure Celery workers in several places to consume
from a specific set of queues.

* Define some queues
---
 .run/Celery.run.xml                           |  1 +
 bin/celery-queues.env                         |  3 +
 bin/docker-worker-celery                      |  4 +
 bin/start-worker                              |  2 +
 .../test/__snapshots__/test_annotation.ambr   | 79 +++++++++++++++++++
 requirements.in                               |  1 +
 requirements.txt                              |  8 ++
 7 files changed, 98 insertions(+)
 create mode 100644 bin/celery-queues.env

diff --git a/.run/Celery.run.xml b/.run/Celery.run.xml
index 82f170a5fa381..9c1dd88226146 100644
--- a/.run/Celery.run.xml
+++ b/.run/Celery.run.xml
@@ -1,6 +1,7 @@
 <component name="ProjectRunConfigurationManager">
   <configuration default="false" name="Celery" type="PythonConfigurationType" factoryName="Python">
     <module name="posthog" />
+    <option name="ENV_FILES" value="$PROJECT_DIR$/bin/celery-queues.env" />
     <option name="INTERPRETER_OPTIONS" value="" />
     <option name="PARENT_ENVS" value="true" />
     <envs>
diff --git a/bin/celery-queues.env b/bin/celery-queues.env
new file mode 100644
index 0000000000000..97920a2fb9262
--- /dev/null
+++ b/bin/celery-queues.env
@@ -0,0 +1,3 @@
+# Default set of queues to be used by Celery.
+# Important: Add new queues to make Celery consume tasks from them.
+CELERY_WORKER_QUEUES=celery,email,insight_export,insight_refresh,gevent
\ No newline at end of file
diff --git a/bin/docker-worker-celery b/bin/docker-worker-celery
index 1f5166e95245b..4c1ff24c1c3ba 100755
--- a/bin/docker-worker-celery
+++ b/bin/docker-worker-celery
@@ -77,6 +77,10 @@ FLAGS+=("-n node@%h")
 # https://github.com/heroku/heroku-buildpack-python/blob/main/vendor/WEB_CONCURRENCY.sh
 [[ -n "${WEB_CONCURRENCY}" ]]    && FLAGS+=" --concurrency $WEB_CONCURRENCY"
 
+if [[ -z "${CELERY_WORKER_QUEUES}" ]]; then
+  source ./bin/celery-queues.env
+fi
+
 echo
 echo "SKIP_ASYNC_MIGRATIONS_SETUP=0 celery -A posthog worker ${FLAGS[*]}"
 echo
diff --git a/bin/start-worker b/bin/start-worker
index d88a235a5bb00..7fba86d7e4764 100755
--- a/bin/start-worker
+++ b/bin/start-worker
@@ -4,6 +4,8 @@ set -e
 # this kills all processes when the last one terminates
 trap 'kill $(jobs -p)' EXIT
 
+source ./bin/celery-queues.env
+
 # start celery worker with heartbeat (-B)
 SKIP_ASYNC_MIGRATIONS_SETUP=0 celery -A posthog worker -B --scheduler redbeat.RedBeatScheduler --without-heartbeat --without-gossip --without-mingle -Ofair -n node@%h &
 
diff --git a/posthog/api/test/__snapshots__/test_annotation.ambr b/posthog/api/test/__snapshots__/test_annotation.ambr
index 31298f4dbf007..2b761c79a0e3b 100644
--- a/posthog/api/test/__snapshots__/test_annotation.ambr
+++ b/posthog/api/test/__snapshots__/test_annotation.ambr
@@ -282,6 +282,85 @@
   LIMIT 1000 /*controller='project_annotations-list',route='api/projects/%28%3FP%3Cparent_lookup_team_id%3E%5B%5E/.%5D%2B%29/annotations/%3F%24'*/
   '''
 # ---
+# name: TestAnnotation.test_retrieving_annotation_is_not_n_plus_1.15
+  '''
+  SELECT "posthog_annotation"."id",
+         "posthog_annotation"."content",
+         "posthog_annotation"."created_at",
+         "posthog_annotation"."updated_at",
+         "posthog_annotation"."dashboard_item_id",
+         "posthog_annotation"."team_id",
+         "posthog_annotation"."organization_id",
+         "posthog_annotation"."created_by_id",
+         "posthog_annotation"."scope",
+         "posthog_annotation"."creation_type",
+         "posthog_annotation"."date_marker",
+         "posthog_annotation"."deleted",
+         "posthog_annotation"."apply_all",
+         "posthog_dashboarditem"."id",
+         "posthog_dashboarditem"."name",
+         "posthog_dashboarditem"."derived_name",
+         "posthog_dashboarditem"."description",
+         "posthog_dashboarditem"."team_id",
+         "posthog_dashboarditem"."filters",
+         "posthog_dashboarditem"."filters_hash",
+         "posthog_dashboarditem"."query",
+         "posthog_dashboarditem"."order",
+         "posthog_dashboarditem"."deleted",
+         "posthog_dashboarditem"."saved",
+         "posthog_dashboarditem"."created_at",
+         "posthog_dashboarditem"."last_refresh",
+         "posthog_dashboarditem"."refreshing",
+         "posthog_dashboarditem"."created_by_id",
+         "posthog_dashboarditem"."is_sample",
+         "posthog_dashboarditem"."short_id",
+         "posthog_dashboarditem"."favorited",
+         "posthog_dashboarditem"."refresh_attempt",
+         "posthog_dashboarditem"."last_modified_at",
+         "posthog_dashboarditem"."last_modified_by_id",
+         "posthog_dashboarditem"."dashboard_id",
+         "posthog_dashboarditem"."layouts",
+         "posthog_dashboarditem"."color",
+         "posthog_dashboarditem"."dive_dashboard_id",
+         "posthog_dashboarditem"."updated_at",
+         "posthog_dashboarditem"."deprecated_tags",
+         "posthog_dashboarditem"."tags",
+         "posthog_user"."id",
+         "posthog_user"."password",
+         "posthog_user"."last_login",
+         "posthog_user"."first_name",
+         "posthog_user"."last_name",
+         "posthog_user"."is_staff",
+         "posthog_user"."is_active",
+         "posthog_user"."date_joined",
+         "posthog_user"."uuid",
+         "posthog_user"."current_organization_id",
+         "posthog_user"."current_team_id",
+         "posthog_user"."email",
+         "posthog_user"."pending_email",
+         "posthog_user"."temporary_token",
+         "posthog_user"."distinct_id",
+         "posthog_user"."is_email_verified",
+         "posthog_user"."requested_password_reset_at",
+         "posthog_user"."has_seen_product_intro_for",
+         "posthog_user"."strapi_id",
+         "posthog_user"."email_opt_in",
+         "posthog_user"."theme_mode",
+         "posthog_user"."partial_notification_settings",
+         "posthog_user"."anonymize_data",
+         "posthog_user"."toolbar_mode",
+         "posthog_user"."events_column_config"
+  FROM "posthog_annotation"
+  LEFT OUTER JOIN "posthog_dashboarditem" ON ("posthog_annotation"."dashboard_item_id" = "posthog_dashboarditem"."id")
+  LEFT OUTER JOIN "posthog_user" ON ("posthog_annotation"."created_by_id" = "posthog_user"."id")
+  WHERE ((("posthog_annotation"."organization_id" = '00000000-0000-0000-0000-000000000000'::uuid
+           AND "posthog_annotation"."scope" = 'organization')
+          OR "posthog_annotation"."team_id" = 2)
+         AND NOT "posthog_annotation"."deleted")
+  ORDER BY "posthog_annotation"."date_marker" DESC
+  LIMIT 1000 /*controller='project_annotations-list',route='api/projects/%28%3FP%3Cparent_lookup_team_id%3E%5B%5E/.%5D%2B%29/annotations/%3F%24'*/
+  '''
+# ---
 # name: TestAnnotation.test_retrieving_annotation_is_not_n_plus_1.2
   '''
   SELECT "posthog_organizationmembership"."id",
diff --git a/requirements.in b/requirements.in
index b924b0df8420a..34b5f74f0a185 100644
--- a/requirements.in
+++ b/requirements.in
@@ -41,6 +41,7 @@ dnspython==2.2.1
 drf-exceptions-hog==0.4.0
 drf-extensions==0.7.0
 drf-spectacular==0.27.0
+gevent==23.9.1
 geoip2==4.6.0
 google-cloud-bigquery==3.11.4
 google-cloud-sqlcommenter==2.0.0
diff --git a/requirements.txt b/requirements.txt
index 0778ff9998a4b..5687cfc4afe91 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -234,6 +234,8 @@ future==0.18.3
     # via lzstring
 geoip2==4.6.0
     # via -r requirements.in
+gevent==23.9.1
+    # via -r requirements.in
 gitdb==4.0.11
     # via gitpython
 gitpython==3.1.40
@@ -263,6 +265,8 @@ googleapis-common-protos==1.60.0
     # via
     #   google-api-core
     #   grpcio-status
+greenlet==3.0.3
+    # via gevent
 grpcio==1.57.0
     # via
     #   google-api-core
@@ -681,6 +685,10 @@ yarl==1.7.2
     # via aiohttp
 zipp==3.17.0
     # via importlib-metadata
+zope-event==5.0
+    # via gevent
+zope-interface==6.1
+    # via gevent
 
 # The following packages are considered to be unsafe in a requirements file:
 # setuptools