them tool get_postgres_foreign_keys,get_postgres_triggers, get_postgres_function_code

memory_mcp_server_stdio.py

@@ -357,7 +357,7 @@ def delete_mongo_document(project_id: str, collection_name: str, document_id: st
     except Exception as e:
         return f"Lỗi khi xóa: {str(e)}"
 # ==========================================
-# CÁC TOOL THAO TÁC VỚI POSTGRESQL (AN TOÀN)
+# CÁC TOOL THAO TÁC VỚI POSTGRESQL (TỔNG QUÁT & BẢO MẬT DỮ LIỆU)
 # ==========================================
 
 def validate_table_name(table_name: str) -> bool:
@@ -366,22 +366,14 @@ def validate_table_name(table_name: str) -> bool:
 
 @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."""
+    """Lấy danh sách tất cả các bảng hiện có trong PostgreSQL."""
     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';
-        """
+        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."
-                
+            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()
@@ -389,50 +381,35 @@ def list_postgres_tables(task_progress: str = "") -> str:
 
 @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."""
+    """Lấy cấu trúc các cột của một bảng trong PostgreSQL."""
     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;
-        """
+        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."
-                
+            if not results: return f"Bảng '{table_name}' không tồn tại."
             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()
-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."""
+def select_postgres_records(table_name: str, query: str = "{}", limit: int = 5) -> str:
+    """Lấy dữ liệu từ PostgreSQL. Tham số query là chuỗi JSON hợp lệ."""
     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:
         query_dict = parse_json_safe(query)
-        # Ép buộc filter bảo mật
-        query_dict["project_id"] = project_id
-        query_dict["is_deleted"] = False
-        
-        # Build câu lệnh SQL an toàn với tham số
         where_clauses = []
         values = []
         for key, value in query_dict.items():
-            if not validate_table_name(key): continue # Bỏ qua tên cột nguy hiểm
+            if not validate_table_name(key): continue 
             where_clauses.append(f"{key} = %s")
             values.append(value)
             
-        where_sql = " AND ".join(where_clauses)
+        where_sql = " AND ".join(where_clauses) if where_clauses else "1=1"
         safe_limit = min(max(1, limit), 20)
         
         sql = f"SELECT * FROM {table_name} WHERE {where_sql} LIMIT {safe_limit}"
@@ -440,68 +417,57 @@ def select_postgres_records(project_id: str, table_name: str, query: str = "{}",
         with pg_conn.cursor(cursor_factory=RealDictCursor) as cursor:
             cursor.execute(sql, tuple(values))
             results = cursor.fetchall()
-            
-            # Xử lý datetime để convert sang chuỗi JSON được
+            # Xử lý format ngày tháng và UUID để đưa vào JSON
             for row in results:
                 for k, v in row.items():
                     if isinstance(v, datetime): row[k] = v.isoformat()
-            
+                    elif hasattr(v, 'hex'): row[k] = str(v) 
             return json.dumps(results, indent=2, ensure_ascii=False)
     except Exception as e:
         pg_conn.rollback()
         return f"Lỗi truy vấn Postgres: {str(e)}"
 
 @mcp.tool()
-def insert_postgres_record(project_id: str, table_name: str, data: str) -> str:
-    """Thêm record mới vào PostgreSQL. Tự động gán project_id và timestamp."""
+def insert_postgres_record(table_name: str, data: str) -> str:
+    """Thêm record mới vào PostgreSQL. Data là JSON string."""
     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:
         doc_dict = parse_json_safe(data)
         if not isinstance(doc_dict, dict): return "Lỗi: Data phải là JSON object."
         
-        # Ép buộc metadata
-        doc_dict["project_id"] = project_id
-        doc_dict["created_at"] = datetime.now()
-        doc_dict["updated_at"] = doc_dict["created_at"]
-        doc_dict["is_deleted"] = False
-        
         columns = []
         values = []
         placeholders = []
-        
         for key, value in doc_dict.items():
             if not validate_table_name(key): continue
             columns.append(key)
             values.append(value)
             placeholders.append("%s")
             
-        sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)}) RETURNING id"
-        
-        with pg_conn.cursor() as cursor:
+        sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)}) RETURNING *"
+        with pg_conn.cursor(cursor_factory=RealDictCursor) as cursor:
             cursor.execute(sql, tuple(values))
-            inserted_id = cursor.fetchone()[0]
-            return f"Thêm thành công! Inserted ID: {inserted_id}"
+            inserted_row = cursor.fetchone()
+            for k, v in inserted_row.items():
+                if isinstance(v, datetime): inserted_row[k] = v.isoformat()
+                elif hasattr(v, 'hex'): inserted_row[k] = str(v)
+            return f"Thêm thành công! Dữ liệu: {json.dumps(inserted_row, ensure_ascii=False)}"
     except Exception as e:
         pg_conn.rollback()
         return f"Lỗi insert Postgres: {str(e)}"
 
 @mcp.tool()
-def update_postgres_record(project_id: str, table_name: str, record_id: str, update_data: str) -> str:
-    """Cập nhật record PostgreSQL. Yêu cầu record_id (khóa chính)."""
+def update_postgres_record(table_name: str, record_id: str, update_data: str, id_column: str = "id") -> str:
+    """Cập nhật record PostgreSQL."""
     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:
         update_dict = parse_json_safe(update_data)
-        update_dict.pop("project_id", None)
-        update_dict.pop("id", None)
-        update_dict["updated_at"] = datetime.now()
+        update_dict.pop(id_column, None) # Chống sửa khóa chính
         
         set_clauses = []
         values = []
-        
         for key, value in update_dict.items():
             if not validate_table_name(key): continue
             set_clauses.append(f"{key} = %s")
@@ -509,37 +475,108 @@ def update_postgres_record(project_id: str, table_name: str, record_id: str, upd
             
         if not set_clauses: return "Lỗi: Không có dữ liệu hợp lệ để cập nhật."
         
-        # Thêm thông tin cho WHERE clause
-        values.extend([record_id, project_id])
-        
-        sql = f"UPDATE {table_name} SET {', '.join(set_clauses)} WHERE id = %s AND project_id = %s AND is_deleted = False"
-        
+        values.append(record_id)
+        sql = f"UPDATE {table_name} SET {', '.join(set_clauses)} WHERE {id_column} = %s"
         with pg_conn.cursor() as cursor:
             cursor.execute(sql, tuple(values))
-            if cursor.rowcount == 0:
-                return "Không tìm thấy record hoặc không có quyền sửa."
+            if cursor.rowcount == 0: return "Không tìm thấy record để sửa."
             return "Cập nhật thành công!"
     except Exception as e:
         pg_conn.rollback()
         return f"Lỗi update Postgres: {str(e)}"
 
 @mcp.tool()
-def delete_postgres_record(project_id: str, table_name: str, record_id: str) -> str:
-    """Xóa mềm (Soft Delete) record PostgreSQL."""
+def delete_postgres_record(table_name: str, record_id: str, id_column: str = "id", status_column: str = "is_active") -> str:
+    """Xóa Mềm (Soft Delete) an toàn bằng cách cập nhật cờ trạng thái (mặc định is_active = False)."""
     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:
-        sql = f"UPDATE {table_name} SET is_deleted = True, updated_at = %s WHERE id = %s AND project_id = %s"
-        
+        # Cập nhật is_active = False thay vì xóa cứng
+        sql = f"UPDATE {table_name} SET {status_column} = False WHERE {id_column} = %s"
         with pg_conn.cursor() as cursor:
-            cursor.execute(sql, (datetime.now(), record_id, project_id))
+            cursor.execute(sql, (record_id,))
             if cursor.rowcount == 0:
                 return "Không tìm thấy record hợp lệ để xóa."
-            return "Đã xóa (mềm) thành công record."
+            return f"Đã xóa mềm thành công (set {status_column} = False)."
+    except Exception as e:
+        pg_conn.rollback()
+        return f"Lỗi delete Postgres: Có thể bảng này không có cột {status_column}. Chi tiết: {str(e)}"
+
+@mcp.tool()
+def get_postgres_foreign_keys(table_name: str) -> str:
+    """Lấy danh sách các Khóa ngoại (Ràng buộc liên kết bảng) của một 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:
+        sql = """
+            SELECT
+                kcu.column_name,
+                ccu.table_name AS foreign_table_name,
+                ccu.column_name AS foreign_column_name
+            FROM information_schema.table_constraints AS tc
+            JOIN information_schema.key_column_usage AS kcu
+              ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema
+            JOIN information_schema.constraint_column_usage AS ccu
+              ON ccu.constraint_name = tc.constraint_name AND ccu.table_schema = tc.table_schema
+            WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.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 có khóa ngoại nào."
+            return json.dumps(results, indent=2, ensure_ascii=False)
+    except Exception as e:
+        pg_conn.rollback()
+        return f"Lỗi lấy khóa ngoại: {str(e)}"
+
+@mcp.tool()
+def get_postgres_triggers(table_name: str) -> str:
+    """Lấy danh sách các Trigger (Logic tự động chạy) được gắn trên một 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:
+        sql = """
+            SELECT event_manipulation AS event, action_statement AS function_call
+            FROM information_schema.triggers
+            WHERE event_object_table = %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 có Trigger nào."
+            return json.dumps(results, indent=2, ensure_ascii=False)
+    except Exception as e:
+        pg_conn.rollback()
+        return f"Lỗi lấy triggers: {str(e)}"
+
+@mcp.tool()
+def get_postgres_function_code(function_name: str) -> str:
+    """Lấy mã nguồn (source code) của một function hoặc stored procedure trong PostgreSQL. 
+    Lưu ý: Chỉ truyền tên hàm, KHÔNG truyền dấu ngoặc đơn (). Ví dụ: truyền 'apply_stock_transaction' thay vì 'apply_stock_transaction()'."""
+    if pg_conn is None: return "Lỗi: PostgreSQL chưa kết nối."
+    
+    # Xử lý chuỗi an toàn: Bỏ dấu () nếu AI lỡ truyền vào
+    clean_function_name = function_name.replace("()", "").strip()
+    
+    try:
+        # Sử dụng hàm hệ thống pg_get_functiondef để lấy toàn bộ code
+        sql = """
+            SELECT pg_get_functiondef(p.oid) AS function_code
+            FROM pg_proc p
+            JOIN pg_namespace n ON p.pronamespace = n.oid
+            WHERE p.proname = %s AND n.nspname = 'public';
+        """
+        with pg_conn.cursor(cursor_factory=RealDictCursor) as cursor:
+            cursor.execute(sql, (clean_function_name,))
+            results = cursor.fetchall()
+            
+            if not results: 
+                return f"Không tìm thấy code của function '{clean_function_name}'. Có thể hàm không tồn tại hoặc sai tên."
+                
+            return results[0]['function_code']
     except Exception as e:
         pg_conn.rollback()
-        return f"Lỗi delete Postgres: {str(e)}"
+        return f"Lỗi lấy code function: {str(e)}"
 
 if __name__ == "__main__":
     mcp.run()
\ No newline at end of file