From 10a461212308badd211963e5d5178a11bb9d3f2a Mon Sep 17 00:00:00 2001 From: Zachary Collins Date: Mon, 18 Sep 2023 11:38:13 -0700 Subject: [PATCH] ref(hc): MigrationTest case infers database connection from table_hint --- src/sentry/testutils/cases.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/src/sentry/testutils/cases.py b/src/sentry/testutils/cases.py index 2b0416193a9451..37ce3565040f66 100644 --- a/src/sentry/testutils/cases.py +++ b/src/sentry/testutils/cases.py @@ -24,8 +24,9 @@ from django.contrib.auth.models import AnonymousUser from django.core import signing from django.core.cache import cache -from django.db import DEFAULT_DB_ALIAS, connection, connections +from django.db import DEFAULT_DB_ALIAS, connection, connections, router from django.db.migrations.executor import MigrationExecutor +from django.db.migrations.loader import MigrationLoader from django.http import HttpRequest from django.test import TestCase as DjangoTestCase from django.test import TransactionTestCase as DjangoTransactionTestCase @@ -2242,7 +2243,25 @@ def migrate_to(self): @property def connection(self): - return "default" + # Infer connection from table_name + m = MigrationLoader(connections["default"]).get_migration(self.app, self.migrate_to[0][1]) + hinted_tables = { + table + for op in m.operations + for tables in op.hints.get("tables", []) + for table in tables + } + + hinted_migrations = set() + for table_name in hinted_tables: + for connection_name in connections: + if router.allow_migrate(connection_name, self.app, table_name=table_name): + hinted_migrations.add(connection_name) + + assert ( + len(hinted_migrations) == 1 + ), f"Could not determine migration test connection from tables hints Found {hinted_migrations}" + return next(iter(hinted_migrations)) def setUp(self): super().setUp()