Django测试中临时切换至生产数据库获取数据的实现方案

核心思路 #

利用Django的多数据库配置+自定义上下文管理器，实现测试过程中临时切换到生产数据库只读查询，完成后自动切回测试库。

步骤1：配置生产数据库（只读权限） #

先在settings.py中添加生产数据库的配置，必须设置只读限制，从根源避免误写生产数据：

DATABASES = {
    'default': {
        # 你的测试数据库配置（如SQLite/测试PostgreSQL）
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': BASE_DIR / 'db_test.sqlite3',
    },
    'production': {
        'ENGINE': 'django.db.backends.postgresql', # 替换为你的生产库引擎
        'NAME': os.getenv('PROD_DB_NAME'),
        'USER': os.getenv('PROD_DB_READ_USER'),
        'PASSWORD': os.getenv('PROD_DB_READ_PWD'),
        'HOST': os.getenv('PROD_DB_HOST'),
        'PORT': os.getenv('PROD_DB_PORT'),
        # 强制只读模式（PostgreSQL示例，其他数据库需调整）
        'OPTIONS': {
            'options': '-c default_transaction_read_only=on'
        }
    }
}

注意：生产库的访问凭证用环境变量管理，不要硬编码到代码里。

步骤2：实现上下文管理器 #

在测试基类中编写activate_production_database方法，作为上下文管理器自动处理数据库切换与恢复：

from django.db import connections, router
from django.test import TestCase

class BaseReadOnlyProductionTestCase(TestCase):
    def activate_production_database(self):
        # 保存当前的数据库路由规则
        original_read_router = router.db_for_read
        
        # 定义临时路由：所有读操作走生产库，禁止写操作
        class ProductionReadOnlyRouter:
            def db_for_read(self, model, **hints):
                return 'production'
            
            def db_for_write(self, model, **hints):
                raise ValueError("禁止向生产数据库写入数据")

        try:
            # 替换路由规则
            router.db_for_read = ProductionReadOnlyRouter().db_for_read
            # 确保生产库连接已建立
            connections['production'].ensure_connection()
            yield
        finally:
            # 恢复原路由规则
            router.db_for_read = original_read_router
            # 关闭生产库连接，释放资源
            connections['production'].close()

步骤3：在测试中使用 #

直接在setUp或测试方法中用with语句调用上下文管理器：

class MetadataTest(BaseReadOnlyProductionTestCase):
    # 可选：类级别缓存生产数据，避免每个测试方法重复读取
    _cached_metadata_titles = None

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        with cls().activate_production_database():
            cls._cached_metadata_titles = list(
                MetadataTable.objects.values_list('title', flat=True)
            )

    def setUp(self):
        super().setUp()
        # 从缓存读取数据，批量写入测试库
        MetadataTable.objects.bulk_create([
            MetadataTable(title=t) for t in self._cached_metadata_titles
        ])

    def test_metadata_exists(self):
        self.assertEqual(MetadataTable.objects.count(), len(self._cached_metadata_titles))

关键注意事项 #

• 权限加固：生产库的测试用户必须只有SELECT权限，即使代码层面限制了写操作，也要从数据库权限层再做防护。
• 性能优化：用类级别缓存（setUpClass）代替每个测试方法读取生产库，减少对生产环境的访问压力。
• 数据一致性：如果生产数据更新频繁，可根据测试需求决定是否每次运行测试都刷新缓存，或固定读取某一快照。

内容的提问来源于stack exchange，提问作者Michael Scott Asato Cuthbert