Jelajahi Sumber

hue 生成sql插件服务

xiaoguang_li 1 bulan lalu
melakukan
dc8cf63595

+ 0 - 0
agent/__init__.py


+ 251 - 0
agent/db_agent.py

@@ -0,0 +1,251 @@
+from typing import Literal, Annotated, Optional
+
+from langchain.agents.middleware import after_model, before_model, AgentMiddleware
+from langchain_core.messages import SystemMessage
+from langchain_openai import ChatOpenAI
+from langchain.agents import AgentState
+from langchain.agents import create_agent
+from langchain.tools import tool
+from langchain.tools import ToolRuntime
+from langgraph.runtime import Runtime
+from langsmith import traceable
+
+from config.settings import model
+
+from pydantic import BaseModel, Field
+
+llm = ChatOpenAI(
+    api_key=model.get("api_key", ""),
+    base_url=model.get("base_url", "https://dashscope.aliyuncs.com/compatible-mode/v1"),
+    model=model.get("name", "qwen-max")
+)
+
+SYSTEM_CHAT_PROMPT = """你是一个全能助手、数据专家。解答问题全面、简明、清晰、易理解。
+【输出约束】
+必须只返回 JSON,禁止任何解释、文字、markdown、注释。
+必须返回如下结构:
+{{
+    "content": "ai 答复结果",
+    "type": "chat"
+}}
+content 字段必须是 ai 的答复结果,type 字段必须是 chat。
+"""
+
+SYSTEM_SQL_PROMPT = """
+你是一个资深数据工程师、SQL构建专家、 SQL审核专家。
+
+你的目标是:
+根据用户给出的数据需求,生成一条正确、可执行、语义准确、性能合理的 SQL,
+并对该 SQL 进行自我验证与改进,直到通过验证或无法继续优化为止, 生成SQL要符合指定的dialect:{dialect}
+
+【重要!必须严格输出以下完整 JSON,不能缺少任何字段!】
+【必须输出的固定结构】
+{{
+    "content": "这里放生成的 SQL 字符串",
+    "validation_passed": 当前 SQL 是否通过验证, true 或 false, 
+    "feedback": "若未通过验证,对 SQL 问题的明确说明",
+    "attempts": 当前已经尝试生成 SQL 的次数,
+    "type": "sql"  <-- 必须固定输出这个值,不能修改!
+}}
+【输出规则】
+1. 必须只返回 JSON,禁止任何文字解释
+2. 必须包含上面所有 5 个字段
+3. type 字段必须固定为 "sql",不能是其他值
+4. 即使无法生成 SQL,也要严格按上面结构返回 JSON
+
+【工作流程】
+
+1. 如果 content 为空:
+   - 根据用户的数据需求生成一条 SQL
+   - 对该 SQL 进行验证, 并根据验证结果更新 validation_passed 和 feedback, 如果验证未通过,根据 feedback 对 SQL 进行改进,生成新的 SQL, attempts 增加 1, 从头重新执行,重新生成改进SQL,直到生成的 SQL 通过验证或达到最大尝试次数为止
+
+2. 如果 content 不为空:
+   - 对当前 sql 进行严格验证,并根据验证结果更新 validation_passed 和 feedback
+   - 如果验证未通过,根据 feedback 对 SQL 进行改进,生成新的 SQL, attempts 增加 1, 从头重新执行,重新生成改进SQL,直到生成的 SQL 通过验证或达到最大尝试次数为止
+
+3. SQL 验证必须覆盖以下维度:
+   - 是否满足用户的数据需求(语义正确性)
+   - 表与字段是否来自提供的 schema
+   - 过滤条件是否完整且合理(如时间、状态、范围等)
+   - 聚合与分组逻辑是否正确
+   - SQL 语法是否符合指定方言并可执行
+   - 是否存在明显的性能风险(如不必要的全表扫描、冗余子查询)
+
+4. 如果验证未通过:
+   - 将 validation_passed 设为 false
+   - 在 feedback 中给出明确、可操作的问题描述
+   - 基于 feedback 改进 SQL
+   - attempts 增加 1
+   - 从头重新执行,重新生成改进SQL
+
+5. 如果验证通过:
+   - 将 validation_passed 设为 true
+   - feedback 设为通过原因或简要说明
+   - 返回当前 SQL 作为最终结果
+
+【重要输出约束】
+
+- 每一次输出都必须同时包含以下字段:
+  - content
+  - validation_passed
+  - feedback
+- 不要输出多余的解释文本
+- 不要描述你的思考过程
+- 严格按照规定的输出结构返回结果
+
+【SQL 方言约束】
+
+这次生成 SQL 的方言 dialect 是 {dialect}。
+
+如果 dialect = "mysql":
+- 生成的 SQL 必须严格符合 MySQL 5.7 及更高版本 语法规范,并可直接执行
+
+如果 dialect = "hive":
+- 生成的 SQL 必须严格符合 Apache Hive SQL 语法规范,并可直接执行
+
+如果 dialect = "impala":
+- 生成的 SQL 必须严格符合 Cloudera Impala SQL 语法规范,并可直接执行
+
+如存在语法或行为歧义,优先选择最保守、最通用的写法。
+在生成和验证过程中,如不满足以上要求,
+必须将 validation_passed 设为 false,并在 feedback 中说明原因。
+
+【Schema 说明】
+
+本次使用的schema是: {schema}
+schema 中包含当前数据库的完整 Schema 描述,
+包括数据库名、表名、字段名、字段类型及示例值。
+
+你在生成和验证 SQL 时必须遵守以下规则:
+
+1. 只能使用 schema 中明确列出的数据库、表和字段
+2. 不得假设、编造或引用任何 schema 中不存在的表或字段
+3. 表名必须与 schema 中的 Table 名完全一致
+4. 字段名必须来自对应表的字段列表
+5. 示例值仅用于帮助理解字段语义,不代表完整数据
+6. 生成 SQL 时,表名请使用 schema 中给出的完整表名(包含 schema 前缀)
+
+在 SQL 验证阶段:
+- 如果 SQL 中使用了 schema 中不存在的表或字段,
+  必须将 validation_passed 设为 false,
+  并在 feedback 中说明具体问题。
+"""
+
+MAX_ATTEMPTS = 3
+
+class SQLAgentResponse(BaseModel):
+    content: Optional[str] = Field(
+        description="ai 答复结果 或 当前生成或优化后的 SQL", default=None
+    )
+    validation_passed: Optional[bool] = Field(
+        description="SQL 是否通过验证", default=False
+    )
+    feedback: Optional[str] = Field(
+        description="验证未通过时的问题说明或通过原因", default=None
+    )
+    attempts: Optional[int] = Field(
+        description="尝试执行次数", default=0
+    )
+    type: str = Field(
+        description="结果类型,sql或chat", default="chat"
+    )
+
+class SQLAgentState(AgentState[SQLAgentResponse]):
+    pass
+
+class SQLAgentContext:
+    dialect: Literal["mysql", "hive", "impala"]  # "mysql" | "hive"
+    schema: str # 数据库的 schema 信息
+
+class MyMiddleware(AgentMiddleware):
+    def before_agent(self, state:SQLAgentState, runtime: Runtime[SQLAgentContext]):
+        try:
+            question = state["messages"][-1].content
+            message = llm.invoke(input=[
+                ("system", "你是一个需求分析专家,分析用户的输入问题,确定用户需求是否为生成sql需求,是则返回yes,否则返回no"),
+                ("user", question)
+            ])
+            if "yes" == message.content.lower():
+                print("用户需求被判断为生成 SQL 相关,继续执行 agent")
+                runtime.context["generate_sql"] = True
+                runtime.context["system_prompt"] = SYSTEM_SQL_PROMPT.format(dialect =runtime.context["dialect"], schema= runtime.context["schema"]);
+            else:
+                print("用户需求被判断为非生成 SQL 相关,继续执行 chat")
+                runtime.context["generate_sql"] = False
+                runtime.context["system_prompt"] = SYSTEM_CHAT_PROMPT
+        except Exception as e:
+            print(f"需求分析模型调用失败,错误信息: {e}")
+            state["structured_response"] = {"content": f"需求分析失败,error: {e}", "type": "chat"}
+            state["jump_to"] = "end"
+
+        return state
+
+    def before_model(self, state, runtime: Runtime[SQLAgentContext]):
+        # print("2. before_model 每次模型前都跑")
+        return state
+
+    def wrap_model_call(self, request, handler):
+        if request.runtime.context["system_prompt"]:
+            request = request.override(system_message = SystemMessage(content=request.runtime.context["system_prompt"]))
+        # print("3. wrap_model_call 每次模型都会走")
+        return handler(request)
+
+    def after_model(self, state, runtime: Runtime[SQLAgentContext]):
+        # print("4. after_model 每次模型后都跑")
+        return state
+
+    def after_agent(self, state, runtime: Runtime[SQLAgentContext]):
+        # print("5. after_agent 只跑一次")
+        return state
+
+@after_model
+def attempts_guard_middleware(state: SQLAgentState, runtime: Runtime[SQLAgentContext]):
+    """
+    控制最大执行步数
+    """
+    resp: SQLAgentResponse = state["structured_response"]
+    print("---------  attempts_guard_middleware  ----------")
+    if resp.attempts >= MAX_ATTEMPTS:
+        # 强制收敛
+        resp.validation_passed = True
+        if not resp.feedback:
+            resp.feedback = (
+                f"已达到最大尝试次数 {MAX_ATTEMPTS},最后一次feedback: {resp.feedback},"
+                "返回当前最优 SQL"
+            )
+    else:
+        resp.attempts = resp.attempts+1
+
+    return state
+
+agent = create_agent(
+    model=llm,
+    # system_prompt = SYSTEM_PROMPT,
+    response_format= SQLAgentResponse,
+    state_schema= SQLAgentState,
+    context_schema= SQLAgentContext,
+    # tools= [get_feature],
+    middleware= [MyMiddleware()],
+    debug=False,
+    # checkpointer= memory,
+)
+
+@traceable
+def invoke(msg: str, schema:str, dialect: str = "hive"):
+    result= agent.invoke(
+        input= {
+            "messages": [
+                ("user", msg)
+            ]
+        },
+        context= {
+            "dialect": dialect,
+            "schema": schema
+        }
+    )
+    # response = result["structured_response"]
+    return result
+
+if __name__ == "__main__":
+    # invoke("My Question is who are you?", "mysql", "mysql")
+    invoke("Could you help me calculate the average of the total number of payments made using the most preferred payment method for each product category, where the most preferred payment method in a category is the one with the highest number of payments? ", "mysql", "mysql" )

+ 23 - 0
agent/middlewares.py

@@ -0,0 +1,23 @@
+from langchain.agents.middleware import AgentMiddleware
+
+
+class MyMiddleware(AgentMiddleware):
+    def before_agent(self, state, runtime):
+        print("1. before_agent 只跑一次")
+        return state
+
+    def before_model(self, state, runtime):
+        print("2. before_model 每次模型前都跑")
+        return state
+
+    def wrap_model_call(self, request, handler):
+        print("3. wrap_model_call 每次模型都会走")
+        return handler(request)
+
+    def after_model(self, state, runtime):
+        print("4. after_model 每次模型后都跑")
+        return state
+
+    def after_agent(self, state, runtime):
+        print("5. after_agent 只跑一次")
+        return state

+ 0 - 0
api/__init__.py


+ 95 - 0
api/routes.py

@@ -0,0 +1,95 @@
+import traceback
+
+from fastapi import APIRouter, HTTPException, Query, Body
+from typing import Dict, Any
+from services.metadata_service import MetadataService
+from services.chat_service import ChatService
+from threading import Lock
+
+# 创建子路由
+router = APIRouter(prefix="", tags=["metadata"])
+
+#  服务实例
+_metadata_svc: MetadataService | None = None
+_chat_svc: ChatService | None = None
+_init_lock = Lock()
+
+def get_chat_svc() -> ChatService:
+    global _chat_svc
+    if _chat_svc is None:
+        with _init_lock:
+            if _chat_svc is None:
+                _chat_svc = ChatService()
+    return _chat_svc
+
+def get_metadata_svc() -> MetadataService:
+    global _metadata_svc
+    if _metadata_svc is None:
+        with _init_lock:
+            if _metadata_svc is None:
+                _metadata_svc = MetadataService()
+    return _metadata_svc
+
+@router.get("/sources")
+def list_sources():
+    return {"sources": get_metadata_svc().list_sources()}
+
+
+@router.get("/databases")
+def list_databases(source: str = Query(...)):
+    dbs = get_metadata_svc().get_databases(source)
+    if not dbs:
+        raise HTTPException(404, detail=f"Source '{source}' not found or no databases")
+    return {"source": source, "databases": dbs}
+
+
+@router.get("/tables")
+def list_tables(source: str = Query(...), database: str = Query(...)):
+    tables = get_metadata_svc().get_tables(source, database)
+    if not tables:
+        raise HTTPException(404, detail="Database or source not found")
+    return {"source": source, "database": database, "tables": tables}
+
+
+@router.get("/columns")
+def list_columns(
+    source: str = Query(...),
+    database: str = Query(...),
+    table: str = Query(...)
+):
+    cols = get_metadata_svc().get_columns(source, database, table)
+    if not cols:
+        raise HTTPException(404, detail="Table not found")
+    return {
+        "source": source,
+        "database": database,
+        "table": table,
+        "columns": cols
+    }
+
+
+@router.post("/generate_sql")
+def generate_sql(payload: Dict[str, Any] = Body(...)):
+    """
+    请求体示例:
+    {
+        "source": "hive_prod",
+        "database": "sales",
+        "requirement": "最近7天销售额最高的10个用户"
+    }
+    """
+    source = payload.get("source")
+    database = payload.get("database")
+    requirement = payload.get("requirement")
+
+    if not all([source, database, requirement]):
+        raise HTTPException(400, "Missing 'source', 'database' or 'requirement'")
+
+    try:
+        schema = get_metadata_svc().build_mschema(source, database)
+        result = get_chat_svc().generate_sql(requirement, schema, source)
+        response = result["structured_response"]
+        return response
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(500, detail=str(e))

+ 34 - 0
config.yaml

@@ -0,0 +1,34 @@
+# 服务配置
+server:
+  host: "0.0.0.0"
+  port: 8000
+  log_level: "info"
+
+# 数据库配置
+databases:
+  - name: "hive"
+    type: "hive"
+    url: "jdbc:hive2://10.26.30.138:10000/default"
+    username: "tjrd"
+    driver: "org.apache.hive.jdbc.HiveDriver"
+    jars: "jdbc/hive-jdbc-2.1.1-standalone.jar,jdbc/hadoop-common-3.0.0.jar"
+
+  - name: "impala"
+    type: "impala"
+    host: "qd01-cloud-datanode009.ps.easou.com"
+    port: 21050
+    user: tjrd
+
+  - name: "mysql"
+    type: "mysql"
+    host: "mysql.busydata.cn"
+    port: 3310
+    username: "esbook"
+    password: "esbook_easou"
+    database: "bookcp_v2"  # 可选,若不指定则查所有库
+
+model:
+  api_key: "sk-20dbd7b84d38523062b7388e0829402e7a4c3f32da394af4b56ee79bfc91408d"
+  base_url: "http://198.11.179.174:8080/v1"
+  name: "gpt-5.5"
+  type: "gpt"

+ 0 - 0
config/__init__.py


+ 14 - 0
config/settings.py

@@ -0,0 +1,14 @@
+import yaml
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).parent.parent
+CONFIG_FILE = PROJECT_ROOT / "config.yaml"
+
+with open(CONFIG_FILE, "r", encoding="utf-8") as f:
+    _config = yaml.safe_load(f)
+
+server = _config.get("server", {})
+
+databases = _config.get("databases", [])
+
+model = _config.get("model", {})

TEMPAT SAMPAH
jdbc/hadoop-common-3.0.0.jar


TEMPAT SAMPAH
jdbc/hive-jdbc-2.1.1-standalone.jar


+ 33 - 0
main.py

@@ -0,0 +1,33 @@
+import os
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from config.settings import server as conf
+from api.routes import router
+import uvicorn
+
+# 1. 创建 FastAPI 应用
+app = FastAPI(title="Chat-Based SQL Generator Backend")
+
+# 2. 注册路由器
+app.include_router(router)
+
+#origins=["http://hue.busydata.cn"]
+# 3. 配置跨域请求(根据需要调整)
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=False, allow_methods=["*"], allow_headers=["*"],)
+
+# 4. 启动
+if __name__ == "__main__":
+    os.putenv("LANGSMITH_TRACING", "true")
+    os.putenv("LANGSMITH_ENDPOINT", "https://api.smith.langchain.com")
+    os.putenv("LANGSMITH_API_KEY", "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")
+    os.putenv("LANGSMITH_PROJECT", "chat_service")
+
+    uvicorn.run(
+        "main:app",
+        host= conf.get("host", "0.0.0.0"),
+        port= conf.get("port", 8000),
+        log_level= conf.get("log_level", "info"),
+        reload=True  # 开发时可开启
+    )

+ 0 - 0
providers/__init__.py


+ 56 - 0
providers/base.py

@@ -0,0 +1,56 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any
+
+class MetadataProvider(ABC):
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.name = config["name"]
+
+    def get_type(self) -> str:
+        return self.config["type"]
+
+    @abstractmethod
+    def has_table(self, table_name:str, schema:str) -> bool:
+        """检查表是否存在"""
+        pass
+
+    @abstractmethod
+    def get_table_comment(self, table_name:str, schema:str) -> str:
+        """检查表是否存在"""
+        pass
+
+    def get_pk_constraint(self, table_name:str, schema:str) -> Dict:
+        return {}
+
+    def get_unique_constraints(self, table_name:str, schema:str) -> Dict:
+        return {}
+
+    def get_foreign_keys(self, table_name:str, schema:str) -> list[Dict]:
+        return []
+
+    @abstractmethod
+    def fectch_distinct_values(self, column_name: str, table_name:str, schema:str, max_num: int = 5) -> list[str]:
+        """检查表是否存在"""
+        pass
+
+    @abstractmethod
+    def get_databases(self) -> List[str]:
+        """返回所有数据库/Schema 名称"""
+        pass
+
+    @abstractmethod
+    def get_tables(self, database: str) -> List[str]:
+        """返回某数据库下的所有表名"""
+        pass
+
+    @abstractmethod
+    def get_columns(self, table: str, database: str) -> List[Dict[str, Any]]:
+        """
+        返回字段列表,每个字段为 dict:
+        {
+            "name": "id",
+            "type": "INT",
+            "comment": "用户ID"
+        }
+        """
+        pass

+ 117 - 0
providers/hive.py

@@ -0,0 +1,117 @@
+from os.path import split
+
+from .base import MetadataProvider
+from typing import List, Dict, Any
+import time
+import threading
+import jaydebeapi
+from pathlib import Path
+
+class HiveMetadataProvider(MetadataProvider):
+    def __init__(self, config: Dict):
+        super().__init__(config)
+        self.cache: Dict[str, Dict[str, Dict[str, Any]]] = {}
+        # 创建一个线程锁
+        self._refresh_lock = threading.Lock()
+
+    def _get_connection(self):
+        """
+        创建 Hive JDBC 连接
+        """
+        BASE_DIR = Path(__file__).resolve().parent.parent
+        jdbc_jars = [f"{Path(BASE_DIR)}/{jar.strip()}" for jar in str(self.config["jars"]).split(",")]
+
+        return jaydebeapi.connect(
+            jclassname= self.config.get("driver", "org.apache.hive.jdbc.HiveDriver"),
+            url= self.config["url"],
+            driver_args=[self.config.get("username", "tjrd"), self.config.get("password", "")],
+            jars=jdbc_jars,
+        )
+
+    def get_databases(self) -> List[str]:
+
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute("SHOW DATABASES")
+            return [row[0] for row in cursor.fetchall()]
+        finally:
+            conn.close()
+
+    def get_tables(self, database: str) -> List[str]:
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute(f"SHOW TABLES IN `{database}`")
+            return [row[0] for row in cursor.fetchall()]
+        finally:
+            conn.close()
+
+    def get_columns(self, table: str, database: str) -> List[Dict[str, Any]]:
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute(f"DESCRIBE `{database}`.`{table}`")
+            columns = []
+            for row in cursor.fetchall():
+                # Hive JDBC DESCRIBE: col_name, data_type, comment
+                if not row or row[0].startswith("#") or str(row[0]).strip() == "" or any(c.get("name") == row[0] for c in columns):
+                    continue
+                columns.append({
+                    "name": row[0],
+                    "type": row[1] if len(row) > 1 else "",
+                    "comment": row[2] if len(row) > 2 else "",
+                })
+            return columns
+        except Exception as e:
+            print(f"Error fetching columns for {database}.{table}: {e}")
+            return []
+        finally:
+            conn.close()
+
+    def has_table(self, table_name: str, schema: str) -> bool:
+        self.refresh_db_cache(schema)
+        return schema in self.cache and table_name in self.cache[schema]
+
+    def get_table_comment(self, table_name: str, schema: str) -> str:
+        self.refresh_db_cache(schema)
+        return self.cache.get(schema, {}).get(table_name, {}).get("comment", "")
+
+    def fectch_distinct_values(self, column_name: str, table_name: str, schema: str, max_num: int = 5) -> list[str]:
+        """检查表是否存在"""
+        return []
+
+    def get_table_comment_db(self, table_name: str, schema: str) -> str:
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(f"DESCRIBE FORMATTED `{schema}`.`{table_name}`")
+            for row in cursor.fetchall():
+                if len(row) >= 3 and str(row[1]).strip() == "comment":
+                    return str(row[2]).strip()
+        return ""
+
+    def refresh_db_cache(self, schema: str = ""):
+        db_cache = self.cache.get(schema, None)
+        time_secs = int(time.time())
+        if db_cache is not None and time_secs - db_cache["cache_time"] < 600: # 10 minutes, 缓存有有效
+            return
+
+        with self._refresh_lock:
+            if db_cache is not None and time_secs - db_cache["cache_time"] < 600:   # double check  缓存有有效
+                return
+
+            # 重新构建缓存
+            table_cache: Dict[str, Dict[str, Any]] = {}
+            tables = self.get_tables(schema)
+            for table in tables:
+                comment = self.get_table_comment_db(table, schema)
+                columns = {}  # self.get_columns(table, database)
+                table_cache[table] = {"comment": comment, "columns": columns, "table": table, "schema": schema}
+                print(f"{schema}.{table} comment: {comment}, columns: {len(columns)}")
+            # 更新缓存
+            self.cache[schema] = {"cache_time": int(time.time()), "tables": table_cache}
+
+    def refresh_cache(self):
+        databases = self.get_databases()
+        for database in databases:
+            self.refresh_db_cache(database)

+ 97 - 0
providers/impala.py

@@ -0,0 +1,97 @@
+from .base import MetadataProvider
+from typing import List, Dict, Any
+import time
+import threading
+from impala.dbapi import connect
+import getpass
+
+class ImpalaMetadataProvider(MetadataProvider):
+    def __init__(self, config: Dict):
+        super().__init__(config)
+        self.cache: Dict[str, Dict[str, Dict[str, Any]]] = {}
+        # 创建一个线程锁
+        self._refresh_lock = threading.Lock()
+
+    def _get_connection(self):
+        return connect(
+            host=self.config["host"],
+            port=self.config["port"],
+            auth_mechanism='NOSASL' # NOSASL 方式,传递的user参数不生效,而实际是使用服务启动用户
+        )
+
+    def get_databases(self) -> List[str]:
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            if "tjrd" == getpass.getuser():
+                cursor.execute("INVALIDATE METADATA")
+            cursor.execute("SHOW SCHEMAS")
+            return [row[0] for row in cursor.fetchall()]
+
+    def get_tables(self, database: str) -> List[str]:
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(f"SHOW TABLES IN `{database}`")
+            return [row[0] for row in cursor.fetchall()]
+
+    def get_columns(self, table: str, database: str) -> List[Dict[str, Any]]:
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(f"DESCRIBE `{database}`.`{table}`")
+            columns = []
+            for row in cursor.fetchall():
+                # Impala: (name, type, comment)
+                if not row or row[0].startswith("#") or str(row[0]).strip() == "" or any(c.get("name") == row[0] for c in columns) :
+                    continue
+                columns.append({
+                    "name": row[0],
+                    "type": row[1] if len(row) > 1 else "",
+                    "comment": row[2] if len(row) > 2 else "",
+                })
+            return columns
+
+    def has_table(self, table_name: str, schema: str) -> bool:
+        self.refresh_db_cache(schema)
+        return schema in self.cache and table_name in self.cache[schema]
+
+    def get_table_comment(self, table_name: str, schema: str) -> str:
+        self.refresh_db_cache(schema)
+        return self.cache.get(schema, {}).get(table_name, {}).get("comment", "")
+
+    def fectch_distinct_values(self, column_name: str, table_name: str, schema: str, max_num: int = 5) -> list[str]:
+        """检查表是否存在"""
+        return []
+
+    def get_table_comment_db(self, table_name: str, schema: str) -> str:
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(f"DESCRIBE FORMATTED `{schema}`.`{table_name}`")
+            for row in cursor.fetchall():
+                if len(row) >= 3 and str(row[1]).strip() == "comment":
+                    return str(row[2]).strip()
+        return ""
+
+    def refresh_db_cache(self, schema: str = ""):
+        db_cache = self.cache.get(schema, None)
+        time_secs = int(time.time())
+        if db_cache is not None and time_secs - db_cache["cache_time"] < 600: # 10 minutes, 缓存有有效
+            return
+
+        with self._refresh_lock:
+            if db_cache is not None and time_secs - db_cache["cache_time"] < 600:   # double check  缓存有有效
+                return
+
+            # 重新构建缓存
+            table_cache: Dict[str, Dict[str, Any]] = {}
+            tables = self.get_tables(schema)
+            for table in tables:
+                comment = self.get_table_comment_db(table, schema)
+                columns = {}  # self.get_columns(table, database)
+                table_cache[table] = {"comment": comment, "columns": columns, "table": table, "schema": schema}
+                print(f"{schema}.{table} comment: {comment}, columns: {len(columns)}")
+            # 更新缓存
+            self.cache[schema] = {"cache_time": int(time.time()), "tables": table_cache}
+
+    def refresh_cache(self):
+        databases = self.get_databases()
+        for database in databases:
+            self.refresh_db_cache(database)

+ 75 - 0
providers/mysql.py

@@ -0,0 +1,75 @@
+from .base import MetadataProvider
+from typing import List, Dict, Any
+from sqlalchemy import create_engine, text, inspect
+
+class MySQLMetadataProvider(MetadataProvider):
+
+    def __init__(self, config: Dict):
+        super().__init__(config)
+        url = f"mysql+pymysql://{config['username']}:{config.get('password','')}@{config['host']}:{config['port']}/{config.get('database','')}?charset=utf8mb4"
+        self.engine = create_engine(url, echo=True ) # 设为 True 可打印 SQL 日志
+        self.inspector = inspect(self.engine)
+
+    def _get_connection(self):
+        return self.engine.connect()
+        # return pymysql.connect(
+        #     host=self.config["host"],
+        #     port=self.config["port"],
+        #     user=self.config["username"],
+        #     password=self.config.get("password", ""),
+        #     database=self.config.get("database", ""),  # 可为空
+        #     charset='utf8mb4'
+        # )
+
+    def get_databases(self) -> List[str]:
+        # with self._get_connection() as conn:
+        #     # cursor = conn.cursor()
+        #     # cursor.execute("SHOW DATABASES")
+        #     # return [row[0] for row in cursor.fetchall() if not row[0].startswith('information_schema')]
+        #     result = conn.execute(text("SHOW DATABASES"))
+        #     return [row[0] for row in result if not row[0].startswith('information_schema')]
+        return [row for row in self.inspector.get_schema_names() if not row.startswith('information_schema')]
+
+    def get_tables(self, database: str) -> List[str]:
+        return self.inspector.get_table_names(schema=database)
+
+    def get_columns(self, table: str, database: str) -> List[Dict[str, Any]]:
+        return self.inspector.get_columns(table_name=table, schema=database)
+        # with self._get_connection() as conn:
+        #     cursor = conn.cursor()
+        #     # 查询 information_schema 获取字段详情
+        #     query = """
+        #         SELECT COLUMN_NAME, DATA_TYPE, COLUMN_COMMENT
+        #         FROM information_schema.COLUMNS
+        #         WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
+        #         ORDER BY ORDINAL_POSITION
+        #     """
+        #     cursor.execute(query, (database, table))
+        #     return [
+        #         {
+        #             "name": row[0],
+        #             "type": row[1],
+        #             "comment": row[2] or ""
+        #         }
+        #         for row in cursor.fetchall()
+        #     ]
+
+    def has_table(self, table_name:str, schema:str) -> bool:
+        return self.inspector.has_table(table_name=table_name, schema=schema)
+
+    def get_table_comment(self, table_name:str, schema:str) -> str:
+        return self.inspector.get_table_comment(table_name=table_name, schema=schema)['text']
+
+    def get_pk_constraint(self, table_name:str, schema:str) -> Dict:
+        return self.inspector.get_pk_constraint(table_name=table_name, schema=schema)
+
+    def get_unique_constraints(self, table_name:str, schema:str) -> Dict:
+        return self.inspector.get_unique_constraints(table_name=table_name, schema=schema)
+
+    def get_foreign_keys(self, table_name:str, schema:str) -> list[Dict]:
+        return self.inspector.get_foreign_keys(table_name=table_name, schema=schema)
+
+    def fectch_distinct_values(self, column_name:str, table_name:str, schema:str, max_num:int = 5) -> list[str]:
+        with self._get_connection() as conn:
+            result = conn.execute(text(f"select distinct `{column_name}` from `{schema}`.`{table_name}` where `{column_name}` is not null limit {max_num}"))
+            return [row[0] for row in result]

+ 26 - 0
requirements.txt

@@ -0,0 +1,26 @@
+# 基础框架
+pydantic
+fastapi
+uvicorn
+sqlalchemy
+
+# 配置
+PyYAML
+
+# JDBC(Hive / Impala)
+JPype1
+JayDeBeApi
+
+# Impala
+impyla
+
+# MySQL
+PyMySQL
+
+# HTTP
+requests
+
+# LLM / Agent
+langchain
+langchain-openai
+llama-index

+ 0 - 0
services/__init__.py


+ 11 - 0
services/chat_service.py

@@ -0,0 +1,11 @@
+from agent.db_agent import invoke
+
+class ChatService:
+
+    @staticmethod
+    def generate_sql(requirement: str, schema: str, source_type: str = "hive") -> object:
+        try:
+            result = invoke(requirement, schema, source_type)
+            return result
+        except Exception as e:
+            raise RuntimeError(f"LLM service error: {str(e)}")

+ 56 - 0
services/metadata_service.py

@@ -0,0 +1,56 @@
+from typing import List, Dict, Any
+from config.settings import databases
+from providers.base import MetadataProvider
+from providers.hive import HiveMetadataProvider
+from providers.impala import ImpalaMetadataProvider
+from providers.mysql import MySQLMetadataProvider
+from services.schema.schema_engine import SchemaEngine
+
+
+class MetadataService:
+    def __init__(self):
+        self.providers: Dict[str, MetadataProvider] = {}
+        for db_cfg in databases:
+            name = db_cfg["name"]
+            db_type = db_cfg["type"]
+            try:
+                if db_type == "hive":
+                    self.providers[db_type] = HiveMetadataProvider(db_cfg)
+                elif db_type == "impala":
+                    self.providers[db_type] = ImpalaMetadataProvider(db_cfg)
+                elif db_type == "mysql":
+                    self.providers[db_type] = MySQLMetadataProvider(db_cfg)
+                else:
+                    print(f"⚠️ Unsupported DB type: {db_type}")
+                database_list = self.providers[db_type].get_databases()
+                print( str(database_list))
+                # tables = self.providers[db_type].get_tables(database_list[0])
+                # print( str(tables))
+                # columns = self.providers[db_type].get_columns(tables[0], database_list[0])
+                # print( str(columns))
+
+            except Exception as e:
+                print(f"❌ Failed to init {name}: {e}")
+
+    def get_databases(self, source: str) -> List[str]:
+        provider = self.providers.get(source)
+        return provider.get_databases() if provider else []
+
+    def get_tables(self, source: str, database: str) -> List[str]:
+        provider = self.providers.get(source)
+        return provider.get_tables(database) if provider else []
+
+    def get_columns(self, source: str, database: str, table: str) -> List[Dict[str, str]]:
+        provider = self.providers.get(source)
+        return provider.get_columns(table, database) if provider else []
+
+    def list_sources(self) -> List[str]:
+        return list(self.providers.keys())
+
+    def build_mschema(self, source: str, database: str) -> str:
+        provider = self.providers.get(source)
+        return SchemaEngine(schema=database, metadata_provider=provider, db_name=database, include_tables=None).mschema.to_mschema()
+
+if __name__ == "__main__":
+    svc = MetadataService()
+    print(svc.build_mschema("hive", "default"))

+ 0 - 0
services/schema/__init__.py


+ 186 - 0
services/schema/m_schema.py

@@ -0,0 +1,186 @@
+
+from .utils import examples_to_str, read_json, write_json
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+
+
+class MSchema:
+    def __init__(self, db_id: str = 'Anonymous', schema: Optional[str] = None):
+        self.db_id = db_id
+        self.schema = schema
+        self.tables = {}
+        self.foreign_keys = []
+
+    def add_table(self, name, fields={}, comment=None):
+        self.tables[name] = {"fields": fields.copy(), 'examples': [], 'comment': comment}
+
+    def add_field(self, table_name: str, field_name: str, field_type: str = "",
+            primary_key: bool = False, nullable: bool = True, default: Any = None,
+            autoincrement: bool = False, comment: str = "", examples: list = [], **kwargs):
+        self.tables[table_name]["fields"][field_name] = {
+            "type": field_type,
+            "primary_key": primary_key,
+            "nullable": nullable,
+            "default": default if default is None else f'{default}',
+            "autoincrement": autoincrement,
+            "comment": comment,
+            "examples": examples.copy(),
+            **kwargs}
+
+    def add_foreign_key(self, table_name, field_name, ref_schema, ref_table_name, ref_field_name):
+        self.foreign_keys.append([table_name, field_name, ref_schema, ref_table_name, ref_field_name])
+
+    def get_field_type(self, field_type, simple_mode=True)->str:
+        if not simple_mode:
+            return field_type
+        else:
+            return field_type.split("(")[0]
+
+    def has_table(self, table_name: str) -> bool:
+        if table_name in self.tables.keys():
+            return True
+        else:
+            return False
+
+    def has_column(self, table_name: str, field_name: str) -> bool:
+        if self.has_table(table_name):
+            if field_name in self.tables[table_name]["fields"].keys():
+                return True
+            else:
+                return False
+        else:
+            return False
+
+    def get_field_info(self, table_name: str, field_name: str) -> Dict:
+        try:
+            return self.tables[table_name]['fields'][field_name]
+        except:
+            return {}
+
+    def single_table_mschema(self, table_name: str, selected_columns: List = None,
+                             example_num=3, show_type_detail=False) -> str:
+        table_info = self.tables.get(table_name, {})
+        output = []
+        table_comment = table_info.get('comment', '')
+        if table_comment is not None and table_comment != 'None' and len(table_comment) > 0:
+            if self.schema is not None and len(self.schema) > 0:
+                output.append(f"# Table: {self.schema}.{table_name}, {table_comment}")
+            else:
+                output.append(f"# Table: {table_name}, {table_comment}")
+        else:
+            if self.schema is not None and len(self.schema) > 0:
+                output.append(f"# Table: {self.schema}.{table_name}")
+            else:
+                output.append(f"# Table: {table_name}")
+
+        field_lines = []
+        # 处理表中的每一个字段
+        for field_name, field_info in table_info['fields'].items():
+            if selected_columns is not None and field_name.lower() not in selected_columns:
+                continue
+
+            raw_type = self.get_field_type(field_info['type'], not show_type_detail)
+            field_line = f"({field_name}:{raw_type.upper()}"
+            if field_info['comment'] != '':
+                field_line += f", {field_info['comment'].strip()}"
+            else:
+                pass
+
+            ## 打上主键标识
+            is_primary_key = field_info.get('primary_key', False)
+            if is_primary_key:
+                field_line += f", Primary Key"
+
+            # 如果有示例,添加上
+            if len(field_info.get('examples', [])) > 0 and example_num > 0:
+                examples = field_info['examples']
+                examples = [s for s in examples if s is not None]
+                examples = examples_to_str(examples)
+                if len(examples) > example_num:
+                    examples = examples[:example_num]
+
+                if raw_type in ['DATE', 'TIME', 'DATETIME', 'TIMESTAMP']:
+                    examples = [examples[0]]
+                elif len(examples) > 0 and max([len(s) for s in examples]) > 20:
+                    if max([len(s) for s in examples]) > 50:
+                        examples = []
+                    else:
+                        examples = [examples[0]]
+                else:
+                    pass
+                if len(examples) > 0:
+                    example_str = ', '.join([str(example) for example in examples])
+                    field_line += f", Examples: [{example_str}]"
+                else:
+                    pass
+            else:
+                field_line += ""
+            field_line += ")"
+
+            field_lines.append(field_line)
+        output.append('[')
+        output.append(',\n'.join(field_lines))
+        output.append(']')
+
+        return '\n'.join(output)
+
+    def to_mschema(self, selected_tables: List = None, selected_columns: List = None,
+                   example_num=3, show_type_detail=False) -> str:
+        """
+        convert to a MSchema string.
+        selected_tables: 默认为None,表示选择所有的表
+        selected_columns: 默认为None,表示所有列全选,格式['table_name.column_name']
+        """
+        output = []
+
+        output.append(f"【DB_ID】 {self.db_id}")
+        output.append(f"【Schema】")
+
+        if selected_tables is not None:
+            selected_tables = [s.lower() for s in selected_tables]
+        if selected_columns is not None:
+            selected_columns = [s.lower() for s in selected_columns]
+            selected_tables = [s.split('.')[0].lower() for s in selected_columns]
+
+        # 依次处理每一个表
+        for table_name, table_info in self.tables.items():
+            if selected_tables is None or table_name.lower() in selected_tables:
+                cur_table_type = table_info.get('type', 'table')
+                column_names = list(table_info['fields'].keys())
+                if selected_columns is not None:
+                    cur_selected_columns = [c.lower() for c in column_names if f"{table_name}.{c}".lower() in selected_columns]
+                else:
+                    cur_selected_columns = selected_columns
+                output.append(self.single_table_mschema(table_name, cur_selected_columns, example_num, show_type_detail))
+
+        # 添加外键信息,选择table_type为view时不展示外键
+        if self.foreign_keys:
+            output.append("【Foreign keys】")
+            for fk in self.foreign_keys:
+                ref_schema = fk[2]
+                table1, column1, _, table2, column2 = fk
+                if selected_tables is None or \
+                        (table1.lower() in selected_tables and table2.lower() in selected_tables):
+                    if ref_schema == self.schema:
+                        output.append(f"{fk[0]}.{fk[1]}={fk[3]}.{fk[4]}")
+
+        return '\n'.join(output)
+
+    def dump(self):
+        schema_dict = {
+            "db_id": self.db_id,
+            "schema": self.schema,
+            "tables": self.tables,
+            "foreign_keys": self.foreign_keys
+        }
+        return schema_dict
+
+    def save(self, file_path: str):
+        schema_dict = self.dump()
+        write_json(file_path, schema_dict)
+
+    def load(self, file_path: str):
+        data = read_json(file_path)
+        self.db_id = data.get("db_id", "Anonymous")
+        self.schema = data.get("schema", None)
+        self.tables = data.get("tables", {})
+        self.foreign_keys = data.get("foreign_keys", [])

+ 122 - 0
services/schema/schema_engine.py

@@ -0,0 +1,122 @@
+import json, os
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+
+from providers.base import MetadataProvider
+from .utils import examples_to_str
+from .m_schema import MSchema
+
+
+class SchemaEngine:
+    def __init__(self, schema: Optional[str] = None, metadata_provider: MetadataProvider = None,
+                 ignore_tables: Optional[List[str]] = None, include_tables: Optional[List[str]] = None,
+                 sample_rows_in_table_info: int = 3, indexes_in_table_info: bool = False,
+                 custom_table_info: Optional[dict] = None, view_support: bool = False, max_string_length: int = 300,
+                 mschema: Optional[MSchema] = None, db_name: Optional[str] = ''):
+
+        self._db_name = db_name
+        self._inspector = metadata_provider
+        self._dialect = metadata_provider.get_type()
+
+        # Dictionary to store table names and their corresponding schema
+        self._tables_schemas: Dict[str, str] = {}
+
+        # If a schema is specified, filter by that schema and store that value for every table.
+        if schema:
+            # self._usable_tables = [
+            #     table_name for table_name in self._usable_tables
+            #     if self._inspector.has_table(table_name, schema)
+            # ]
+            # for table_name in self._usable_tables:
+            #     self._tables_schemas[table_name] = schema
+
+            all_tables = []
+            tables = self._inspector.get_tables(database=schema)
+            all_tables.extend(tables)
+            for table in tables:
+                self._tables_schemas[table] = schema
+            self._usable_tables = all_tables
+
+        else:
+            all_tables = []
+            # Iterate through all available schemas
+            for s in self.get_schema_names():
+                tables = self._inspector.get_tables(database=s)
+                all_tables.extend(tables)
+                for table in tables:
+                    self._tables_schemas[table] = s
+            self._usable_tables = all_tables
+
+
+        if mschema is not None:
+            self._mschema = mschema
+        else:
+            self._mschema = MSchema(db_id=db_name, schema=None)
+            self.init_mschema()
+
+    @property
+    def mschema(self) -> MSchema:
+        """Return M-Schema"""
+        return self._mschema
+
+    def get_pk_constraint(self, table_name: str) -> Dict:
+        return self._inspector.get_pk_constraint(table_name, self._tables_schemas[table_name] )
+
+    def get_table_comment(self, table_name: str):
+        try:
+            return self._inspector.get_table_comment(table_name, self._tables_schemas[table_name])
+        except:    # sqlite does not support comments
+            return ''
+
+    def default_schema_name(self) -> Optional[str]:
+        return self._inspector.default_schema_name
+
+    def get_schema_names(self) -> List[str]:
+        return self._inspector.get_databases()
+
+    def get_foreign_keys(self, table_name: str):
+        return self._inspector.get_foreign_keys(table_name, self._tables_schemas[table_name])
+
+    def get_unique_constraints(self, table_name: str):
+        return self._inspector.get_unique_constraints(table_name, self._tables_schemas[table_name])
+
+    def fectch_distinct_values(self, column_name: str, table_name: str, max_num: int = 5):
+        return self._inspector.fectch_distinct_values(column_name, table_name, self._tables_schemas[table_name], max_num)
+
+    def init_mschema(self):
+        for table_name in self._usable_tables:
+            table_comment = self.get_table_comment(table_name)
+            table_comment = '' if table_comment is None else table_comment.strip()
+            table_with_schema = self._tables_schemas[table_name] + '.' + table_name
+            self._mschema.add_table(table_with_schema, fields={}, comment=table_comment)
+            pks = self.get_pk_constraint(table_name)
+
+            fks = self.get_foreign_keys(table_name)
+            for fk in fks:
+                referred_schema = fk['referred_schema']
+                for c, r in zip(fk['constrained_columns'], fk['referred_columns']):
+                    self._mschema.add_foreign_key(table_with_schema, c, referred_schema, fk['referred_table'], r)
+
+            fields = self._inspector.get_columns(table=table_name, database=self._tables_schemas[table_name])
+            for field in fields:
+                field_type = f"{field['type']!s}"
+                field_name = field['name']
+                primary_key = field_name in pks
+                field_comment = field.get("comment", None)
+                field_comment = "" if field_comment is None else field_comment.strip()
+                autoincrement = field.get('autoincrement', False)
+                default = field.get('default', None)
+                nullable = field.get('nullable', False)
+                if default is not None:
+                    default = f'{default}'
+
+                try:
+                    examples = self.fectch_distinct_values(field_name, table_name, 5)
+                except:
+                    examples = []
+                examples = examples_to_str(examples)
+
+                self._mschema.add_field(
+                    table_with_schema, field_name, field_type=field_type, primary_key=primary_key,
+                    nullable=nullable, default=default, autoincrement=autoincrement,
+                    comment=field_comment, examples=examples
+                )

+ 80 - 0
services/schema/utils.py

@@ -0,0 +1,80 @@
+import datetime
+import decimal
+import re
+import json
+
+
+def write_json(path, data):
+    with open(path, 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+
+def read_json(path):
+    with open(path, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+    return data
+
+
+def read_text(filename)->str:
+    data = []
+    with open(filename, 'r', encoding='utf-8') as file:
+        for line in file.readlines():
+            line = line.strip()
+            data.append(line)
+    return data
+
+
+def save_raw_text(filename, content):
+    with open(filename, 'w', encoding='utf-8') as file:
+        file.write(content)
+
+
+def read_map_file(path):
+    data = {}
+    with open(path, 'r', encoding='utf-8') as f:
+        for line in f.readlines():
+            line = line.strip().split('\t')
+            data[line[0]] = line[1].split('、')
+            data[line[0]].append(line[0])
+    return data
+
+
+def save_json(target_file,js,indent=4):
+    with open(target_file, 'w', encoding='utf-8') as f:
+        json.dump(js, f, ensure_ascii=False, indent=indent)
+
+def is_email(string):
+    pattern = r'^[\w\.-]+@[\w\.-]+\.\w+$'
+    match = re.match(pattern, string)
+    if match:
+        return True
+    else:
+        return False
+
+
+def examples_to_str(examples: list) -> list[str]:
+    """
+    from examples to a list of str
+    """
+    values = examples
+    for i in range(len(values)):
+        if isinstance(values[i], datetime.date):
+            values = [values[i]]
+            break
+        elif isinstance(values[i], datetime.datetime):
+            values = [values[i]]
+            break
+        elif isinstance(values[i], decimal.Decimal):
+            values[i] = str(float(values[i]))
+        elif is_email(str(values[i])):
+            values = []
+            break
+        elif 'http://' in str(values[i]) or 'https://' in str(values[i]):
+            values = []
+            break
+        elif values[i] is not None and not isinstance(values[i], str):
+            pass
+        elif values[i] is not None and '.com' in values[i]:
+            pass
+
+    return [str(v) for v in values if v is not None and len(str(v)) > 0]

+ 18 - 0
services/schema/xiyan.py

@@ -0,0 +1,18 @@
+from sqlalchemy import create_engine
+from services.schema.schema_engine import SchemaEngine
+
+
+def mschema():
+    """
+    阿里析言中组织Schema的方式
+    :return:
+    """
+    db_engine = create_engine(f'sqlite:///ecommerce.db')
+    schema_engine = SchemaEngine(engine=db_engine, db_name='ecommerce.db')
+    mschema = schema_engine.mschema
+    mschema_str = mschema.to_mschema()
+    print(mschema_str)
+    return mschema_str
+
+if __name__ == '__main__':
+    mschema()