From 6238224e8ca99a40a6b0c89cac44eb3c12bb1090 Mon Sep 17 00:00:00 2001
From: sunxyw <31698606+sunxyw@users.noreply.github.com>
Date: Wed, 4 Jan 2023 16:44:26 +0800
Subject: [PATCH] add postgres support (#236)

---
 src/ZM/Framework.php                  |  2 +-
 src/ZM/Store/Database/DBPool.php      | 22 ++++++++++++++++++++
 src/ZM/Store/Database/DBWrapper.php   |  1 +
 src/ZM/Store/Database/PgSQLDriver.php | 30 +++++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 src/ZM/Store/Database/PgSQLDriver.php

diff --git a/src/ZM/Framework.php b/src/ZM/Framework.php
index 6e04c51a..35263e04 100644
--- a/src/ZM/Framework.php
+++ b/src/ZM/Framework.php
@@ -324,7 +324,7 @@ class Framework
             }
             $properties['db[' . $name . ']'] = match ($db['type']) {
                 'sqlite' => $db['type'] . '://' . $db['dbname'],
-                'mysql' => $db['type'] . '://' . $db['host'] . ':' . $db['port'] . '/' . $db['dbname'],
+                'mysql', 'pgsql' => $db['type'] . '://' . $db['host'] . ':' . $db['port'] . '/' . $db['dbname'],
                 default => '未知数据库类型',
             };
         }
diff --git a/src/ZM/Store/Database/DBPool.php b/src/ZM/Store/Database/DBPool.php
index 2da79b7f..9fbd0796 100644
--- a/src/ZM/Store/Database/DBPool.php
+++ b/src/ZM/Store/Database/DBPool.php
@@ -53,6 +53,19 @@ class DBPool
                 $args = [];
                 $connect_str = str_replace(array_keys($table), array_values($table), $connect_str);
                 break;
+            case 'postgres':
+            case 'pgsql':
+                $connect_str = 'pgsql:host={host};port={port};dbname={dbname};user={username};password={password}';
+                $table = [
+                    '{host}' => $config['host'],
+                    '{port}' => $config['port'],
+                    '{dbname}' => $config['dbname'],
+                    '{username}' => $config['username'],
+                    '{password}' => $config['password'],
+                ];
+                $connect_str = str_replace(array_keys($table), array_values($table), $connect_str);
+                $args = [];
+                break;
             default:
                 throw new DBException('type ' . $config['type'] . ' not supported yet');
         }
@@ -76,6 +89,15 @@ class DBPool
             case 'mysql':
                 // TODO: 编写验证 MySQL 连接有效性的功能
                 break;
+            case 'postgres':
+            case 'pgsql':
+                $pool = self::$pools[$name]->get();
+                $a = $pool->query('select version();')->fetchAll()[0][0] ?? '';
+                if (str_starts_with($a, 'PostgreSQL')) {
+                    logger()->debug('pgsql ' . $name . ' connected');
+                }
+                self::$pools[$name]->put($pool);
+                break;
         }
     }
 
diff --git a/src/ZM/Store/Database/DBWrapper.php b/src/ZM/Store/Database/DBWrapper.php
index 6fdf322f..5a1c636f 100644
--- a/src/ZM/Store/Database/DBWrapper.php
+++ b/src/ZM/Store/Database/DBWrapper.php
@@ -600,6 +600,7 @@ class DBWrapper
         return match ($type) {
             'mysql' => MySQLDriver::class,
             'sqlite' => SQLiteDriver::class,
+            'pgsql', 'postgres' => PgSQLDriver::class,
             default => throw new DBException('Unknown database type: ' . $type),
         };
     }
diff --git a/src/ZM/Store/Database/PgSQLDriver.php b/src/ZM/Store/Database/PgSQLDriver.php
new file mode 100644
index 00000000..0925bc50
--- /dev/null
+++ b/src/ZM/Store/Database/PgSQLDriver.php
@@ -0,0 +1,30 @@
+<?php
+
+declare(strict_types=1);
+
+namespace ZM\Store\Database;
+
+use Doctrine\DBAL\Driver\AbstractPostgreSQLDriver;
+use Doctrine\DBAL\Driver\Connection;
+
+class PgSQLDriver extends AbstractPostgreSQLDriver
+{
+    public function connect(array $params, $username = null, $password = null, array $driverOptions = [])
+    {
+        logger()->debug('Requiring new connection');
+        return new DBConnection($params);
+    }
+
+    public function getName()
+    {
+        return 'pdo_pgsql_pool';
+    }
+
+    public function getDatabase(Connection $conn)
+    {
+        if ($conn instanceof DBConnection) {
+            return config('database')[$conn->getPoolName()]['dbname'] ?? '';
+        }
+        return '';
+    }
+}