第一次提交
This commit is contained in:
285
backend/app/services/data_source_connector.py
Normal file
285
backend/app/services/data_source_connector.py
Normal file
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
数据源连接器服务
|
||||
"""
|
||||
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
|
||||
# 但实际返回的是具体的连接器实现
|
||||
Reference in New Issue
Block a user