From 882862534637736f352be7ce28b7fbdddbc40016 Mon Sep 17 00:00:00 2001
From: caoqianming <caoqianming@foxmail.com>
Date: Tue, 27 Aug 2024 17:26:05 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20base=20dbconnection=E6=94=AF=E6=8C=81pg?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/utils/sql.py | 37 +++++++++++++++++++++++++------------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/apps/utils/sql.py b/apps/utils/sql.py
index c919ab80..177be62d 100644
--- a/apps/utils/sql.py
+++ b/apps/utils/sql.py
@@ -53,26 +53,39 @@ def query_one_dict(sql, params=None):
         return dict(zip(columns, row))
     
 import pymysql
+import psycopg2
 
 class DbConnection:
-    def __init__(self, host, user, password, database):
+    def __init__(self, host, user, password, database, dbtype='mysql'):
+        if dbtype not in ['mysql', 'pg']:
+            raise ValueError('dbtype must be mysql or pg')
+        self.dbtype = dbtype
         self.host = host
         self.user = user
         self.password = password
         self.database = database
         self.conn = None
-        self.cursor = None
+
+    def _connect(self):
+        if self.dbtype == 'mysql':
+            return pymysql.connect(
+                host=self.host,
+                user=self.user,
+                password=self.password,
+                database=self.database
+            )
+        elif self.dbtype == 'pg':
+            return psycopg2.connect(
+                host=self.host,
+                user=self.user,
+                password=self.password,
+                database=self.database,
+            )
 
     def __enter__(self):
-        self.conn = pymysql.connect(
-            host=self.host,
-            user=self.user,
-            password=self.password,
-            database=self.database
-        )
-        self.cursor = self.conn.cursor()
-        return self.cursor
+        self.conn = self._connect()
+        return self.conn.cursor()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.cursor.close()
-        self.conn.close()
\ No newline at end of file
+        if self.conn:
+            self.conn.close()