286 lines
9.5 KiB
Python
286 lines
9.5 KiB
Python
|
"""
|
|||
|
|
数据源连接器服务
|
|||
|
|
"""
|
|||
|
|
from typing import Dict, Any, List, Optional
|
|||
|
|
import logging
|
|||
|
|
import json
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DataSourceConnector:
|
|||
|
|
"""数据源连接器基类"""
|
|||
|
|
|
|||
|
|
def __init__(self, source_type: str, config: Dict[str, Any]):
|
|||
|
|
"""
|
|||
|
|
初始化数据源连接器
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
source_type: 数据源类型
|
|||
|
|
config: 连接配置
|
|||
|
|
"""
|
|||
|
|
self.source_type = source_type
|
|||
|
|
self.config = config
|
|||
|
|
|
|||
|
|
def test_connection(self) -> Dict[str, Any]:
|
|||
|
|
"""
|
|||
|
|
测试连接
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
连接测试结果
|
|||
|
|
"""
|
|||
|
|
raise NotImplementedError("子类必须实现test_connection方法")
|
|||
|
|
|
|||
|
|
def query(self, query_params: Dict[str, Any]) -> Any:
|
|||
|
|
"""
|
|||
|
|
查询数据
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
query_params: 查询参数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
查询结果
|
|||
|
|
"""
|
|||
|
|
raise NotImplementedError("子类必须实现query方法")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MySQLConnector(DataSourceConnector):
|
|||
|
|
"""MySQL连接器"""
|
|||
|
|
|
|||
|
|
def test_connection(self) -> Dict[str, Any]:
|
|||
|
|
try:
|
|||
|
|
import pymysql
|
|||
|
|
connection = pymysql.connect(
|
|||
|
|
host=self.config.get('host'),
|
|||
|
|
port=self.config.get('port', 3306),
|
|||
|
|
user=self.config.get('user'),
|
|||
|
|
password=self.config.get('password'),
|
|||
|
|
database=self.config.get('database'),
|
|||
|
|
connect_timeout=5
|
|||
|
|
)
|
|||
|
|
connection.close()
|
|||
|
|
return {"status": "success", "message": "连接成功"}
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"MySQL连接失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def query(self, query_params: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|||
|
|
try:
|
|||
|
|
import pymysql
|
|||
|
|
sql = query_params.get('sql')
|
|||
|
|
if not sql:
|
|||
|
|
raise ValueError("缺少SQL查询语句")
|
|||
|
|
|
|||
|
|
connection = pymysql.connect(
|
|||
|
|
host=self.config.get('host'),
|
|||
|
|
port=self.config.get('port', 3306),
|
|||
|
|
user=self.config.get('user'),
|
|||
|
|
password=self.config.get('password'),
|
|||
|
|
database=self.config.get('database')
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with connection.cursor(pymysql.cursors.DictCursor) as cursor:
|
|||
|
|
cursor.execute(sql)
|
|||
|
|
result = cursor.fetchall()
|
|||
|
|
return result
|
|||
|
|
finally:
|
|||
|
|
connection.close()
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"MySQL查询失败: {str(e)}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PostgreSQLConnector(DataSourceConnector):
|
|||
|
|
"""PostgreSQL连接器"""
|
|||
|
|
|
|||
|
|
def test_connection(self) -> Dict[str, Any]:
|
|||
|
|
try:
|
|||
|
|
import psycopg2
|
|||
|
|
connection = psycopg2.connect(
|
|||
|
|
host=self.config.get('host'),
|
|||
|
|
port=self.config.get('port', 5432),
|
|||
|
|
user=self.config.get('user'),
|
|||
|
|
password=self.config.get('password'),
|
|||
|
|
database=self.config.get('database'),
|
|||
|
|
connect_timeout=5
|
|||
|
|
)
|
|||
|
|
connection.close()
|
|||
|
|
return {"status": "success", "message": "连接成功"}
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"PostgreSQL连接失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def query(self, query_params: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|||
|
|
try:
|
|||
|
|
import psycopg2
|
|||
|
|
import psycopg2.extras
|
|||
|
|
sql = query_params.get('sql')
|
|||
|
|
if not sql:
|
|||
|
|
raise ValueError("缺少SQL查询语句")
|
|||
|
|
|
|||
|
|
connection = psycopg2.connect(
|
|||
|
|
host=self.config.get('host'),
|
|||
|
|
port=self.config.get('port', 5432),
|
|||
|
|
user=self.config.get('user'),
|
|||
|
|
password=self.config.get('password'),
|
|||
|
|
database=self.config.get('database')
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
|||
|
|
cursor.execute(sql)
|
|||
|
|
result = cursor.fetchall()
|
|||
|
|
return [dict(row) for row in result]
|
|||
|
|
finally:
|
|||
|
|
connection.close()
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"PostgreSQL查询失败: {str(e)}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class APIConnector(DataSourceConnector):
|
|||
|
|
"""API连接器"""
|
|||
|
|
|
|||
|
|
def test_connection(self) -> Dict[str, Any]:
|
|||
|
|
try:
|
|||
|
|
import httpx
|
|||
|
|
url = self.config.get('base_url')
|
|||
|
|
if not url:
|
|||
|
|
raise ValueError("缺少base_url配置")
|
|||
|
|
|
|||
|
|
headers = self.config.get('headers', {})
|
|||
|
|
timeout = self.config.get('timeout', 10)
|
|||
|
|
|
|||
|
|
response = httpx.get(url, headers=headers, timeout=timeout)
|
|||
|
|
response.raise_for_status()
|
|||
|
|
|
|||
|
|
return {"status": "success", "message": "连接成功", "status_code": response.status_code}
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"API连接失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def query(self, query_params: Dict[str, Any]) -> Any:
|
|||
|
|
try:
|
|||
|
|
import httpx
|
|||
|
|
method = query_params.get('method', 'GET').upper()
|
|||
|
|
endpoint = query_params.get('endpoint', '')
|
|||
|
|
params = query_params.get('params', {})
|
|||
|
|
data = query_params.get('data', {})
|
|||
|
|
headers = self.config.get('headers', {})
|
|||
|
|
timeout = self.config.get('timeout', 10)
|
|||
|
|
|
|||
|
|
base_url = self.config.get('base_url', '').rstrip('/')
|
|||
|
|
url = f"{base_url}/{endpoint.lstrip('/')}"
|
|||
|
|
|
|||
|
|
if method == 'GET':
|
|||
|
|
response = httpx.get(url, params=params, headers=headers, timeout=timeout)
|
|||
|
|
elif method == 'POST':
|
|||
|
|
response = httpx.post(url, json=data, headers=headers, timeout=timeout)
|
|||
|
|
elif method == 'PUT':
|
|||
|
|
response = httpx.put(url, json=data, headers=headers, timeout=timeout)
|
|||
|
|
elif method == 'DELETE':
|
|||
|
|
response = httpx.delete(url, headers=headers, timeout=timeout)
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"不支持的HTTP方法: {method}")
|
|||
|
|
|
|||
|
|
response.raise_for_status()
|
|||
|
|
return response.json() if response.content else {}
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"API查询失败: {str(e)}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class JSONFileConnector(DataSourceConnector):
|
|||
|
|
"""JSON文件连接器"""
|
|||
|
|
|
|||
|
|
def test_connection(self) -> Dict[str, Any]:
|
|||
|
|
try:
|
|||
|
|
import os
|
|||
|
|
file_path = self.config.get('file_path')
|
|||
|
|
if not file_path:
|
|||
|
|
raise ValueError("缺少file_path配置")
|
|||
|
|
|
|||
|
|
if not os.path.exists(file_path):
|
|||
|
|
raise FileNotFoundError(f"文件不存在: {file_path}")
|
|||
|
|
|
|||
|
|
return {"status": "success", "message": "文件存在"}
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"JSON文件连接失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def query(self, query_params: Dict[str, Any]) -> Any:
|
|||
|
|
try:
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
file_path = self.config.get('file_path')
|
|||
|
|
|
|||
|
|
if not os.path.exists(file_path):
|
|||
|
|
raise FileNotFoundError(f"文件不存在: {file_path}")
|
|||
|
|
|
|||
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|||
|
|
data = json.load(f)
|
|||
|
|
|
|||
|
|
# 支持简单的查询过滤
|
|||
|
|
filter_path = query_params.get('path')
|
|||
|
|
if filter_path:
|
|||
|
|
# 支持JSONPath风格的路径查询
|
|||
|
|
parts = filter_path.split('.')
|
|||
|
|
result = data
|
|||
|
|
for part in parts:
|
|||
|
|
if isinstance(result, dict):
|
|||
|
|
result = result.get(part)
|
|||
|
|
elif isinstance(result, list):
|
|||
|
|
try:
|
|||
|
|
index = int(part)
|
|||
|
|
result = result[index]
|
|||
|
|
except (ValueError, IndexError):
|
|||
|
|
return None
|
|||
|
|
else:
|
|||
|
|
return None
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
return data
|
|||
|
|
except Exception as e:
|
|||
|
|
raise Exception(f"JSON文件查询失败: {str(e)}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 连接器工厂
|
|||
|
|
_connector_classes = {
|
|||
|
|
'mysql': MySQLConnector,
|
|||
|
|
'postgresql': PostgreSQLConnector,
|
|||
|
|
'api': APIConnector,
|
|||
|
|
'json': JSONFileConnector,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def create_connector(source_type: str, config: Dict[str, Any]):
|
|||
|
|
"""
|
|||
|
|
创建数据源连接器
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
source_type: 数据源类型
|
|||
|
|
config: 连接配置
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
数据源连接器实例
|
|||
|
|
"""
|
|||
|
|
connector_class = _connector_classes.get(source_type)
|
|||
|
|
if not connector_class:
|
|||
|
|
raise ValueError(f"不支持的数据源类型: {source_type}")
|
|||
|
|
|
|||
|
|
return connector_class(source_type, config)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 为了兼容API，创建一个统一的DataSourceConnector包装类
|
|||
|
|
class DataSourceConnectorWrapper:
|
|||
|
|
"""统一的数据源连接器包装类（用于API调用）"""
|
|||
|
|
|
|||
|
|
def __init__(self, source_type: str, config: Dict[str, Any]):
|
|||
|
|
self.connector = create_connector(source_type, config)
|
|||
|
|
self.source_type = source_type
|
|||
|
|
self.config = config
|
|||
|
|
|
|||
|
|
def test_connection(self) -> Dict[str, Any]:
|
|||
|
|
return self.connector.test_connection()
|
|||
|
|
|
|||
|
|
def query(self, query_params: Dict[str, Any]) -> Any:
|
|||
|
|
return self.connector.query(query_params)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 导出时使用包装类，这样API可以统一使用DataSourceConnector
|
|||
|
|
# 但实际返回的是具体的连接器实现
|