Changing OrmCache getSchema() to take a connection instance

src/ORM/OrmCache.php

@@ -15,7 +15,9 @@
 namespace Cake\ORM;
 
 use Cake\Cache\Cache;
+use Cake\Datasource\ConnectionInterface;
 use Cake\Datasource\ConnectionManager;
+use RuntimeException;
 
 /**
  * ORM Cache.
@@ -31,14 +33,14 @@ class OrmCache
     /**
      * Schema
      *
-     * @var \Cake\Datasource\SchemaInterface
+     * @var \Cake\Database\Schema\Collection|\Cake\Database\Schema\CachedCollection
      */
     protected $_schema;
 
     /**
      * Constructor
      *
-     * @param string $connection Connection name
+     * @param string|\Cake\Datasource\ConnectionInterface $connection Connection name to get the schema for or a connection instance
      */
     public function __construct($connection)
     {
@@ -90,17 +92,31 @@ class OrmCache
     /**
      * Helper method to get the schema collection.
      *
-     * @param string $connection Connection name to get the schema for
-     * @return false|\Cake\Database\Schema\Collection|\Cake\Database\Schema\CachedCollection
+     * @param string|\Cake\Datasource\ConnectionInterface $connection Connection name to get the schema for or a connection instance
+     * @return \Cake\Database\Schema\Collection|\Cake\Database\Schema\CachedCollection
      */
     public function getSchema($connection)
     {
-        /* @var \Cake\Database\Connection $source */
-        $source = ConnectionManager::get($connection);
+        if (is_string($connection)) {
+            /* @var \Cake\Database\Connection $source */
+            $source = ConnectionManager::get($connection);
+        } elseif (!$connection instanceof ConnectionInterface) {
+            throw new RuntimeException(sprtinf(
+                'OrmCache::getSchema() expects `%s`, `%s` given.',
+                ConnectionInterface::class,
+                is_object($connection) ? get_class($connection) : gettype($connection)
+            ));
+        }
+
+        if ($connection instanceof ConnectionInterface) {
+            $source = $connection;
+            $connection = $source->configName();
+        }
+
         if (!method_exists($source, 'schemaCollection')) {
             throw new RuntimeException(sprintf(
-                'The "%s" connection is not compatible with orm caching, ' .
-                'as it does not implement a "schemaCollection()" method.',
+                'The "%s" connection is not compatible with schema ' .
+                'caching, as it does not implement a "schemaCollection()" method.',
                 $connection
             ));
         }

tests/TestCase/ORM/OrmCacheTest.php

@@ -16,6 +16,7 @@ namespace Cake\Test\TestCase\ORM;
 
 use Cake\Cache\Cache;
 use Cake\Cache\CacheEngine;
+use Cake\Database\Schema\CachedCollection;
 use Cake\Datasource\ConnectionManager;
 use Cake\ORM\OrmCache;
 use Cake\TestSuite\TestCase;
@@ -198,4 +199,19 @@ class OrmCacheTest extends TestCase
         $ormCache = new OrmCache('test');
         $ormCache->clear('articles');
     }
+
+    /**
+     * Tests getting a schema config from a connection instance
+     *
+     * @return void
+     */
+    public function testGetSchemaWithConnectionInstance()
+    {
+        $ds = ConnectionManager::get('test');
+
+        $ormCache = new OrmCache($ds);
+        $result = $ormCache->getSchema($ds);
+
+        $this->assertInstanceOf(CachedCollection::class, $result);
+    }
 }