Repository: PostHog/posthog
Pull request: fix optimize table timeout on 0002
## Changes

closes https://github.com/PostHog/posthog/issues/8511 

OPTIMIZE TABLE times out for larger customers. 

## How did you test this code?

<!-- If the answer is manually, please include a quick step-by-step on how to test this PR. -->

*Please describe.*


Problem statement:
Async events migration to handle large events volumes: timeout error
We timed out on a self-hosted instance with 45G of events, that's not the biggest. 
The timeout should be configurable from the UI probably.
Ideally we'll want to test to the size of the largest users.

Hints from issue discussion:

Commit message: refactor

Patch to review:
diff --git a/posthog/async_migrations/definition.py b/posthog/async_migrations/definition.py
index d78782bae57b8..d1dd86d5431a6 100644
--- a/posthog/async_migrations/definition.py
+++ b/posthog/async_migrations/definition.py
@@ -1,6 +1,8 @@
 from datetime import datetime
 from typing import Callable, List, Optional, Tuple
 
+from constance import config
+
 from posthog.constants import AnalyticsDBMS
 from posthog.version_requirement import ServiceVersionRequirement
 
@@ -60,7 +62,12 @@ def simple_op(
         )
 
     @classmethod
-    def get_db_op(cls, sql: str, database: AnalyticsDBMS = AnalyticsDBMS.CLICKHOUSE, timeout_seconds: int = 60):
+    def get_db_op(
+        cls,
+        sql: str,
+        database: AnalyticsDBMS = AnalyticsDBMS.CLICKHOUSE,
+        timeout_seconds: int = getattr(config, "ASYNC_MIGRATIONS_DEFAULT_TIMEOUT"),
+    ):
         from posthog.async_migrations.utils import execute_op_clickhouse, execute_op_postgres
 
         # timeout is currently CH only
diff --git a/posthog/async_migrations/migrations/0002_events_sample_by.py b/posthog/async_migrations/migrations/0002_events_sample_by.py
index 08a31124e7f4c..fca319ad8cd87 100644
--- a/posthog/async_migrations/migrations/0002_events_sample_by.py
+++ b/posthog/async_migrations/migrations/0002_events_sample_by.py
@@ -104,6 +104,17 @@ def operations(self):
 
         last_partition_op = [generate_insert_into_op(self._partitions[-1] if len(self._partitions) > 0 else 0)]
 
+        def optimize_table_fn():
+            default_timeout = getattr(config, "ASYNC_MIGRATIONS_DEFAULT_TIMEOUT")
+            sync_execute(
+                "OPTIMIZE TABLE {EVENTS_TABLE_NAME} FINAL",
+                settings={
+                    "max_execution_time": default_timeout,
+                    "send_timeout": default_timeout,
+                    "receive_timeout": default_timeout,
+                },
+            )
+
         post_insert_ops = [
             AsyncMigrationOperation.simple_op(
                 database=AnalyticsDBMS.CLICKHOUSE,
@@ -120,12 +131,7 @@ def operations(self):
                     ON CLUSTER {CLICKHOUSE_CLUSTER}
                 """,
             ),
-            AsyncMigrationOperation.simple_op(
-                database=AnalyticsDBMS.CLICKHOUSE,
-                sql=f"OPTIMIZE TABLE {EVENTS_TABLE_NAME} FINAL",
-                rollback="",
-                resumable=True,
-            ),
+            AsyncMigrationOperation.simple_op(database=AnalyticsDBMS.CLICKHOUSE, fn=optimize_table_fn, rollback=""),
             AsyncMigrationOperation.simple_op(
                 database=AnalyticsDBMS.CLICKHOUSE,
                 sql=f"ATTACH TABLE {EVENTS_TABLE_NAME}_mv ON CLUSTER {CLICKHOUSE_CLUSTER}",
diff --git a/posthog/settings/dynamic_settings.py b/posthog/settings/dynamic_settings.py
index f6065c345e487..7377043cb605f 100644
--- a/posthog/settings/dynamic_settings.py
+++ b/posthog/settings/dynamic_settings.py
@@ -36,6 +36,11 @@
         "Used to disable automatic rollback of failed async migrations.",
         bool,
     ),
+    "ASYNC_MIGRATIONS_DEFAULT_TIMEOUT": (
+        get_from_env("ASYNC_MIGRATIONS_DEFAULT_TIMEOUT_SECONDS", 60 * 30, type_cast=int),
+        "Sets the default timeout for completing an async migration operation.",
+        int,
+    ),
     "EMAIL_ENABLED": (
         get_from_env("EMAIL_ENABLED", True, type_cast=str_to_bool),
         "Whether email service is enabled or not.",
@@ -93,6 +98,7 @@
     "AUTO_START_ASYNC_MIGRATIONS",
     "ASYNC_MIGRATIONS_ROLLBACK_TIMEOUT",
     "ASYNC_MIGRATIONS_DISABLE_AUTO_ROLLBACK",
+    "ASYNC_MIGRATIONS_DEFAULT_TIMEOUT",
     "EMAIL_ENABLED",
     "EMAIL_HOST",
     "EMAIL_PORT",


Code review comments:
diff_hunk: @@ -60,7 +62,12 @@ def simple_op(
         )
 
     @classmethod
-    def get_db_op(cls, sql: str, database: AnalyticsDBMS = AnalyticsDBMS.CLICKHOUSE, timeout_seconds: int = 60):
+    def get_db_op(
+        cls,
+        sql: str,
+        database: AnalyticsDBMS = AnalyticsDBMS.CLICKHOUSE,
+        timeout_seconds: int = getattr(config, "ASYNC_MIGRATIONS_DEFAULT_TIMEOUT"),
path: posthog/async_migrations/definition.py
text: @user1:
change it above too :55

diff_hunk: @@ -120,12 +131,7 @@ def operations(self):
                     ON CLUSTER {CLICKHOUSE_CLUSTER}
                 """,
             ),
-            AsyncMigrationOperation.simple_op(
-                database=AnalyticsDBMS.CLICKHOUSE,
-                sql=f"OPTIMIZE TABLE {EVENTS_TABLE_NAME} FINAL",
-                rollback="",
-                resumable=True,
-            ),
+            AsyncMigrationOperation.simple_op(database=AnalyticsDBMS.CLICKHOUSE, fn=optimize_table_fn, rollback=""),
path: posthog/async_migrations/migrations/0002_events_sample_by.py
text: @user1:
This isn't `simple_op` anymore & add something as `debug_context` too, I just landed that PR.
```suggestion
            AsyncMigrationOperation(fn=optimize_table_fn),
```

diff_hunk: @@ -104,6 +104,17 @@ def operations(self):
 
         last_partition_op = [generate_insert_into_op(self._partitions[-1] if len(self._partitions) > 0 else 0)]
 
+        def optimize_table_fn():
+            default_timeout = getattr(config, "ASYNC_MIGRATIONS_DEFAULT_TIMEOUT")
+            sync_execute(
path: posthog/async_migrations/migrations/0002_events_sample_by.py
text: @user1:
we're missing query ID here too + any other future changes there, so instead add additional_setting or sth like that to `execute_op_clickhouse` and use that fn here too.