feat: migrate core data stores to MySQL with compatibility fixes
This commit is contained in:
@@ -17,6 +17,46 @@ _MYSQL_USER = os.getenv("MYSQL_USER", "root")
|
||||
_MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "")
|
||||
_MYSQL_DATABASE = os.getenv("MYSQL_DATABASE", "ai_cs")
|
||||
|
||||
|
||||
class _CompatResult:
|
||||
def __init__(self, rows=None, rowcount: int = 0, lastrowid: int = 0):
|
||||
self._rows = rows or []
|
||||
self.rowcount = rowcount
|
||||
self.lastrowid = lastrowid
|
||||
|
||||
def fetchall(self):
|
||||
return self._rows
|
||||
|
||||
def fetchone(self):
|
||||
return self._rows[0] if self._rows else None
|
||||
|
||||
|
||||
class _PyMySQLCompatConn:
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
if exc_type:
|
||||
try:
|
||||
self._conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
self._conn.close()
|
||||
|
||||
def execute(self, query: str, args=None):
|
||||
cur = self._conn.cursor()
|
||||
cur.execute(query, args or ())
|
||||
rows = cur.fetchall() if cur.description else []
|
||||
res = _CompatResult(rows=rows, rowcount=cur.rowcount, lastrowid=getattr(cur, "lastrowid", 0))
|
||||
cur.close()
|
||||
return res
|
||||
|
||||
def commit(self):
|
||||
self._conn.commit()
|
||||
|
||||
def _is_mysql() -> bool:
|
||||
return _DB_TYPE in ("mysql", "mariadb")
|
||||
|
||||
@@ -27,7 +67,7 @@ def _sql(query: str) -> str:
|
||||
def _get_conn() -> sqlite3.Connection:
|
||||
if _is_mysql():
|
||||
import pymysql
|
||||
return pymysql.connect(
|
||||
conn = pymysql.connect(
|
||||
host=_MYSQL_HOST,
|
||||
port=_MYSQL_PORT,
|
||||
user=_MYSQL_USER,
|
||||
@@ -37,6 +77,7 @@ def _get_conn() -> sqlite3.Connection:
|
||||
cursorclass=pymysql.cursors.DictCursor,
|
||||
autocommit=False,
|
||||
)
|
||||
return _PyMySQLCompatConn(conn)
|
||||
os.makedirs(os.path.dirname(_DB_PATH), exist_ok=True)
|
||||
conn = sqlite3.connect(_DB_PATH)
|
||||
conn.row_factory = sqlite3.Row
|
||||
|
||||
Reference in New Issue
Block a user