From 5e46eef8da639f59c5c788727b92fa523ba1f8cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E9=93=AD=E6=98=95?= <l@hyperf.io>
Date: Fri, 26 Feb 2021 19:09:30 +0800
Subject: [PATCH] Allow custom driver which you can used to register your own
 driver (#3316)

* Allow custom driver which you can used to register your own driver

* Added test cases

* Update CHANGELOG-2.1.md
---
 CHANGELOG-2.1.md                              |  1 +
 .../src/Exception/InvalidDriverException.php  | 16 ++++
 src/db/src/Pool/PoolFactory.php               | 11 ++-
 src/db/tests/Cases/AbstractTestCase.php       |  3 +-
 src/db/tests/Cases/CustomDriverTest.php       | 89 +++++++++++++++++++
 5 files changed, 118 insertions(+), 2 deletions(-)
 create mode 100644 src/db/src/Exception/InvalidDriverException.php
 create mode 100644 src/db/tests/Cases/CustomDriverTest.php

diff --git a/CHANGELOG-2.1.md b/CHANGELOG-2.1.md
index 78f8ad33f..1c7cfc099 100644
--- a/CHANGELOG-2.1.md
+++ b/CHANGELOG-2.1.md
@@ -8,6 +8,7 @@
 
 - [#3310](https://github.com/hyperf/hyperf/pull/3310) Added `Blueprint::comment()` which you can set comment of table for migration.
 - [#3311](https://github.com/hyperf/hyperf/pull/3311) Added `RouteCollector::getRouteParser` which you can get `RouteParser` from `RouteCollector`.
+- [#3316](https://github.com/hyperf/hyperf/pull/3316) Allow custom driver which you can used to register your own driver for `hyperf/db`.
 
 ## Optimized
 
diff --git a/src/db/src/Exception/InvalidDriverException.php b/src/db/src/Exception/InvalidDriverException.php
new file mode 100644
index 000000000..cf1fbb014
--- /dev/null
+++ b/src/db/src/Exception/InvalidDriverException.php
@@ -0,0 +1,16 @@
+<?php
+
+declare(strict_types=1);
+/**
+ * This file is part of Hyperf.
+ *
+ * @link     https://www.hyperf.io
+ * @document https://hyperf.wiki
+ * @contact  group@hyperf.io
+ * @license  https://github.com/hyperf/hyperf/blob/master/LICENSE
+ */
+namespace Hyperf\DB\Exception;
+
+class InvalidDriverException extends RuntimeException
+{
+}
diff --git a/src/db/src/Pool/PoolFactory.php b/src/db/src/Pool/PoolFactory.php
index 92563403c..4b4875d79 100644
--- a/src/db/src/Pool/PoolFactory.php
+++ b/src/db/src/Pool/PoolFactory.php
@@ -13,6 +13,7 @@ namespace Hyperf\DB\Pool;
 
 use Hyperf\Contract\ConfigInterface;
 use Hyperf\DB\Exception\DriverNotFoundException;
+use Hyperf\DB\Exception\InvalidDriverException;
 use Psr\Container\ContainerInterface;
 
 class PoolFactory
@@ -42,7 +43,11 @@ class PoolFactory
         $driver = $config->get(sprintf('db.%s.driver', $name), 'pdo');
         $class = $this->getPoolName($driver);
 
-        return $this->pools[$name] = make($class, [$this->container, $name]);
+        $pool = make($class, [$this->container, $name]);
+        if (! $pool instanceof Pool) {
+            throw new InvalidDriverException(sprintf('Driver %s is not invalid.', $driver));
+        }
+        return $this->pools[$name] = $pool;
     }
 
     protected function getPoolName(string $driver)
@@ -54,6 +59,10 @@ class PoolFactory
                 return PDOPool::class;
         }
 
+        if (class_exists($driver)) {
+            return $driver;
+        }
+
         throw new DriverNotFoundException(sprintf('Driver %s is not found.', $driver));
     }
 }
diff --git a/src/db/tests/Cases/AbstractTestCase.php b/src/db/tests/Cases/AbstractTestCase.php
index 8b12cae94..4eca20021 100644
--- a/src/db/tests/Cases/AbstractTestCase.php
+++ b/src/db/tests/Cases/AbstractTestCase.php
@@ -43,6 +43,8 @@ abstract class AbstractTestCase extends TestCase
     protected function getContainer($options = [])
     {
         $container = Mockery::mock(Container::class);
+        ApplicationContext::setContainer($container);
+
         $container->shouldReceive('get')->with(ConfigInterface::class)->andReturn(new Config([
             'db' => [
                 'default' => [
@@ -86,7 +88,6 @@ abstract class AbstractTestCase extends TestCase
             return new DB($factory, $params['poolName']);
         });
         $container->shouldReceive('has')->with(StdoutLoggerInterface::class)->andReturn(false);
-        ApplicationContext::setContainer($container);
         return $container;
     }
 }
diff --git a/src/db/tests/Cases/CustomDriverTest.php b/src/db/tests/Cases/CustomDriverTest.php
new file mode 100644
index 000000000..ff55d4634
--- /dev/null
+++ b/src/db/tests/Cases/CustomDriverTest.php
@@ -0,0 +1,89 @@
+<?php
+
+declare(strict_types=1);
+/**
+ * This file is part of Hyperf.
+ *
+ * @link     https://www.hyperf.io
+ * @document https://hyperf.wiki
+ * @contact  group@hyperf.io
+ * @license  https://github.com/hyperf/hyperf/blob/master/LICENSE
+ */
+namespace HyperfTest\DB\Cases;
+
+use Hyperf\Config\Config;
+use Hyperf\Contract\ConfigInterface;
+use Hyperf\Contract\ConnectionInterface;
+use Hyperf\DB\Pool\Pool;
+use Hyperf\DB\Pool\PoolFactory;
+use Hyperf\Utils\ApplicationContext;
+use Psr\Container\ContainerInterface;
+
+/**
+ * @internal
+ * @coversNothing
+ */
+class CustomDriverTest extends AbstractTestCase
+{
+    public function testCustomDriver()
+    {
+        $container = \Mockery::mock(ContainerInterface::class);
+        ApplicationContext::setContainer($container);
+        $container->shouldReceive('get')->with(ConfigInterface::class)->andReturn(new Config([
+            'db' => [
+                'custom' => [
+                    'driver' => CustomPool::class,
+                    'host' => '127.0.0.1',
+                    'password' => '',
+                    'database' => 'hyperf',
+                    'pool' => [
+                        'max_connections' => 20,
+                    ],
+                    'options' => [],
+                ],
+            ],
+        ]));
+        $factory = new PoolFactory($container);
+
+        $pool = $factory->getPool('custom');
+
+        $this->assertInstanceOf(CustomPool::class, $pool);
+        $this->assertInstanceOf(CustomConnection::class, $pool->get());
+    }
+}
+
+class CustomPool extends Pool
+{
+    protected function createConnection(): ConnectionInterface
+    {
+        return new CustomConnection();
+    }
+}
+
+class CustomConnection implements ConnectionInterface
+{
+    public function getConnection()
+    {
+        // TODO: Implement getConnection() method.
+    }
+
+    public function reconnect(): bool
+    {
+        // TODO: Implement reconnect() method.
+    }
+
+    public function check(): bool
+    {
+        // TODO: Implement check() method.
+    }
+
+    public function close(): bool
+    {
+        // TODO: Implement close() method.
+    }
+
+    public function release(): void
+    {
+        // TODO: Implement release() method.
+    }
+}