From 09b6933032a86dbb4b1f18307cedf41f061afc3e Mon Sep 17 00:00:00 2001
From: Sandy Spicer <sandy@posthog.com>
Date: Wed, 5 Jun 2024 12:32:02 -0700
Subject: [PATCH] feat(dev): add xdist for local concurrency in pytest (#22684)

---
 bin/tests                                     |   9 +-
 docker/clickhouse/config.xml                  |   2 +-
 ee/clickhouse/materialized_columns/columns.py |  10 +-
 .../test/test_trends_data_warehouse_query.py  |   3 +-
 posthog/settings/data_stores.py               |  26 ++-
 posthog/test/base.py                          |  21 ++-
 requirements-dev.in                           |   1 +
 requirements-dev.txt                          | 154 +++++++++++++++---
 8 files changed, 184 insertions(+), 42 deletions(-)

diff --git a/bin/tests b/bin/tests
index 21120dbc3fb..395f01c519a 100755
--- a/bin/tests
+++ b/bin/tests
@@ -30,4 +30,11 @@ PG_PASSWORD="${PGPASSWORD:=posthog}"
 PG_PORT="${PGPORT:=5432}"
 PGOPTIONS='--client-min-messages=warning' psql posthog -d "postgres://${PG_USER}:${PG_PASSWORD}@${PG_HOST}:${PG_PORT}" -c "drop database if exists test_posthog" 1> /dev/null
 
-nodemon -w ./posthog -w ./hogvm/python -w ./ee --ext py --exec "OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES pytest --reuse-db --durations-min=2.0 ${MIGRATIONS} -s $* --snapshot-update; mypy -p posthog | mypy-baseline filter"
+if [ -z "$XDIST_WORKERS" ]
+then
+    TEST_CONCURRENCY=""
+else
+    TEST_CONCURRENCY="-n $XDIST_WORKERS"
+fi
+
+nodemon -w ./posthog -w ./hogvm/python -w ./ee --ext py --exec "OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES pytest --reuse-db --durations-min=2.0 ${MIGRATIONS} ${TEST_CONCURRENCY} -s $* --snapshot-update; mypy -p posthog | mypy-baseline filter"
\ No newline at end of file
diff --git a/docker/clickhouse/config.xml b/docker/clickhouse/config.xml
index 8a6c12eac00..bc1c691f8b4 100644
--- a/docker/clickhouse/config.xml
+++ b/docker/clickhouse/config.xml
@@ -285,7 +285,7 @@
     -->
 
     <!-- Maximum number of concurrent queries. -->
-    <max_concurrent_queries>100</max_concurrent_queries>
+    <max_concurrent_queries>200</max_concurrent_queries>
 
     <!-- Maximum memory usage (resident set size) for server process.
          Zero value or unset means default. Default is "max_server_memory_usage_to_ram_ratio" of available
diff --git a/ee/clickhouse/materialized_columns/columns.py b/ee/clickhouse/materialized_columns/columns.py
index 0ac9cf18d02..9961359a99c 100644
--- a/ee/clickhouse/materialized_columns/columns.py
+++ b/ee/clickhouse/materialized_columns/columns.py
@@ -82,7 +82,7 @@ def materialize(
             {column_name} VARCHAR MATERIALIZED {TRIM_AND_EXTRACT_PROPERTY.format(table_column=table_column)}
         """,
             {"property": property},
-            settings={"alter_sync": 1},
+            settings={"alter_sync": 2 if TEST else 1},
         )
         sync_execute(
             f"""
@@ -91,7 +91,7 @@ def materialize(
             ADD COLUMN IF NOT EXISTS
             {column_name} VARCHAR
         """,
-            settings={"alter_sync": 1},
+            settings={"alter_sync": 2 if TEST else 1},
         )
     else:
         sync_execute(
@@ -102,13 +102,13 @@ def materialize(
             {column_name} VARCHAR MATERIALIZED {TRIM_AND_EXTRACT_PROPERTY.format(table_column=table_column)}
         """,
             {"property": property},
-            settings={"alter_sync": 1},
+            settings={"alter_sync": 2 if TEST else 1},
         )
 
     sync_execute(
         f"ALTER TABLE {table} {execute_on_cluster} COMMENT COLUMN {column_name} %(comment)s",
         {"comment": f"column_materializer::{table_column}::{property}"},
-        settings={"alter_sync": 1},
+        settings={"alter_sync": 2 if TEST else 1},
     )
 
     if create_minmax_index:
@@ -130,7 +130,7 @@ def add_minmax_index(table: TablesWithMaterializedColumns, column_name: str):
             ADD INDEX {index_name} {column_name}
             TYPE minmax GRANULARITY 1
             """,
-            settings={"alter_sync": 1},
+            settings={"alter_sync": 2 if TEST else 1},
         )
     except ServerException as err:
         if "index with this name already exists" not in str(err):
diff --git a/posthog/hogql_queries/insights/trends/test/test_trends_data_warehouse_query.py b/posthog/hogql_queries/insights/trends/test/test_trends_data_warehouse_query.py
index 15a3060004c..a606d560cf1 100644
--- a/posthog/hogql_queries/insights/trends/test/test_trends_data_warehouse_query.py
+++ b/posthog/hogql_queries/insights/trends/test/test_trends_data_warehouse_query.py
@@ -26,6 +26,7 @@ from posthog.settings import (
     OBJECT_STORAGE_BUCKET,
     OBJECT_STORAGE_ENDPOINT,
     OBJECT_STORAGE_SECRET_ACCESS_KEY,
+    XDIST_SUFFIX,
 )
 import s3fs
 from pyarrow import parquet as pq
@@ -39,7 +40,7 @@ from posthog.hogql_queries.legacy_compatibility.filter_to_query import (
     clean_entity_properties,
 )
 
-TEST_BUCKET = "test_storage_bucket-posthog.hogql.datawarehouse.trendquery"
+TEST_BUCKET = "test_storage_bucket-posthog.hogql.datawarehouse.trendquery" + XDIST_SUFFIX
 
 
 class TestTrendsDataWarehouseQuery(ClickhouseTestMixin, BaseTest):
diff --git a/posthog/settings/data_stores.py b/posthog/settings/data_stores.py
index fa92489bd83..b2985242c52 100644
--- a/posthog/settings/data_stores.py
+++ b/posthog/settings/data_stores.py
@@ -124,8 +124,25 @@ replica_opt_in = os.environ.get("READ_REPLICA_OPT_IN", "")
 READ_REPLICA_OPT_IN: list[str] = get_list(replica_opt_in)
 
 
+# Xdist Settings
+# When running concurrent tests, PYTEST_XDIST_WORKER gets set to "gw0" ... "gwN"
+# We use this setting to create multiple databases to achieve test isolation
+PYTEST_XDIST_WORKER: str | None = os.getenv("PYTEST_XDIST_WORKER")
+PYTEST_XDIST_WORKER_NUM: int | None = None
+SUFFIX = ""
+XDIST_SUFFIX = ""
+try:
+    if PYTEST_XDIST_WORKER is not None:
+        XDIST_SUFFIX = f"_{PYTEST_XDIST_WORKER}"
+        PYTEST_XDIST_WORKER_NUM = int("".join([x for x in PYTEST_XDIST_WORKER if x.isdigit()]))
+except:
+    pass
+
+if TEST:
+    SUFFIX = "_test" + XDIST_SUFFIX
+
 # Clickhouse Settings
-CLICKHOUSE_TEST_DB: str = "posthog_test"
+CLICKHOUSE_TEST_DB: str = "posthog" + SUFFIX
 
 CLICKHOUSE_HOST: str = os.getenv("CLICKHOUSE_HOST", "localhost")
 CLICKHOUSE_OFFLINE_CLUSTER_HOST: str | None = os.getenv("CLICKHOUSE_OFFLINE_CLUSTER_HOST", None)
@@ -222,8 +239,6 @@ KAFKA_SASL_MECHANISM = os.getenv("KAFKA_SASL_MECHANISM", None)
 KAFKA_SASL_USER = os.getenv("KAFKA_SASL_USER", None)
 KAFKA_SASL_PASSWORD = os.getenv("KAFKA_SASL_PASSWORD", None)
 
-SUFFIX = "_test" if TEST else ""
-
 KAFKA_EVENTS_PLUGIN_INGESTION: str = (
     f"{KAFKA_PREFIX}events_plugin_ingestion{SUFFIX}"  # can be overridden in settings.py
 )
@@ -241,7 +256,10 @@ TOKENS_HISTORICAL_DATA = os.getenv("TOKENS_HISTORICAL_DATA", "").split(",")
 
 # The last case happens when someone upgrades Heroku but doesn't have Redis installed yet. Collectstatic gets called before we can provision Redis.
 if TEST or DEBUG or IS_COLLECT_STATIC:
-    REDIS_URL = os.getenv("REDIS_URL", "redis://localhost/")
+    if PYTEST_XDIST_WORKER_NUM is not None:
+        REDIS_URL = os.getenv("REDIS_URL", f"redis://localhost/{PYTEST_XDIST_WORKER_NUM}")
+    else:
+        REDIS_URL = os.getenv("REDIS_URL", "redis://localhost/")
 else:
     REDIS_URL = os.getenv("REDIS_URL", "")
 
diff --git a/posthog/test/base.py b/posthog/test/base.py
index 59a0c3c171f..f93ad085cbb 100644
--- a/posthog/test/base.py
+++ b/posthog/test/base.py
@@ -401,14 +401,21 @@ def cleanup_materialized_columns():
         # EE not available? Skip
         return
 
+    def optionally_drop(table, filter=None):
+        drops = ",".join(
+            [
+                f"DROP COLUMN {column_name}"
+                for column_name in get_materialized_columns(table).values()
+                if filter is None or filter(column_name)
+            ]
+        )
+        if drops:
+            sync_execute(f"ALTER TABLE {table} {drops}")
+
     default_columns = default_materialised_columns()
-    for column_name in get_materialized_columns("events").values():
-        if column_name not in default_columns:
-            sync_execute(f"ALTER TABLE events DROP COLUMN {column_name}")
-    for column_name in get_materialized_columns("person").values():
-        sync_execute(f"ALTER TABLE person DROP COLUMN {column_name}")
-    for column_name in get_materialized_columns("groups").values():
-        sync_execute(f"ALTER TABLE groups DROP COLUMN {column_name}")
+    optionally_drop("events", lambda name: name not in default_columns)
+    optionally_drop("person")
+    optionally_drop("groups")
 
 
 def also_test_with_materialized_columns(
diff --git a/requirements-dev.in b/requirements-dev.in
index 66be8378356..691e790df2b 100644
--- a/requirements-dev.in
+++ b/requirements-dev.in
@@ -44,6 +44,7 @@ pytest-icdiff==0.6
 pytest-mock==3.11.1
 pytest-split==0.8.1
 pytest-watch==4.2.0
+pytest-xdist==3.6.1
 python-dateutil>=2.8.2
 responses==0.23.1
 syrupy~=4.6.0
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 3afb2e935ea..6a0b72290d4 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,70 +1,111 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements-dev.in -o requirements-dev.txt
 aiohttp==3.9.3
-    # via aioresponses
+    # via
+    #   -c requirements.txt
+    #   aioresponses
 aioresponses==0.7.6
+    # via -r requirements-dev.in
 aiosignal==1.2.0
-    # via aiohttp
+    # via
+    #   -c requirements.txt
+    #   aiohttp
 annotated-types==0.5.0
-    # via pydantic
+    # via
+    #   -c requirements.txt
+    #   pydantic
 argcomplete==2.0.0
     # via datamodel-code-generator
 asgiref==3.7.2
-    # via django
+    # via
+    #   -c requirements.txt
+    #   django
 async-timeout==4.0.2
     # via
+    #   -c requirements.txt
     #   aiohttp
     #   redis
 attrs==23.2.0
     # via
+    #   -c requirements.txt
     #   aiohttp
     #   jsonschema
     #   referencing
 black==23.9.1
-    # via datamodel-code-generator
+    # via
+    #   -c requirements.txt
+    #   -r requirements-dev.in
+    #   datamodel-code-generator
 boto3-stubs==1.34.84
+    # via -r requirements-dev.in
 botocore-stubs==1.34.84
     # via boto3-stubs
 certifi==2019.11.28
-    # via requests
+    # via
+    #   -c requirements.txt
+    #   requests
 cffi==1.14.5
-    # via cryptography
+    # via
+    #   -c requirements.txt
+    #   cryptography
 chardet==5.2.0
     # via prance
 charset-normalizer==2.1.0
-    # via requests
+    # via
+    #   -c requirements.txt
+    #   requests
 click==8.1.7
-    # via black
+    # via
+    #   -c requirements.txt
+    #   black
 colorama==0.4.4
     # via pytest-watch
 coverage==5.5
     # via pytest-cov
 cryptography==37.0.2
-    # via types-paramiko
+    # via
+    #   -c requirements.txt
+    #   types-paramiko
 datamodel-code-generator==0.25.6
+    # via -r requirements-dev.in
 django==4.2.11
     # via
+    #   -c requirements.txt
     #   django-stubs
     #   django-stubs-ext
 django-stubs==4.2.7
-    # via djangorestframework-stubs
+    # via
+    #   -r requirements-dev.in
+    #   djangorestframework-stubs
 django-stubs-ext==5.0.0
     # via django-stubs
 djangorestframework-stubs==3.14.5
+    # via -r requirements-dev.in
 dnspython==2.2.1
-    # via email-validator
+    # via
+    #   -c requirements.txt
+    #   email-validator
 docopt==0.6.2
     # via pytest-watch
 email-validator==2.0.0.post2
     # via pydantic
 exceptiongroup==1.2.1
-    # via pytest
+    # via
+    #   -c requirements.txt
+    #   pytest
+execnet==2.1.1
+    # via pytest-xdist
 faker==17.5.0
+    # via -r requirements-dev.in
 fakeredis==2.11.0
+    # via -r requirements-dev.in
 flaky==3.7.0
+    # via -r requirements-dev.in
 freezegun==1.2.2
+    # via -r requirements-dev.in
 frozenlist==1.3.0
     # via
+    #   -c requirements.txt
     #   aiohttp
     #   aiosignal
 genson==1.2.2
@@ -73,6 +114,7 @@ icdiff==2.0.5
     # via pytest-icdiff
 idna==2.8
     # via
+    #   -c requirements.txt
     #   email-validator
     #   requests
     #   yarl
@@ -86,12 +128,14 @@ jinja2==3.1.4
     # via datamodel-code-generator
 jsonschema==4.20.0
     # via
+    #   -c requirements.txt
     #   openapi-schema-validator
     #   openapi-spec-validator
 jsonschema-path==0.3.2
     # via openapi-spec-validator
 jsonschema-specifications==2023.12.1
     # via
+    #   -c requirements.txt
     #   jsonschema
     #   openapi-schema-validator
 lazy-object-proxy==1.10.0
@@ -102,45 +146,66 @@ markupsafe==2.1.5
     # via jinja2
 multidict==6.0.2
     # via
+    #   -c requirements.txt
     #   aiohttp
     #   yarl
 mypy==1.10.0
+    # via -r requirements-dev.in
 mypy-baseline==0.7.0
+    # via -r requirements-dev.in
 mypy-boto3-s3==1.34.65
     # via boto3-stubs
 mypy-extensions==1.0.0
     # via
+    #   -c requirements.txt
+    #   -r requirements-dev.in
     #   black
     #   mypy
 openapi-schema-validator==0.6.2
     # via openapi-spec-validator
 openapi-spec-validator==0.7.1
+    # via -r requirements-dev.in
 packaging==23.1
     # via
+    #   -c requirements.txt
+    #   -r requirements-dev.in
     #   black
     #   datamodel-code-generator
     #   prance
     #   pytest
 parameterized==0.9.0
+    # via -r requirements-dev.in
 pathable==0.4.3
     # via jsonschema-path
 pathspec==0.12.1
-    # via black
+    # via
+    #   -c requirements.txt
+    #   black
 platformdirs==3.11.0
-    # via black
+    # via
+    #   -c requirements.txt
+    #   black
 pluggy==0.13.1
     # via pytest
 pprintpp==0.4.0
     # via pytest-icdiff
 prance==23.6.21.0
+    # via -r requirements-dev.in
 pycparser==2.20
-    # via cffi
+    # via
+    #   -c requirements.txt
+    #   cffi
 pydantic==2.5.3
-    # via datamodel-code-generator
+    # via
+    #   -c requirements.txt
+    #   datamodel-code-generator
 pydantic-core==2.14.6
-    # via pydantic
+    # via
+    #   -c requirements.txt
+    #   pydantic
 pytest==7.4.4
     # via
+    #   -r requirements-dev.in
     #   pytest-asyncio
     #   pytest-cov
     #   pytest-django
@@ -149,42 +214,62 @@ pytest==7.4.4
     #   pytest-mock
     #   pytest-split
     #   pytest-watch
+    #   pytest-xdist
     #   syrupy
 pytest-asyncio==0.21.1
+    # via -r requirements-dev.in
 pytest-cov==4.1.0
+    # via -r requirements-dev.in
 pytest-django==4.5.2
+    # via -r requirements-dev.in
 pytest-env==0.8.2
+    # via -r requirements-dev.in
 pytest-icdiff==0.6
+    # via -r requirements-dev.in
 pytest-mock==3.11.1
+    # via -r requirements-dev.in
 pytest-split==0.8.1
+    # via -r requirements-dev.in
 pytest-watch==4.2.0
+    # via -r requirements-dev.in
+pytest-xdist==3.6.1
+    # via -r requirements-dev.in
 python-dateutil==2.8.2
     # via
+    #   -c requirements.txt
+    #   -r requirements-dev.in
     #   faker
     #   freezegun
 pyyaml==6.0.1
     # via
+    #   -c requirements.txt
     #   datamodel-code-generator
     #   jsonschema-path
     #   responses
 redis==4.5.4
-    # via fakeredis
+    # via
+    #   -c requirements.txt
+    #   fakeredis
 referencing==0.31.1
     # via
+    #   -c requirements.txt
     #   jsonschema
     #   jsonschema-path
     #   jsonschema-specifications
 requests==2.32.0
     # via
+    #   -c requirements.txt
     #   djangorestframework-stubs
     #   jsonschema-path
     #   prance
     #   responses
 responses==0.23.1
+    # via -r requirements-dev.in
 rfc3339-validator==0.1.4
     # via openapi-schema-validator
 rpds-py==0.16.2
     # via
+    #   -c requirements.txt
     #   jsonschema
     #   referencing
 ruamel-yaml==0.18.6
@@ -192,22 +277,30 @@ ruamel-yaml==0.18.6
 ruamel-yaml-clib==0.2.8
     # via ruamel-yaml
 ruff==0.4.3
+    # via -r requirements-dev.in
 six==1.16.0
     # via
+    #   -c requirements.txt
     #   prance
     #   python-dateutil
     #   rfc3339-validator
 sortedcontainers==2.4.0
-    # via fakeredis
+    # via
+    #   -c requirements.txt
+    #   fakeredis
 sqlparse==0.4.4
-    # via django
+    # via
+    #   -c requirements.txt
+    #   django
 syrupy==4.6.0
+    # via -r requirements-dev.in
 toml==0.10.1
     # via
     #   coverage
     #   datamodel-code-generator
 tomli==2.0.1
     # via
+    #   -c requirements.txt
     #   black
     #   django-stubs
     #   mypy
@@ -215,27 +308,39 @@ tomli==2.0.1
 types-awscrt==0.20.9
     # via botocore-stubs
 types-freezegun==1.1.10
+    # via -r requirements-dev.in
 types-markdown==3.3.9
+    # via -r requirements-dev.in
 types-paramiko==3.4.0.20240423
+    # via -r requirements-dev.in
 types-python-dateutil==2.8.3
+    # via -r requirements-dev.in
 types-pytz==2023.3.0.0
     # via
+    #   -r requirements-dev.in
     #   django-stubs
     #   types-tzlocal
 types-pyyaml==6.0.1
     # via
+    #   -r requirements-dev.in
     #   django-stubs
     #   djangorestframework-stubs
     #   responses
 types-redis==4.3.20
+    # via -r requirements-dev.in
 types-requests==2.26.1
-    # via djangorestframework-stubs
+    # via
+    #   -r requirements-dev.in
+    #   djangorestframework-stubs
 types-retry==0.9.9.4
+    # via -r requirements-dev.in
 types-s3transfer==0.10.1
     # via boto3-stubs
 types-tzlocal==5.1.0.1
+    # via -r requirements-dev.in
 typing-extensions==4.7.1
     # via
+    #   -c requirements.txt
     #   asgiref
     #   black
     #   boto3-stubs
@@ -248,9 +353,12 @@ typing-extensions==4.7.1
     #   pydantic-core
 urllib3==1.26.18
     # via
+    #   -c requirements.txt
     #   requests
     #   responses
 watchdog==2.1.8
     # via pytest-watch
 yarl==1.7.2
-    # via aiohttp
+    # via
+    #   -c requirements.txt
+    #   aiohttp