zcdsj
/
factory
5
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: base dbconnection支持pg

This commit is contained in:
caoqianming 2024-08-27 17:26:05 +08:00
parent 1b19c02754
commit 8828625346
1 changed files with 25 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
apps/utils/sql.py
View File

@ -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()
if self.conn:
self.conn.close()