import os
|
import psycopg2
|
from psycopg2 import OperationalError
|
from dotenv import load_dotenv
|
|
# 加载环境变量
|
load_dotenv()
|
|
class DatabaseConnection:
|
def __init__(self):
|
self.host = os.getenv('DB_HOST')
|
self.port = os.getenv('DB_PORT')
|
self.database = os.getenv('DB_NAME')
|
self.user = os.getenv('DB_USER')
|
self.password = os.getenv('DB_PASSWORD')
|
self.connection = None
|
|
def connect(self):
|
try:
|
self.connection = psycopg2.connect(
|
host=self.host,
|
port=self.port,
|
database=self.database,
|
user=self.user,
|
password=self.password
|
)
|
|
|
with self.connection.cursor() as cursor:
|
cursor.execute("SET timezone = 'UTC'")
|
return True
|
except OperationalError as e:
|
print(f"数据库连接失败: {e}")
|
return False
|
|
def disconnect(self):
|
if self.connection:
|
self.connection.close()
|
self.connection = None
|
|
def get_connection(self):
|
if not self.connection:
|
if not self.connect():
|
return None
|
return self.connection
|
|
def is_connected(self):
|
if not self.connection:
|
return False
|
try:
|
# 执行一个简单的查询来测试连接
|
with self.connection.cursor() as cursor:
|
cursor.execute("SELECT 1")
|
cursor.fetchone()
|
return True
|
except:
|
return False
|
