Merge 2d413619b539d571f4095ec90724c18d37990923 into a81b7855de814be8c3e4fb1c9341e2c507f34b6b

src/Control/Director.php

@@ -87,7 +87,7 @@ class Director implements TemplateGlobalProvider
     /**
      * List of rules that must only use the primary database and not a replica
      */
-    private static array $must_use_primary_db_rules = [
+    private static array $rule_patterns_must_use_primary_db = [
         'dev',
         'Security',
     ];
@@ -306,8 +306,9 @@ class Director implements TemplateGlobalProvider
 
         // Check if primary database must be used based on request rules
         // Note this check must happend before the rules are processed as
-        // $shiftOnSuccess is true during $request->match() below
-        $primaryDbOnlyRules = Director::config()->uninherited('must_use_primary_db_rules');
+        // $shiftOnSuccess param is passed as true in `$request->match($pattern, true)` later on in
+        // this method, which modifies `$this->dirParts`, thus affecting `$request->match($rule)` directly below
+        $primaryDbOnlyRules = Director::config()->uninherited('rule_patterns_must_use_primary_db');
         foreach ($primaryDbOnlyRules as $rule) {
             if ($request->match($rule)) {
                 DB::setMustUsePrimary();

src/Core/CoreKernel.php

@@ -6,6 +6,7 @@ use SilverStripe\Control\HTTPResponse_Exception;
 use SilverStripe\Dev\Install\DatabaseAdapterRegistry;
 use SilverStripe\ORM\DB;
 use Exception;
+use InvalidArgumentException;
 use LogicException;
 use SilverStripe\Dev\Deprecation;
 use SilverStripe\ORM\Connect\NullDatabase;
@@ -15,6 +16,7 @@ use SilverStripe\ORM\Connect\NullDatabase;
  */
 class CoreKernel extends BaseKernel
 {
+
     protected bool $bootDatabase = true;
 
     /**
@@ -76,7 +78,7 @@ class CoreKernel extends BaseKernel
     }
 
     /**
-     * Load default database configuration from the $database and $databaseConfig globals
+     * Load database configuration from the $database and $databaseConfig globals
      */
     protected function bootDatabaseGlobals()
     {
@@ -87,9 +89,9 @@ class CoreKernel extends BaseKernel
         global $databaseConfig;
         global $database;
 
-        for ($i = 0; $i <= 99; $i++) {
+        for ($i = 0; $i <= DB::MAX_REPLICAS; $i++) {
             if ($i === 0) {
-                $key = 'default';
+                $key = DB::CONN_PRIMARY;
             } else {
                 $key = DB::getReplicaConfigKey($i);
                 if (!DB::hasConfig($key)) {
@@ -120,20 +122,20 @@ class CoreKernel extends BaseKernel
     }
 
     /**
-     * Load default database configuration from environment variables
+     * Load database configuration from environment variables
      */
     protected function bootDatabaseEnvVars()
     {
         if (!$this->bootDatabase) {
             return;
         }
-        // Set default database config
+        // Set primary database config
         $databaseConfig = $this->getDatabaseConfig();
         $databaseConfig['database'] = $this->getDatabaseName();
-        DB::setConfig($databaseConfig);
+        DB::setConfig($databaseConfig, DB::CONN_PRIMARY);
 
         // Set database replicas config
-        for ($i = 1; $i <= 99; $i++) {
+        for ($i = 1; $i <= DB::MAX_REPLICAS; $i++) {
             $envKey = $this->getReplicaEnvKey('SS_DATABASE_SERVER', $i);
             if (!Environment::hasEnv($envKey)) {
                 break;
@@ -156,15 +158,24 @@ class CoreKernel extends BaseKernel
 
     private function getDatabaseReplicaConfig(int $replica)
     {
+        if ($replica <= 0) {
+            throw new InvalidArgumentException('Replica number must be greater than 0');
+        }
         return $this->getSingleDataBaseConfig($replica);
     }
 
     /**
      * Convert a database key to a replica key
      * e.g. SS_DATABASE_SERVER -> SS_DATABASE_SERVER_REPLICA_01
+     *
+     * @param string $key - The key to look up in the environment
+     * @param int $replica - Replica number
      */
     private function getReplicaEnvKey(string $key, int $replica): string
     {
+        if ($replica <= 0) {
+            throw new InvalidArgumentException('Replica number must be greater than 0');
+        }
         // Do not allow replicas to define keys that could lead to unexpected behaviour if
         // they do not match the primary database configuration
         if (in_array($key, ['SS_DATABASE_CLASS', 'SS_DATABASE_NAME', 'SS_DATABASE_CHOOSE_NAME'])) {
@@ -177,10 +188,13 @@ class CoreKernel extends BaseKernel
     /**
      * Reads a single database configuration variable from the environment
      * For replica databases, it will first attempt to find replica-specific configuration
-     * before falling back to the default configuration.
+     * before falling back to the primary configuration.
      *
      * Replicate specific configuration has `_REPLICA_01` appended to the key
      * where 01 is the replica number.
+     *
+     * @param string $key - The key to look up in the environment
+     * @param int $replica - Replica number. Passing 0 will return the primary database configuration
      */
     private function getDatabaseConfigVariable(string $key, int $replica): string
     {
@@ -193,6 +207,9 @@ class CoreKernel extends BaseKernel
         return '';
     }
 
+    /**
+     * @param int $replica - Replica number. Passing 0 will return the primary database configuration
+     */
     private function getSingleDataBaseConfig(int $replica): array
     {
         $databaseConfig = [

src/Dev/SapphireTest.php

@@ -397,8 +397,7 @@ abstract class SapphireTest extends TestCase implements TestOnly
     public static function setUpBeforeClass(): void
     {
         // Disallow the use of DB replicas in tests
-        // This prevents replicas from being used if there are environment variables set to use them
-        DB::setCanUseReplicas(false);
+        DB::setMustUsePrimary();
 
         // Start tests
         static::start();

src/ORM/DB.php

@@ -15,6 +15,7 @@ use SilverStripe\ORM\Connect\DBConnector;
 use SilverStripe\ORM\Connect\DBSchemaManager;
 use SilverStripe\ORM\Connect\Query;
 use SilverStripe\ORM\Queries\SQLExpression;
+use SilverStripe\Core\CoreKernel;
 
 /**
  * Global database interface, complete with static methods.
@@ -22,6 +23,25 @@ use SilverStripe\ORM\Queries\SQLExpression;
  */
 class DB
 {
+    /**
+     * The 'default' connection is a dynamic connection that will use either a replica connection if one is
+     * available and not forced to use the 'primary' connection, or the 'primary' connection
+     *
+     * Prior to CMS 5.4 when replica connections were introduced, the 'default' connection
+     * was not dynamic and instead represented what the 'primary' connection is now
+     */
+    public const CONN_DEFAULT = 'default';
+
+    /**
+     * The 'primary' connection name, which is the main database connection and is used for all write
+     * operations and for read operations when the 'default' connection is forced to use the 'primary' connection
+     */
+    public const CONN_PRIMARY = 'primary';
+
+    /**
+     * The maximum number of replicas databases that can be configured
+     */
+    public const MAX_REPLICAS = 99;
 
     /**
      * This constant was added in SilverStripe 2.4 to indicate that SQL-queries
@@ -76,6 +96,28 @@ class DB
      */
     private static $connection_attempted = false;
 
+    /**
+     * Only use the primary database connection for the rest of the current request
+     *
+     * @internal
+     */
+    private static bool $mustUsePrimary = false;
+
+    /**
+     * Used by DB::withPrimary() to count the number of times it has been called
+     * Uses an int instead of a bool to allow for nested calls
+     *
+     * @internal
+     */
+    private static int $withPrimaryCount = 0;
+
+    /**
+     * The key of the replica config to use for this request
+     *
+     * @internal
+     */
+    private static string $replicaConfigKey = '';
+
     /**
      * Set the global database connection.
      * Pass an object that's a subclass of SS_Database.  This object will be used when {@link DB::query()}
@@ -87,7 +129,7 @@ class DB
      * be accessed through DB::get_conn($name).  This is useful when you have an application that
      * needs to connect to more than one database.
      */
-    public static function set_conn(Database $connection, $name = 'default')
+    public static function set_conn(Database $connection, $name = DB::CONN_DEFAULT)
     {
         DB::$connections[$name] = $connection;
     }
@@ -99,10 +141,10 @@ class DB
      * the default connection is returned.
      * @return Database
      */
-    public static function get_conn($name = 'default')
+    public static function get_conn($name = DB::CONN_DEFAULT)
     {
         // Allow default to connect to replica if configured
-        if ($name === 'default') {
+        if ($name === DB::CONN_DEFAULT) {
             $name = DB::getDefaultConnectionName();
         }
 
@@ -111,12 +153,6 @@ class DB
             return DB::$connections[$name];
         }
 
-        // Ensure that primary config is set, clone it from default config
-        if ($name === 'primary' && !static::hasConfig('primary')) {
-            $config = static::getConfig('default');
-            static::setConfig($config, 'primary');
-        }
-
         // lazy connect
         $config = static::getConfig($name);
         if ($config) {
@@ -127,45 +163,16 @@ class DB
     }
 
     /**
-     * Only use the primary database connection for the rest of the current request
-     *
-     * @internal
+     * Get whether the primary database connection must be used for the rest of the current request
+     * and not a replica connection
      */
-    private static bool $mustUsePrimary = false;
-
-    /**
-     * Used by DB::withPrimary() to count the number of times it has been called
-     * Uses an int instead of a bool to allow for nested calls     *
-     *
-     * @internal
-     */
-    private static int $withPrimaryCount = 0;
-
-    /**
-     * The key of the previously choosen random replica config
-     *
-     * @internal
-     */
-    private static string $randomReplicaConfigKey = '';
-
-    /**
-     * Flag to determine if replicas can be used
-     *
-     * @internal
-     */
-    private static bool $canUseReplicas = true;
-
-    /**
-     * Set the flag to allow or disallow the use of replica database connections
-     * For instance disable when running unit-tests
-     */
-    public static function setCanUseReplicas(bool $canUseReplicas): void
+    public static function getMustUsePrimary(): bool
     {
-        DB::$canUseReplicas = $canUseReplicas;
+        return DB::$mustUsePrimary;
     }
 
     /**
-     * Set the flag to only use the primary database connection for the current request
+     * Set to use the primary database connection for rest of the current request
      * meaning that replia connections will no longer be used
      *
      * This intentioally does not have a parameter to set this back to false, as this it to prevent
@@ -192,58 +199,6 @@ class DB
         return $result;
     }
 
-    /**
-     * Get the name of the database connection to use for the given SQL query
-     * The 'default' connection can be either the primary or a replica connection if configured
-     */
-    private static function getDefaultConnectionName(string $sql = ''): string
-    {
-        if (DB::$mustUsePrimary || DB::$withPrimaryCount > 0 || !DB::$canUseReplicas || !DB::hasReplicaConfig()) {
-            return 'primary';
-        }
-        if ($sql) {
-            $dbClass = DB::getConfig('default')['type'];
-            // This must use getServiceSpec() and not Injector::get/create() followed by
-            // getConnector() as this can remove the dbConn from a different connection
-            // under edge case conditions
-            $dbSpec = Injector::inst()->getServiceSpec($dbClass);
-            $connectorService = $dbSpec['properties']['connector'];
-            $connector = Injector::inst()->convertServiceProperty($connectorService);
-            if ($connector->isQueryMutable($sql)) {
-                DB::$mustUsePrimary = true;
-                return 'primary';
-            }
-        }
-        if (DB::$randomReplicaConfigKey) {
-            return DB::$randomReplicaConfigKey;
-        }
-        $name = DB::getRandomReplicaConfigKey();
-        DB::$randomReplicaConfigKey = $name;
-        return $name;
-    }
-
-    /**
-     * Get a random replica database configuration key from the available replica configurations
-     * The replica choosen will be used for the rest of the request, unless the primary connection
-     * is forced
-     */
-    private static function getRandomReplicaConfigKey(): string
-    {
-        $replicaNumbers = [];
-        for ($i = 1; $i < 99; $i++) {
-            $replicaKey = DB::getReplicaConfigKey($i);
-            if (DB::hasConfig($replicaKey)) {
-                $replicaNumbers[] = $i;
-            }
-        }
-        if (count($replicaNumbers) === 0) {
-            throw new RunTimeException('No replica configurations found');
-        }
-        // Choose a random replica
-        $index = rand(0, count($replicaNumbers) - 1);
-        return DB::getReplicaConfigKey($replicaNumbers[$index]);
-    }
-
     /**
      * Retrieves the schema manager for the current database
      *
@@ -251,7 +206,7 @@ class DB
      * the default connection is returned.
      * @return DBSchemaManager
      */
-    public static function get_schema($name = 'default')
+    public static function get_schema($name = DB::CONN_DEFAULT)
     {
         $connection = DB::get_conn($name);
         if ($connection) {
@@ -269,7 +224,7 @@ class DB
      * the default connection is returned.
      * @return string The resulting SQL as a string
      */
-    public static function build_sql(SQLExpression $expression, &$parameters, $name = 'default')
+    public static function build_sql(SQLExpression $expression, &$parameters, $name = DB::CONN_DEFAULT)
     {
         $connection = DB::get_conn($name);
         if ($connection) {
@@ -287,7 +242,7 @@ class DB
      * the default connection is returned.
      * @return DBConnector
      */
-    public static function get_connector($name = 'default')
+    public static function get_connector($name = DB::CONN_DEFAULT)
     {
         $connection = DB::get_conn($name);
         if ($connection) {
@@ -403,8 +358,13 @@ class DB
      * @param string $label identifier for the connection
      * @return Database
      */
-    public static function connect($databaseConfig, $label = 'default')
+    public static function connect($databaseConfig, $label = DB::CONN_DEFAULT)
     {
+        // Allow default to connect to replica if configured
+        if ($label === DB::CONN_DEFAULT) {
+            $label = DB::getDefaultConnectionName();
+        }
+
         // This is used by the "testsession" module to test up a test session using an alternative name
         if ($name = DB::get_alternative_database_name()) {
             $databaseConfig['database'] = $name;
@@ -434,8 +394,11 @@ class DB
      * @param array $databaseConfig
      * @param string $name
      */
-    public static function setConfig($databaseConfig, $name = 'default')
+    public static function setConfig($databaseConfig, $name = DB::CONN_DEFAULT)
     {
+        if ($name === DB::CONN_DEFAULT) {
+            $name = DB::CONN_PRIMARY;
+        }
         static::$configs[$name] = $databaseConfig;
     }
 
@@ -445,8 +408,11 @@ class DB
      * @param string $name
      * @return mixed
      */
-    public static function getConfig($name = 'default')
+    public static function getConfig($name = DB::CONN_DEFAULT)
     {
+        if ($name === DB::CONN_DEFAULT) {
+            $name = DB::CONN_PRIMARY;
+        }
         if (static::hasConfig($name)) {
             return static::$configs[$name];
         }
@@ -455,8 +421,11 @@ class DB
     /**
      * Check if a named connection config exists
      */
-    public static function hasConfig($name = 'default'): bool
+    public static function hasConfig($name = DB::CONN_DEFAULT): bool
     {
+        if ($name === DB::CONN_DEFAULT) {
+            $name = DB::CONN_PRIMARY;
+        }
         return array_key_exists($name, static::$configs);
     }
 
@@ -844,4 +813,63 @@ class DB
     {
         DB::get_schema()->alterationMessage($message, $type);
     }
+
+    /**
+     * Get the name of the database connection to use for the given SQL query
+     * The 'default' connection can be either the primary or a replica connection if configured
+     */
+    private static function getDefaultConnectionName(string $sql = ''): string
+    {
+        if (DB::$mustUsePrimary || DB::$withPrimaryCount > 0 || !DB::hasReplicaConfig()) {
+            return DB::CONN_PRIMARY;
+        }
+        if (DB::isMutableSql($sql)) {
+            DB::$mustUsePrimary = true;
+            return DB::CONN_PRIMARY;
+        }
+        if (DB::$replicaConfigKey) {
+            return DB::$replicaConfigKey;
+        }
+        $name = DB::getRandomReplicaConfigKey();
+        DB::$replicaConfigKey = $name;
+        return $name;
+    }
+
+    /**
+     * Check if the given SQL query is mutable
+     */
+    private static function isMutableSql(string $sql): bool
+    {
+        $dbClass = DB::getConfig(DB::CONN_PRIMARY)['type'];
+        // This must use getServiceSpec() and not Injector::get/create() followed by
+        // getConnector() as this can remove the dbConn from a different connection
+        // under edge case conditions
+        $dbSpec = Injector::inst()->getServiceSpec($dbClass);
+        $connectorService = $dbSpec['properties']['connector'];
+        $connector = Injector::inst()->convertServiceProperty($connectorService);
+        return $connector->isQueryMutable($sql);
+    }
+
+    /**
+     * Get a random replica database configuration key from the available replica configurations
+     * The replica choosen will be used for the rest of the request, unless the primary connection
+     * is forced
+     */
+    private static function getRandomReplicaConfigKey(): string
+    {
+        $replicaNumbers = [];
+        for ($i = 1; $i < DB::MAX_REPLICAS; $i++) {
+            $replicaKey = DB::getReplicaConfigKey($i);
+            if (!DB::hasConfig($replicaKey)) {
+                break;
+            }
+            $replicaNumbers[] = $i;
+        }
+        if (count($replicaNumbers) === 0) {
+            throw new RunTimeException('No replica configurations found');
+        }
+        // Choose a random replica
+        $index = rand(0, count($replicaNumbers) - 1);
+        return DB::getReplicaConfigKey($replicaNumbers[$index]);
+    }
 }

src/ORM/DataObject.php

@@ -142,10 +142,10 @@ class DataObject extends ViewableData implements DataObjectInterface, i18nEntity
     private static $default_classname = null;
 
     /**
-     * Whether this DataObject class must only use the primary database and not a replica
+     * Whether this DataObject class must only use the primary database and not a read-only replica
      * Note that this will be only be enforced when using DataQuery::execute() or
-     * another method that uses calls DataQuery::execute() internally e.g. as DataObject::get()
-     * This will not be enforced when using another method to query data e.g. SQLSelect or DB::query()
+     * another method that uses calls DataQuery::execute() internally e.g. DataObject::get()
+     * This will not be enforced when using low-level ORM functionality to query data e.g. SQLSelect or DB::query()
      */
     private static bool $must_use_primary_db = false;
 

src/ORM/DataQuery.php

@@ -443,18 +443,6 @@ class DataQuery
         }
     }
 
-    /**
-     * Calls a callback on either the primary database or a replica, with respect to the configured
-     * value of `must_use_primary_db` on the current dataClass
-     */
-    private function withCorrectDatabase(callable $callback): mixed
-    {
-        if (Config::inst()->get($this->dataClass(), 'must_use_primary_db')) {
-            return DB::withPrimary($callback);
-        }
-        return $callback();
-    }
-
     /**
      * Execute the query and return the result as {@link SS_Query} object.
      *
@@ -1522,4 +1510,16 @@ class DataQuery
 
         return $updated;
     }
+
+    /**
+     * Calls a callback on either the primary database or a replica, with respect to the configured
+     * value of `must_use_primary_db` on the current dataClass
+     */
+    private function withCorrectDatabase(callable $callback): mixed
+    {
+        if (Config::inst()->get($this->dataClass(), 'must_use_primary_db')) {
+            return DB::withPrimary($callback);
+        }
+        return $callback();
+    }
 }

tests/php/ORM/DBReplicaTest.php

@@ -13,14 +13,7 @@ use SilverStripe\ORM\Tests\DBReplicaTest\TestObject;
 use SilverStripe\Security\Group;
 use SilverStripe\Security\Member;
 use SilverStripe\ORM\DataQuery;
-use SilverStripe\Security\Permission;
 
-/**
- * Test for DB replica functionality
- *
- * Note that a lot of these methods test mutiple things at once. This is to increase the run speed of the tests
- * because they're all quite slow due to the need to create and tear-down fixtures for every test
- */
 class DBReplicaTest extends FunctionalTest
 {
     protected static $extra_dataobjects = [
@@ -33,9 +26,9 @@ class DBReplicaTest extends FunctionalTest
     {
         parent::setUp();
         $this->setupConfigsAndConnections(true);
-        // Allow using replicas - this is disabled by default in SapphireTest::setUpBeforeClass()
-        DB::setCanUseReplicas(true);
-        // Reset mustUsePrimary after using mutable sql to create yml fixtures
+        // Set DB:$mustUsePrimary to true to allow using replicas
+        // This is disabled by default in SapphireTest::setUpBeforeClass()
+        // Also reset mustUsePrimary after using mutable sql to create yml fixtures
         // and also because by default an ADMIN user is logged in when using fixtures in SapphireTest::setUp()
         // and also prevent tests from affecting subsequent tests
         (new ReflectionClass(DB::class))->setStaticPropertyValue('mustUsePrimary', false);
@@ -44,106 +37,70 @@ class DBReplicaTest extends FunctionalTest
     protected function tearDown(): void
     {
         $this->setupConfigsAndConnections(false);
-        DB::setCanUseReplicas(false);
+        // Reset DB:$mustUsePrimary to true which is the default set by SapphireTest::setUpBeforeClass()
+        (new ReflectionClass(DB::class))->setStaticPropertyValue('mustUsePrimary', true);
         parent::tearDown();
     }
-
-    /**
-     * Using reflection, set DB::configs and DB::connections with a fake a replica connection
-     * that points to the same connection as the primary connection.
-     */
-    private function setupConfigsAndConnections($includeReplica = true)
-    {
-        $reflector = new ReflectionClass(DB::class);
-        $defaultConfig = DB::getConfig('default');
-        $configs = ['default' => $defaultConfig];
-        if ($includeReplica) {
-            $configs['primary'] = $defaultConfig;
-            $configs['replica_01'] = $defaultConfig;
-        }
-        $reflector->setStaticPropertyValue('configs', $configs);
-        // Make a call to the database to ensure the connection is established
-        TestObject::get()->count();
-        // Create connections
-        $defaultConnection = DB::get_conn('default');
-        $connections = ['default' => $defaultConnection];
-        if ($includeReplica) {
-            $connections['primary'] = $defaultConnection;
-            $connections['replica_01'] = $defaultConnection;
-        }
-        $reflector->setStaticPropertyValue('connections', $connections);
-    }
-
-    /**
-     * Get the last connection name used by the DB class. This shows if a replica was used.
-     */
-    private function getLastConnectionName()
-    {
-        return (new ReflectionClass(DB::class))->getStaticPropertyValue('lastConnectionName');
-    }
     
-    public function testUsesReplica()
+    public function testUsesReplica(): void
     {
         // Assert uses replica by default
         TestObject::get()->count();
         $this->assertSame('replica_01', $this->getLastConnectionName());
         // Assert uses primary when using withPrimary()
         DB::withPrimary(fn() => TestObject::get()->count());
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
         // Assert that withPrimary() was only temporary
         TestObject::get()->count();
         $this->assertSame('replica_01', $this->getLastConnectionName());
         // Assert DB::setMustUsePrimary() forces primary from now on
         DB::setMustUsePrimary();
         TestObject::get()->count();
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
     }
 
-    public function testMutableSql()
+    public function testMutableSql(): void
     {
         // Assert that using mutable sql in an ORM method with a dataclass uses primary
         TestObject::create(['Title' => 'testing'])->write();
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
         // Assert that now all subsequent queries use primary
         TestObject::get()->count();
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
     }
 
-    public function testMutableSqlDbQuery()
+    public function testMutableSqlDbQuery(): void
     {
         // Assert that using mutable sql in DB::query() uses primary
         DB::query('INSERT INTO "DBReplicaTest_TestObject" ("Title") VALUES (\'testing\')');
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
         // Assert that now all subsequent queries use primary
         TestObject::get()->count();
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
     }
 
-    public function testMutableSqlDbPreparedQuery()
+    public function testMutableSqlDbPreparedQuery(): void
     {
         // Assert that using mutable sql in DB::prepared_query() uses primary
         DB::prepared_query('INSERT INTO "DBReplicaTest_TestObject" ("Title") VALUES (?)', ['testing']);
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
         // Assert that now all subsequent queries use primary
         TestObject::get()->count();
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
     }
 
-    public function testSetCurrentUser()
+    /**
+     * @dataProvider provideSetCurrentUser
+     */
+    public function testSetCurrentUser(string $firstName, string $expected): void
     {
-        // Assert non cms user logged in uses replica
-        $member = Member::get()->find('FirstName', 'random');
+        $member = Member::get()->find('FirstName', $firstName);
         Security::setCurrentUser($member);
         TestObject::get()->count();
-        $this->assertSame('replica_01', $this->getLastConnectionName());
-        // Assert cms user logged in uses primary
-        $member = Member::get()->find('FirstName', 'cmsuser');
-        Security::setCurrentUser($member);
-        TestObject::get()->count();
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame($expected, $this->getLastConnectionName());
     }
 
-    public function testDataObjectMustUsePrimaryDb()
+    public function testDataObjectMustUsePrimaryDb(): void
     {
         // Assert that DataList::getIterator() respect DataObject.must_use_primary_db
         foreach (TestObject::get() as $object) {
@@ -153,7 +110,7 @@ class DBReplicaTest extends FunctionalTest
         foreach (Group::get() as $group) {
             $group->Title = 'test2';
         }
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName());
         // Assert that DataQuery methods without params respect DataObject.must_use_primary_db
         $methods = [
             'count',
@@ -165,7 +122,7 @@ class DBReplicaTest extends FunctionalTest
             (new DataQuery(TestObject::class))->$method();
             $this->assertSame('replica_01', $this->getLastConnectionName(), "method is $method");
             (new DataQuery(Group::class))->$method();
-            $this->assertSame('primary', $this->getLastConnectionName(), "method is $method");
+            $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName(), "method is $method");
         }
         // Assert that DataQuery methods with a param respect DataObject.must_use_primary_db
         $methods = [
@@ -179,11 +136,50 @@ class DBReplicaTest extends FunctionalTest
             (new DataQuery(TestObject::class))->$method('ID');
             $this->assertSame('replica_01', $this->getLastConnectionName(), "method is $method");
             (new DataQuery(Group::class))->$method('ID');
-            $this->assertSame('primary', $this->getLastConnectionName(), "method is $method");
+            $this->assertSame(DB::CONN_PRIMARY, $this->getLastConnectionName(), "method is $method");
         }
     }
 
-    public function testRoutes()
+    public static function provideSetCurrentUser(): array
+    {
+        return [
+            'non_cms_user' => [
+                'firstName' => 'random',
+                'expected' => 'replica_01'
+            ],
+            'cms_user' => [
+                'firstName' => 'cmsuser',
+                'expected' => DB::CONN_PRIMARY
+            ],
+        ];
+    }
+
+    public static function provideRoutes(): array
+    {
+        return [
+            'normal_route' => [
+                'path' => 'test',
+                'expected' => 'replica_01'
+            ],
+            'security_route' => [
+                'path' => 'Security/login',
+                'expected' => DB::CONN_PRIMARY
+            ],
+            'dev_route' => [
+                'path' => 'dev/tasks',
+                'expected' => DB::CONN_PRIMARY
+            ],
+            'dev_in_path_but_not_dev_route' => [
+                'path' => 'test/dev',
+                'replica_01'
+            ],
+        ];
+    }
+
+    /**
+     * @dataProvider provideRoutes
+     */
+    public function testRoutes(string $path, string $expected): void
     {
         // Create a custom rule to test our controller that should default to using a replica
         $rules = Config::inst()->get(Director::class, 'rules');
@@ -191,40 +187,79 @@ class DBReplicaTest extends FunctionalTest
         // Ensure that routes staring with '$' are at the bottom of the assoc array index and don't override
         // our new 'test' route
         uksort($rules, fn($a, $b) => str_starts_with($a, '$') ? 1 : (str_starts_with($b, '$') ? -1 : 0));
-        // Assert that normal routes use replica
-        $this->get('/test');
-        $this->assertSame('replica_01', $this->getLastConnectionName());
-        // Assert that Security route is forced to use primary
-        $this->get('/Security/login');
-        $this->assertSame('primary', $this->getLastConnectionName());
-        // Assert that dev route is forced to use primary
-        $this->get('/dev/tasks');
-        $this->assertSame('primary', $this->getLastConnectionName());
+        $this->get($path);
+        $this->assertSame($expected, $this->getLastConnectionName());
     }
 
-    public function testSetCanUseReplicas()
+    public function provideHasReplicaConfig(): array
     {
-        DB::setCanUseReplicas(false);
-        TestObject::get()->count();
-        $this->assertSame('primary', $this->getLastConnectionName());
-        DB::setCanUseReplicas(true);
-        TestObject::get()->count();
-        $this->assertSame('replica_01', $this->getLastConnectionName());
+        return [
+            'no_replica' => [
+                'includeReplica' => false,
+                'expected' => false
+            ],
+            'with_replica' => [
+                'includeReplica' => true,
+                'expected' => true
+            ],
+        ];
     }
 
-    public function testConfigMethods()
+    /**
+     * @dataProvider provideHasReplicaConfig
+     */
+    public function testHasReplicaConfig(bool $includeReplica, bool $expected): void
+    {
+        $this->assertTrue(DB::hasReplicaConfig());
+        $primaryConfig = DB::getConfig(DB::CONN_PRIMARY);
+        $config = [DB::CONN_PRIMARY => $primaryConfig];
+        if ($includeReplica) {
+            $config['replica_01'] = $primaryConfig;
+        }
+        (new ReflectionClass(DB::class))->setStaticPropertyValue('configs', $config);
+        $this->assertSame($expected, DB::hasReplicaConfig());
+    }
+
+    public function testHasConfig(): void
     {
-        // Test DB::hasConfig()
         $this->assertFalse(DB::hasConfig('lorem'));
         DB::setConfig(['type' => 'lorem'], 'lorem');
         $this->assertTrue(DB::hasConfig('lorem'));
-        // Test DB::hasReplicaConfig()
-        $this->assertTrue(DB::hasReplicaConfig());
-        $default = DB::getConfig('default');
-        (new ReflectionClass(DB::class))->setStaticPropertyValue('configs', ['default' => $default]);
-        $this->assertFalse(DB::hasReplicaConfig());
-        // Test DB::getReplicaConfigKey()
+    }
+
+    public function testGetReplicaConfigKey(): void
+    {
         $this->assertSame('replica_03', DB::getReplicaConfigKey(3));
         $this->assertSame('replica_58', DB::getReplicaConfigKey(58));
     }
+
+    /**
+     * Using reflection, set DB::configs and DB::connections with a fake a replica connection
+     * that points to the same connection as the primary connection.
+     */
+    private function setupConfigsAndConnections($includeReplica = true): void
+    {
+        $reflector = new ReflectionClass(DB::class);
+        $primaryConfig = DB::getConfig(DB::CONN_PRIMARY);
+        $configs = [DB::CONN_PRIMARY => $primaryConfig];
+        if ($includeReplica) {
+            $configs['replica_01'] = $primaryConfig;
+        }
+        $reflector->setStaticPropertyValue('configs', $configs);
+        // Create connections
+        $primaryConnection = DB::get_conn(DB::CONN_PRIMARY);
+        $connections = [DB::CONN_PRIMARY => $primaryConnection];
+        if ($includeReplica) {
+            $connections['replica_01'] = $primaryConnection;
+        }
+        $reflector->setStaticPropertyValue('connections', $connections);
+    }
+
+    /**
+     * Get the last connection name used by the DB class. This shows if a replica was used.
+     */
+    private function getLastConnectionName(): string
+    {
+        return (new ReflectionClass(DB::class))->getStaticPropertyValue('lastConnectionName');
+    }
 }