them tool lay cau truc DB

parent e2e6148d
...@@ -364,6 +364,54 @@ def validate_table_name(table_name: str) -> bool: ...@@ -364,6 +364,54 @@ def validate_table_name(table_name: str) -> bool:
"""BẢO MẬT: Chống SQL Injection qua tên bảng. Chỉ cho phép chữ cái, số và gạch dưới.""" """BẢO MẬT: Chống SQL Injection qua tên bảng. Chỉ cho phép chữ cái, số và gạch dưới."""
return bool(re.match(r'^[a-zA-Z0-9_]+$', table_name)) return bool(re.match(r'^[a-zA-Z0-9_]+$', table_name))
@mcp.tool()
def list_postgres_tables(task_progress: str = "") -> str:
"""Lấy danh sách tất cả các bảng (tables) hiện có trong database PostgreSQL. BẮT BUỘC KHÔNG truyền tham số nào vào tool này."""
if pg_conn is None: return "Lỗi: PostgreSQL chưa kết nối."
try:
sql = """
SELECT table_name
FROM information_schema.tables
WHERE table_schema = 'public' AND table_type = 'BASE TABLE';
"""
with pg_conn.cursor(cursor_factory=RealDictCursor) as cursor:
cursor.execute(sql)
results = cursor.fetchall()
if not results:
return "Database hiện tại chưa có bảng nào."
return json.dumps([row['table_name'] for row in results], indent=2, ensure_ascii=False)
except Exception as e:
pg_conn.rollback()
return f"Lỗi lấy danh sách bảng: {str(e)}"
@mcp.tool()
def get_postgres_schema(table_name: str) -> str:
"""Lấy cấu trúc các cột của một bảng trong PostgreSQL. BẮT BUỘC gọi tool này nếu bạn không chắc chắn về tên các cột trong bảng."""
if pg_conn is None: return "Lỗi: PostgreSQL chưa kết nối."
if not validate_table_name(table_name): return "Lỗi: Tên bảng không hợp lệ."
try:
# Truy vấn vào bảng hệ thống của Postgres để lấy danh sách cột
sql = """
SELECT column_name, data_type
FROM information_schema.columns
WHERE table_name = %s;
"""
with pg_conn.cursor(cursor_factory=RealDictCursor) as cursor:
cursor.execute(sql, (table_name,))
results = cursor.fetchall()
if not results:
return f"Bảng '{table_name}' không tồn tại hoặc không có dữ liệu."
return json.dumps(results, indent=2, ensure_ascii=False)
except Exception as e:
pg_conn.rollback()
return f"Lỗi lấy cấu trúc bảng: {str(e)}"
@mcp.tool() @mcp.tool()
def select_postgres_records(project_id: str, table_name: str, query: str = "{}", limit: int = 5) -> str: def select_postgres_records(project_id: str, table_name: str, query: str = "{}", limit: int = 5) -> str:
"""Lấy dữ liệu từ PostgreSQL. Tự động filter theo project_id và bỏ qua is_deleted.""" """Lấy dữ liệu từ PostgreSQL. Tự động filter theo project_id và bỏ qua is_deleted."""
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment