[WIP] feat(metrics): PoC of running metrics indexer migrations against spanner

requirements-base.txt

@@ -102,4 +102,4 @@ phabricator>=0.7.0
 # test dependencies, but unable to move to requirements-test until
 # sentry.utils.pytest and sentry.testutils are moved to tests/
 selenium>=4.1.5
-sqlparse>=0.2.4
+sqlparse>=0.3.0

requirements-dev-frozen.txt

@@ -151,7 +151,7 @@ sniffio==1.2.0
 snuba-sdk==1.0.0
 sortedcontainers==2.4.0
 soupsieve==2.3.2.post1
-sqlparse==0.2.4
+sqlparse==0.4.2
 statsd==3.3
 structlog==21.1.0
 symbolic==8.7.1

requirements-frozen.txt

@@ -104,7 +104,7 @@ sniffio==1.2.0
 snuba-sdk==1.0.0
 sortedcontainers==2.4.0
 soupsieve==2.3.2.post1
-sqlparse==0.2.4
+sqlparse==0.4.2
 statsd==3.3
 structlog==21.1.0
 symbolic==8.7.1

src/sentry/conf/server.py

@@ -145,9 +145,17 @@ def env(key, default="", type=None):
         "PORT": "",
         "AUTOCOMMIT": True,
         "ATOMIC_REQUESTS": False,
-    }
+    },
+    "metrics-spanner": {
+        "ENGINE": "django_spanner",
+        "PROJECT": "search-and-storage",
+        "INSTANCE": "markus-test-spanner-pg",
+        "NAME": "metrics-spanner",
+    },
 }
 
+DATABASE_ROUTERS = ("sentry.db.router.MultiDatabaseRouter",)
+
 if "DATABASE_URL" in os.environ:
     url = urlparse(os.environ["DATABASE_URL"])
 
@@ -321,6 +329,7 @@ def env(key, default="", type=None):
 ]
 
 INSTALLED_APPS = (
+    "django_spanner",
     "django.contrib.admin",
     "django.contrib.auth",
     "django.contrib.contenttypes",

src/sentry/db/models/fields/bounded.py

@@ -2,7 +2,6 @@
 
 from django.conf import settings
 from django.db import models
-from django.db.backends.base.base import BaseDatabaseWrapper
 from django.utils.translation import ugettext_lazy as _
 
 __all__ = (
@@ -60,19 +59,13 @@ def get_prep_value(self, value: int) -> int:
                 assert value <= self.MAX_VALUE
             return cast(int, super().get_prep_value(value))
 
-    class BoundedBigAutoField(models.AutoField):  # type: ignore
+    class BoundedBigAutoField(models.BigAutoField):  # type: ignore
         description = _("Big Integer")
 
         MAX_VALUE = 9223372036854775807
 
-        def db_type(self, connection: BaseDatabaseWrapper) -> str:
-            return "bigserial"
-
-        def get_related_db_type(self, connection: BaseDatabaseWrapper) -> str:
-            return cast(str, BoundedBigIntegerField().db_type(connection))
-
         def get_internal_type(self) -> str:
-            return "BigIntegerField"
+            return "BigAutoField"
 
         def get_prep_value(self, value: int) -> int:
             if value:

src/sentry/db/router.py

@@ -0,0 +1,61 @@
+_db_table_to_db = {
+    "sentry_stringindexer": "metrics-spanner",  # will replace ^ table above
+    "sentry_perfstringindexer": "metrics-spanner",  # will replace ^ table above
+}
+
+
+def db_for_model(model):
+    return db_for_table(model._meta.db_table)
+
+
+def db_for_table(table):
+    return _db_table_to_db.get(table, "default")
+
+
+class MultiDatabaseRouter:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._setup_replicas()
+
+    def _setup_replicas(self):
+        from django.conf import settings
+
+        self._replicas = {}
+        for db, db_conf in settings.DATABASES.items():
+            primary = db_conf.get("REPLICA_OF")
+            if primary:
+                self._replicas[primary] = db
+
+    def db_for_read(self, model, **hints):
+        db = db_for_model(model)
+        if hints.get("replica", False):
+            db = self._replicas.get(db, db)
+        return db
+
+    def db_for_write(self, model, **hints):
+        return db_for_model(model)
+
+    def allow_relation(self, obj1, obj2, **hints):
+        return db_for_model(obj1) == db_for_model(obj2)
+
+    def allow_syncdb(self, db, model):
+        if db_for_model(model) == db:
+            return True
+        return False
+
+    def allow_migrate(self, db, app_label, model=None, **hints):
+        if model:
+            return db_for_model(model) == db
+
+        # We use this hint in our RunSql/RunPython migrations to help resolve databases.
+        if "tables" in hints:
+            dbs = {db_for_table(table) for table in hints["tables"]}
+            if len(dbs) > 1:
+                raise RuntimeError(
+                    "Migration tables resolve to multiple databases. "
+                    f"Got {dbs} when only one database should be used."
+                )
+            return dbs.pop() == db
+        # Assume migrations with no model routing or hints need to run on
+        # the default database.
+        return db == "default"

src/sentry/migrations/0284_metrics_indexer_alter_seq.py

@@ -27,16 +27,4 @@ class Migration(CheckedMigration):
         ("sentry", "0283_extend_externalissue_key"),
     ]
 
-    operations = [
-        migrations.RunSQL(
-            """
-            ALTER SEQUENCE sentry_stringindexer_id_seq START WITH 65536;
-            ALTER SEQUENCE sentry_stringindexer_id_seq RESTART;
-            """,
-            hints={"tables": ["sentry_stringindexer"]},
-            reverse_sql="""
-            ALTER SEQUENCE sentry_stringindexer_id_seq START WITH 1;
-            ALTER SEQUENCE sentry_stringindexer_id_seq RESTART;
-            """,
-        )
-    ]
+    operations = []

src/sentry/migrations/0291_add_new_perf_indexer.py

@@ -58,15 +58,4 @@ class Migration(CheckedMigration):
                 fields=("string", "organization_id"), name="perf_unique_org_string"
             ),
         ),
-        migrations.RunSQL(
-            """
-            ALTER SEQUENCE sentry_perfstringindexer_id_seq START WITH 65536;
-            ALTER SEQUENCE sentry_perfstringindexer_id_seq RESTART;
-            """,
-            hints={"tables": ["sentry_perfstringindexer"]},
-            reverse_sql="""
-            ALTER SEQUENCE sentry_perfstringindexer_id_seq START WITH 1;
-            ALTER SEQUENCE sentry_perfstringindexer_id_seq RESTART;
-             """,
-        ),
     ]

src/sentry/models/counter.py

@@ -122,6 +122,3 @@ def create_counter_function(app_config, using, **kwargs):
         )
     finally:
         cursor.close()
-
-
-post_migrate.connect(create_counter_function, dispatch_uid="create_counter_function", weak=False)