|
|
@@ -1,14 +1,24 @@
|
|
|
-import pymysql
|
|
|
+import os
|
|
|
+from typing import Any, Dict, Optional, Tuple
|
|
|
|
|
|
|
|
|
-from typing import Tuple, Any, Dict, Literal, Optional
|
|
|
-from dbutils.pooled_db import PooledDB, PooledDedicatedDBConnection
|
|
|
-from dbutils.steady_db import SteadyDBCursor
|
|
|
-from pymysql.cursors import DictCursor
|
|
|
+def _database_config() -> dict[str, str]:
|
|
|
+ values = {
|
|
|
+ "host": os.environ.get("AGENT_BROWSER_DATABASE_HOST"),
|
|
|
+ "user": os.environ.get("AGENT_BROWSER_DATABASE_USER"),
|
|
|
+ "password": os.environ.get("AGENT_BROWSER_DATABASE_PASSWORD"),
|
|
|
+ "database": os.environ.get("AGENT_BROWSER_DATABASE_NAME"),
|
|
|
+ }
|
|
|
+ missing = [name for name, value in values.items() if not value]
|
|
|
+ if missing:
|
|
|
+ raise RuntimeError(
|
|
|
+ "legacy browser database configuration is missing: " + ", ".join(missing)
|
|
|
+ )
|
|
|
+ return {name: value for name, value in values.items() if value is not None}
|
|
|
|
|
|
|
|
|
class SyncMySQLHelper(object):
|
|
|
- _pool: PooledDB = None
|
|
|
+ _pool: Any = None
|
|
|
_instance = None
|
|
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
|
@@ -19,50 +29,69 @@ class SyncMySQLHelper(object):
|
|
|
|
|
|
def get_pool(self):
|
|
|
if self._pool is None:
|
|
|
+ required = _database_config()
|
|
|
+ import pymysql
|
|
|
+ from dbutils.pooled_db import PooledDB
|
|
|
+
|
|
|
self._pool = PooledDB(
|
|
|
creator=pymysql,
|
|
|
mincached=10,
|
|
|
maxconnections=20,
|
|
|
blocking=True,
|
|
|
- host='rm-t4na9qj85v7790tf84o.mysql.singapore.rds.aliyuncs.com',
|
|
|
+ host=required["host"],
|
|
|
port=3306,
|
|
|
- user='crawler_admin',
|
|
|
- password='cyber#crawler_2023',
|
|
|
- database='aigc-admin-prod')
|
|
|
+ user=required["user"],
|
|
|
+ password=required["password"],
|
|
|
+ database=required["database"],
|
|
|
+ )
|
|
|
|
|
|
return self._pool
|
|
|
|
|
|
- def fetchone(self, sql: str, data: Optional[Tuple[Any, ...]] = None) -> Dict[str, Any]:
|
|
|
+ @staticmethod
|
|
|
+ def _dict_cursor() -> Any:
|
|
|
+ from pymysql.cursors import DictCursor
|
|
|
+
|
|
|
+ return DictCursor
|
|
|
+
|
|
|
+ def fetchone(
|
|
|
+ self, sql: str, data: Optional[Tuple[Any, ...]] = None
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
pool = self.get_pool()
|
|
|
- with pool.connection() as conn:
|
|
|
- with conn.cursor(DictCursor) as cursor:
|
|
|
+ with pool.connection() as conn:
|
|
|
+ with conn.cursor(self._dict_cursor()) as cursor:
|
|
|
cursor.execute(sql, data)
|
|
|
result = cursor.fetchone()
|
|
|
return result
|
|
|
|
|
|
- def fetchall(self, sql: str, data: Optional[Tuple[Any, ...]] = None) -> Tuple[Dict[str, Any]]:
|
|
|
+ def fetchall(
|
|
|
+ self, sql: str, data: Optional[Tuple[Any, ...]] = None
|
|
|
+ ) -> Tuple[Dict[str, Any]]:
|
|
|
pool = self.get_pool()
|
|
|
- with pool.connection() as conn:
|
|
|
- with conn.cursor(DictCursor) as cursor:
|
|
|
+ with pool.connection() as conn:
|
|
|
+ with conn.cursor(self._dict_cursor()) as cursor:
|
|
|
cursor.execute(sql, data)
|
|
|
result = cursor.fetchall()
|
|
|
return result
|
|
|
|
|
|
- def fetchmany(self,
|
|
|
- sql: str,
|
|
|
- data: Optional[Tuple[Any, ...]] = None,
|
|
|
- size: Optional[int] = None) -> Tuple[Dict[str, Any]]:
|
|
|
+ def fetchmany(
|
|
|
+ self,
|
|
|
+ sql: str,
|
|
|
+ data: Optional[Tuple[Any, ...]] = None,
|
|
|
+ size: Optional[int] = None,
|
|
|
+ ) -> Tuple[Dict[str, Any]]:
|
|
|
pool = self.get_pool()
|
|
|
- with pool.connection() as conn:
|
|
|
- with conn.cursor(DictCursor) as cursor:
|
|
|
+ with pool.connection() as conn:
|
|
|
+ with conn.cursor(self._dict_cursor()) as cursor:
|
|
|
cursor.execute(sql, data)
|
|
|
result = cursor.fetchmany(size=size)
|
|
|
return result
|
|
|
|
|
|
def execute(self, sql: str, data: Optional[Tuple[Any, ...]] = None):
|
|
|
+ import pymysql
|
|
|
+
|
|
|
pool = self.get_pool()
|
|
|
- with pool.connection() as conn:
|
|
|
- with conn.cursor(DictCursor) as cursor:
|
|
|
+ with pool.connection() as conn:
|
|
|
+ with conn.cursor(self._dict_cursor()) as cursor:
|
|
|
try:
|
|
|
cursor.execute(sql, data)
|
|
|
result = conn.commit()
|
|
|
@@ -81,6 +110,3 @@ class SyncMySQLHelper(object):
|
|
|
|
|
|
|
|
|
mysql = SyncMySQLHelper()
|
|
|
-
|
|
|
-
|
|
|
-
|