From 22e689c8f3af4df3c1df507e5e931e03e38cd2a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= <tomas@tomasfarias.dev>
Date: Fri, 16 Aug 2024 10:19:36 +0200
Subject: [PATCH] fix: Start non-worker celery command with app.start()
 (#24375)

Co-authored-by: Julian Bez <julian@posthog.com>
---
 .run/Celery Beat.run.xml                      | 31 +++++++++++++++++++
 .run/Dev.run.xml                              |  1 +
 .../commands/run_autoreload_celery.py         | 25 +++++++--------
 3 files changed, 44 insertions(+), 13 deletions(-)
 create mode 100644 .run/Celery Beat.run.xml

diff --git a/.run/Celery Beat.run.xml b/.run/Celery Beat.run.xml
new file mode 100644
index 00000000000..56899ad0bba
--- /dev/null
+++ b/.run/Celery Beat.run.xml	
@@ -0,0 +1,31 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="Celery Beat" 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>
+      <env name="PYTHONUNBUFFERED" value="1" />
+      <env name="CLICKHOUSE_SECURE" value="False" />
+      <env name="DATABASE_URL" value="postgres://posthog:posthog@localhost:5432/posthog" />
+      <env name="DEBUG" value="1" />
+      <env name="KAFKA_HOSTS" value="localhost" />
+      <env name="SKIP_SERVICE_VERSION_REQUIREMENTS" value="1" />
+      <env name="REPLAY_EMBEDDINGS_ALLOWED_TEAM" value="1,2,3" />
+    </envs>
+    <option name="SDK_HOME" value="$PROJECT_DIR$/env/bin/python" />
+    <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
+    <option name="IS_MODULE_SDK" value="true" />
+    <option name="ADD_CONTENT_ROOTS" value="true" />
+    <option name="ADD_SOURCE_ROOTS" value="true" />
+    <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+    <option name="SCRIPT_NAME" value="$PROJECT_DIR$/manage.py" />
+    <option name="PARAMETERS" value="run_autoreload_celery --type=beat" />
+    <option name="SHOW_COMMAND_LINE" value="false" />
+    <option name="EMULATE_TERMINAL" value="false" />
+    <option name="MODULE_MODE" value="false" />
+    <option name="REDIRECT_INPUT" value="false" />
+    <option name="INPUT_FILE" value="" />
+    <method v="2" />
+  </configuration>
+</component>
\ No newline at end of file
diff --git a/.run/Dev.run.xml b/.run/Dev.run.xml
index 52b305471f1..48c1f35c7da 100644
--- a/.run/Dev.run.xml
+++ b/.run/Dev.run.xml
@@ -3,6 +3,7 @@
     <toRun name="PostHog" type="Python.DjangoServer" />
     <toRun name="Frontend" type="js.build_tools.npm" />
     <toRun name="Plugin Server" type="js.build_tools.npm" />
+    <toRun name="Celery Beat" type="PythonConfigurationType" />
     <toRun name="Celery Threads" type="PythonConfigurationType" />
     <method v="2" />
   </configuration>
diff --git a/posthog/management/commands/run_autoreload_celery.py b/posthog/management/commands/run_autoreload_celery.py
index cefb49df0ed..a9dd349eac4 100644
--- a/posthog/management/commands/run_autoreload_celery.py
+++ b/posthog/management/commands/run_autoreload_celery.py
@@ -1,4 +1,5 @@
 from typing import Literal
+
 import django
 from django.core.management.base import BaseCommand
 
@@ -23,22 +24,20 @@ class Command(BaseCommand):
     def run_celery_worker(type: Literal["worker", "beat"]):
         from posthog.celery import app as celery_app
 
-        args = (
-            [
-                "-A",
-                "posthog",
-                "worker",
-                "--pool=threads",
-                f"--queues={','.join(q.value for q in CeleryQueue)}",
-            ]
-            if type == "worker"
-            else [
-                "-A",
-                "posthog",
+        if type == "beat":
+            args = [
                 "beat",
                 "--scheduler",
                 "redbeat.RedBeatScheduler",
             ]
-        )
+            celery_app.start(argv=args)
+            return
 
+        args = [
+            "-A",
+            "posthog",
+            "worker",
+            "--pool=threads",
+            f"--queues={','.join(q.value for q in CeleryQueue)}",
+        ]
         celery_app.worker_main(args)