From b203584cbe4d018982d1d4c57403ee8499ab843d Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Wed, 23 Jul 2025 11:51:19 +0800 Subject: [PATCH 1/7] perf: Delete embedding --- backend/apps/datasource/crud/ds_embedding.py | 250 ------------------- 1 file changed, 250 deletions(-) delete mode 100644 backend/apps/datasource/crud/ds_embedding.py diff --git a/backend/apps/datasource/crud/ds_embedding.py b/backend/apps/datasource/crud/ds_embedding.py deleted file mode 100644 index 7fe50ee8..00000000 --- a/backend/apps/datasource/crud/ds_embedding.py +++ /dev/null @@ -1,250 +0,0 @@ -import logging -from typing import Dict, List -from langchain.schema import Document -from langchain_openai import OpenAIEmbeddings -from langchain.vectorstores.pgvector import PGVector -from sqlalchemy import Tuple - -from apps.chat.schemas.schema_engine import SchemaEngine -from apps.datasource.models.datasource import CoreDatasource -from apps.db.db import get_uri -from langchain_community.utilities import SQLDatabase -from sqlmodel import Session, select -from apps.system.models.system_model import AiModelDetail -from common.core.db import engine -class DsEmbeddingManage: - def __init__(self, ds: CoreDatasource): - self.ds = ds - db_uri = get_uri(ds) - self.db = SQLDatabase.from_uri(db_uri) - - # Get AI model configuration from database - with Session(engine) as session: - aimodel = session.exec(select(AiModelDetail).where( - AiModelDetail.status == True, - AiModelDetail.api_key.is_not(None) - )).first() - - if not aimodel: - raise RuntimeError("No available AI model configuration found") - - # Initialize embedding model with AI model config - self.embedding_model = OpenAIEmbeddings( - model=aimodel.name, - openai_api_base=aimodel.endpoint if aimodel.endpoint else None, - api_key=aimodel.api_key, - ) - - def getSchema(self): - schema_engine = SchemaEngine(engine=self.db._engine) - mschema = schema_engine.mschema - mschema_str = mschema.to_mschema() - return mschema_str - - def mschema_to_documents(self, mschema_str: str) -> List[Document]: - """Convert M-Schema string to list of Documents - - Args: - mschema_str: M-Schema format string - - Returns: - List of Documents for vector storage - """ - docs = [] - current_table = "" - table_content = [] - table_line = "" - current_table_col_docs = [] - # Split schema into lines - lines = mschema_str.split('\n') - - for line in lines: - # Skip empty lines and header sections - if not line.strip() or line.startswith('【'): - continue - - # Process table header - if line.startswith('# Table:'): - # Store previous table if exists - if current_table and table_content: - doc_content = f"Table: {current_table}\n" + '\n'.join(table_content) - table_doc = Document( - page_content=doc_content, - metadata={ - "type": "table", - "name": current_table, - "origin": table_line - } - ) - docs.append(table_doc) - docs.extend(current_table_col_docs) - - # Start new table - current_table = line.split('# Table: ')[1].split(',')[0].strip() - table_content = [] - current_table_col_docs = [] - table_line = line - # If there's a table comment, add it - if ',' in line: - table_content.append(f"Description: {line.split(',')[1].strip()}. Contains columns: ") - continue - - # Add column information to current table - if line.strip() and not line.startswith('[') and not line.startswith(']') and line.startswith('('): - # Clean up the line - column_line = line.strip('(),') - column_info_list = column_line.split(',') - column_name = column_info_list[0].split(':')[0] - table_content.append(column_name + ",\n") - column_content = column_line - - column_doc = Document( - page_content=column_content, - metadata = { - "type": "column", - "name": column_name, - "table": current_table, - "origin": line - } - ) - current_table_col_docs.append(column_doc) - - - # Don't forget to add the last table - if current_table and table_content: - doc_content = f"Table: {current_table}\n" + '\n'.join(table_content) - docs.append(Document( - page_content=doc_content, - metadata={ - "type": "table", - "name": current_table, - "origin": table_line - } - )) - docs.extend(current_table_col_docs) - - return docs - - def store_schema_embeddings(self): - """Store schema embeddings in PGVector""" - try: - # Get schema string - mschema_str = self.getSchema() - - # Convert to documents - docs = self.mschema_to_documents(mschema_str) - # print(docs) - # Store in PGVector - PGVector.from_documents( - embedding=self.embedding_model, - documents=docs, - collection_name="schema_embeddings", - connection_string=str(self.db._engine.url), - pre_delete_collection=True - ) - - return f"Successfully stored {len(docs)} schema documents in vector store" - - except Exception as e: - raise RuntimeError(f"Failed to store schema embeddings: {str(e)}") - - def get_relevant_schema( - self, - query: str, - collection_name: str = "schema_embeddings", - top_k: int = 5 - ) -> List[Dict]: - """Get relevant schema information based on user query - - Args: - query: User query text - collection_name: Vector collection name - top_k: Number of most relevant items to return - - Returns: - List of dictionaries containing relevant schema information - """ - try: - store = PGVector( - collection_name=collection_name, - connection_string=str(self.db._engine.url), - embedding_function=self.embedding_model, - ) - - # Execute similarity search - docs = store.similarity_search_with_score(query, k=top_k) - - relevant_tables = {} - - for doc, score in docs: - metadata = doc.metadata - item = { - "name": metadata["name"], - "type": metadata["type"], - "origin": metadata["origin"], - "score": float(score) - } - - if metadata["type"] == "table": - table_name = metadata["name"] - item["columns"] = [] - if table_name not in relevant_tables: - relevant_tables[table_name] = item - else: - relevant_tables[table_name]["score"] = max(relevant_tables[table_name]["score"], item["score"]) - relevant_tables[table_name]["origin"] = item["origin"] - - elif metadata["type"] == "column": - table_name = f"{metadata['table']}" - if table_name not in relevant_tables: - relevant_tables[table_name] = { - "name": table_name, - "type": "table", - "score": float(score) * 0.9, - "columns": [item] - } - else: - relevant_tables[table_name]["columns"].append(item) - - # Convert dictionary to sorted list - relevant_tables_list = list(relevant_tables.values()) - # Sort by score in descending order - relevant_tables_list.sort(key=lambda x: x["score"], reverse=True) - - return relevant_tables_list - - except Exception as e: - logging.error(f"Error during schema search: {str(e)}") - raise - - def generate_schema_context(self, relevant_tables_list: List[Dict]) -> str: - - context = "## Relevant Database Schema Information:\n\n【Schema】\n" - - if relevant_tables_list: - for table in sorted(relevant_tables_list, key=lambda x: -x["score"]): - if table['origin'] is not None: - context += f"\n{table['origin']}" - else: - context += f"# Table: {table['name']}" - context += f" (relevance score: {table['score']:.2f})\n" - - columns: List[Dict] = table["columns"] - # context += f"[\n{',\n'.join(f'({column["origin"].strip('(),')} , relevance score: {column["score"]:.2f})' for column in columns)}]\n" - context += "\n[\n" + ",\n".join(column["origin"] for column in columns) + "\n]\n" - return context - - def get_schema_with_query(self, - query: str, - collection_name: str = "schema_embeddings", - top_k: int = 5): - list = self.get_relevant_schema(query, collection_name, top_k) - context = self.generate_schema_context(list) - print(context) - return context - - def execute(self): - """Main execution method""" - return self.store_schema_embeddings() - - From dfe261a464289d7f5135c7068c2618f9c6fa9894 Mon Sep 17 00:00:00 2001 From: junjun Date: Wed, 23 Jul 2025 11:53:47 +0800 Subject: [PATCH 2/7] fix: check ds status --- backend/apps/datasource/crud/datasource.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 0960980b..cde525f5 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -86,6 +86,7 @@ def create_ds(session: SessionDep, trans: Trans, user: CurrentUser, create_ds: C def chooseTables(session: SessionDep, id: int, tables: List[CoreTable]): ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first() + check_status(session, ds, True) sync_table(session, ds, tables) updateNum(session, ds) @@ -131,6 +132,7 @@ def getTables(session: SessionDep, id: int): def getTablesByDs(session: SessionDep, ds: CoreDatasource): + check_status(session, ds, True) tables = get_tables(ds) return tables @@ -244,6 +246,7 @@ def preview(session: SessionDep, id: int, data: TableObj): return {"fields": [], "data": [], "sql": ''} ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first() + check_status(session, ds, True) conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config() sql: str = "" if ds.type == "mysql": From f7164f681cf9b8e5b24ee7ccb5b20f19b1617f9b Mon Sep 17 00:00:00 2001 From: junjun Date: Wed, 23 Jul 2025 11:59:31 +0800 Subject: [PATCH 3/7] fix: check ds status --- frontend/src/views/ds/Datasource.vue | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/frontend/src/views/ds/Datasource.vue b/frontend/src/views/ds/Datasource.vue index 5d4a2360..4dd39791 100644 --- a/frontend/src/views/ds/Datasource.vue +++ b/frontend/src/views/ds/Datasource.vue @@ -95,16 +95,12 @@ const handleEditDatasource = (res: any) => { }) } -const handleQuestion = (ele: any) => { - datasourceApi.check(ele).then((res: any) => { - if (res) { - router.push({ - path: '/chat/index', - query: { - start_chat: ele.id, - }, - }) - } +const handleQuestion = (id: string) => { + router.push({ + path: '/chat/index', + query: { + start_chat: id, + }, }) } @@ -294,7 +290,7 @@ const back = () => { :type-name="ele.type_name" :num="ele.num" :description="ele.description" - @question="handleQuestion(ele)" + @question="handleQuestion" @edit="handleEditDatasource(ele)" @del="deleteHandler(ele)" @data-table-detail="dataTableDetail(ele)" From 8c00b44777795316cf853f60e796313419121757 Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Wed, 23 Jul 2025 12:41:02 +0800 Subject: [PATCH 4/7] perf: dev error msg --- backend/apps/system/middleware/auth.py | 4 ++-- backend/common/core/response_middleware.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py index 89c3b1f3..e298d3a7 100644 --- a/backend/apps/system/middleware/auth.py +++ b/backend/apps/system/middleware/auth.py @@ -35,7 +35,7 @@ class TokenMiddleware(BaseHTTPMiddleware): request.state.current_user = validator[1] request.state.assistant = validator[2] return await call_next(request) - return JSONResponse({"msg": f"Unauthorized:[{validator[1]}]"}, status_code=401) + return JSONResponse({"msg": f"Unauthorized:[{validator[1]}]"}, status_code=401, headers={"Access-Control-Allow-Origin": "*"}) #validate pass tokenkey = settings.TOKEN_KEY token = request.headers.get(tokenkey) @@ -43,7 +43,7 @@ class TokenMiddleware(BaseHTTPMiddleware): if validate_pass: request.state.current_user = data return await call_next(request) - return JSONResponse({"msg": f"Unauthorized:[{data}]"}, status_code=401) + return JSONResponse({"msg": f"Unauthorized:[{data}]"}, status_code=401, headers={"Access-Control-Allow-Origin": "*"}) def is_options(self, request: Request): return request.method == "OPTIONS" diff --git a/backend/common/core/response_middleware.py b/backend/common/core/response_middleware.py index 114bdf99..880117d3 100644 --- a/backend/common/core/response_middleware.py +++ b/backend/common/core/response_middleware.py @@ -55,6 +55,10 @@ class ResponseMiddleware(BaseHTTPMiddleware): "code": 500, "data": None, "msg": str(e) + }, + headers={ + k: v for k, v in response.headers.items() + if k.lower() not in ("content-length", "content-type") } ) @@ -70,7 +74,8 @@ class exception_handler(): "code": exc.status_code, "msg": exc.detail, "data": None - } + }, + headers={"Access-Control-Allow-Origin": "*"} ) @@ -84,6 +89,7 @@ class exception_handler(): "code": 500, "msg": str(exc), "data": None - } + }, + headers={"Access-Control-Allow-Origin": "*"} ) From aa00222247a5ffc45fd42ef30f31ed024939b79f Mon Sep 17 00:00:00 2001 From: junjun Date: Wed, 23 Jul 2025 14:25:13 +0800 Subject: [PATCH 5/7] feat: pg tables --- backend/apps/db/db.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index b461d87e..2c00bcaf 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -126,7 +126,7 @@ def get_tables(ds: CoreDatasource): LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = 0 WHERE n.nspname = current_schema() - AND c.relkind IN ('r', 'v', 'p') + AND c.relkind IN ('r', 'v', 'p', 'm') AND c.relname NOT LIKE 'pg_%' AND c.relname NOT LIKE 'sql_%' ORDER BY c.relname \ From 9beb4f8089d377485401bdbfaa9a3cba1f235930 Mon Sep 17 00:00:00 2001 From: wangjiahao <1522128093@qq.com> Date: Wed, 23 Jul 2025 15:27:28 +0800 Subject: [PATCH 6/7] refactor: resource delete --- frontend/src/views/dashboard/common/ResourceTree.vue | 5 +++-- frontend/src/views/dashboard/preview/SQPreviewShow.vue | 8 ++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/frontend/src/views/dashboard/common/ResourceTree.vue b/frontend/src/views/dashboard/common/ResourceTree.vue index 6610257c..c21678a5 100644 --- a/frontend/src/views/dashboard/common/ResourceTree.vue +++ b/frontend/src/views/dashboard/common/ResourceTree.vue @@ -158,7 +158,7 @@ const afterTreeInit = () => { } const copyLoading = ref(false) -const emit = defineEmits(['nodeClick']) +const emit = defineEmits(['nodeClick', 'deleteCurResource']) function createNewObject() { addOperation({ opt: 'newLeaf' }) @@ -230,7 +230,8 @@ const operation = (opt: string, data: SQTreeNode) => { dashboardApi.delete_resource({ id: data.id }).then(() => { ElMessage.success(t('dashboard.delete_success')) getTree() - dashboardStore.setCurComponent(null) + dashboardStore.canvasDataInit() + emit('deleteCurResource') }) }) } else if (opt === 'rename') { diff --git a/frontend/src/views/dashboard/preview/SQPreviewShow.vue b/frontend/src/views/dashboard/preview/SQPreviewShow.vue index d89b887b..84bb5e81 100644 --- a/frontend/src/views/dashboard/preview/SQPreviewShow.vue +++ b/frontend/src/views/dashboard/preview/SQPreviewShow.vue @@ -48,6 +48,13 @@ const mounted = computed(() => { function createNew() { resourceTreeRef.value?.createNewObject() } + +const stateInit = () => { + state.canvasDataPreview = [] + state.canvasStylePreview = {} + state.canvasViewInfoPreview = {} + state.dashboardInfo = {} +} const loadCanvasData = (params: any) => { dataInitState.value = false load_resource_prepare( @@ -100,6 +107,7 @@ defineExpose({ :cur-canvas-type="'dashboard'" :show-position="showPosition" @node-click="resourceNodeClick" + @delete-cur-resource="stateInit" /> Date: Wed, 23 Jul 2025 15:59:09 +0800 Subject: [PATCH 7/7] feat: Log file --- .env.example | 6 + backend/alembic/env.py | 2 - backend/apps/chat/task/llm.py | 37 +++--- backend/apps/datasource/crud/datasource.py | 6 +- backend/apps/system/crud/user.py | 4 +- backend/apps/system/middleware/auth.py | 6 +- backend/common/core/config.py | 5 + backend/common/core/response_middleware.py | 8 +- backend/common/core/sqlbot_cache.py | 19 ++- backend/common/utils/utils.py | 146 ++++++++++++++++++++- backend/common/utils/whitelist.py | 3 +- backend/main.py | 7 +- 12 files changed, 200 insertions(+), 49 deletions(-) diff --git a/.env.example b/.env.example index e60c0f21..10f5bd88 100644 --- a/.env.example +++ b/.env.example @@ -14,6 +14,12 @@ FIRST_SUPERUSER_PASSWORD=123456 # Change this to your pwd TOKEN_KEY="X-SQLBOT-TOKEN" DEFAULT_PWD="SQLBot@123456" +ASSISTANT_TOKEN_KEY="X-SQLBOT-ASSISTANT-TOKEN" + +LOG_LEVEL="INFO" # DEBUG, INFO, WARNING, ERROR +LOG_DIR="/opt/sqlbot/logs" +LOG_FORMAT="%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s" +SQL_DEBUG=False # Postgres POSTGRES_SERVER=localhost diff --git a/backend/alembic/env.py b/backend/alembic/env.py index 1c821a4f..db6e9413 100755 --- a/backend/alembic/env.py +++ b/backend/alembic/env.py @@ -5,7 +5,6 @@ from os.path import abspath, dirname sys.path.insert(0, dirname(dirname(abspath(__file__)))) import os -from logging.config import fileConfig from alembic import context from sqlalchemy import engine_from_config, pool @@ -16,7 +15,6 @@ config = context.config # Interpret the config file for Python logging. # This line sets up loggers basically. -fileConfig(config.config_file_name) # add your model's MetaData object here # for 'autogenerate' support diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index b4786662..c1d718ce 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -1,6 +1,5 @@ import json import sqlparse -import logging import traceback import warnings from typing import Any, List, Optional, Union, Dict @@ -35,7 +34,7 @@ from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, ge from apps.system.schemas.system_schema import AssistantOutDsSchema from common.core.config import settings from common.core.deps import CurrentAssistant, SessionDep, CurrentUser -from common.utils.utils import extract_nested_json +from common.utils.utils import SQLBotLogUtil, extract_nested_json warnings.filterwarnings("ignore") @@ -223,7 +222,7 @@ class LLMService: res = self.llm.stream(analysis_msg) token_usage = {} for chunk in res: - print(chunk) + SQLBotLogUtil.info(chunk) reasoning_content_chunk = '' if 'reasoning_content' in chunk.additional_kwargs: reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') @@ -273,7 +272,7 @@ class LLMService: res = self.llm.stream(predict_msg) token_usage = {} for chunk in res: - print(chunk) + SQLBotLogUtil.info(chunk) reasoning_content_chunk = '' if 'reasoning_content' in chunk.additional_kwargs: reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') @@ -323,7 +322,7 @@ class LLMService: token_usage = {} res = self.llm.stream(guess_msg) for chunk in res: - print(chunk) + SQLBotLogUtil.info(chunk) reasoning_content_chunk = '' if 'reasoning_content' in chunk.additional_kwargs: reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') @@ -390,7 +389,7 @@ class LLMService: token_usage = {} res = self.llm.stream(datasource_msg) for chunk in res: - print(chunk) + SQLBotLogUtil.info(chunk) reasoning_content_chunk = '' if 'reasoning_content' in chunk.additional_kwargs: reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') @@ -474,7 +473,7 @@ class LLMService: token_usage = {} res = self.llm.stream(self.sql_message) for chunk in res: - print(chunk) + SQLBotLogUtil.info(chunk) reasoning_content_chunk = '' if 'reasoning_content' in chunk.additional_kwargs: reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') @@ -543,7 +542,7 @@ class LLMService: res = self.llm.stream(msg) token_usage = {} for chunk in res: - print(chunk) + SQLBotLogUtil.info(chunk) reasoning_content_chunk = '' if 'reasoning_content' in chunk.additional_kwargs: reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') @@ -567,7 +566,7 @@ class LLMService: # 'content': msg.content} for msg # in # analysis_msg]).decode()) - print(full_filter_text) + SQLBotLogUtil.info(full_filter_text) return full_filter_text def generate_chart(self): @@ -582,7 +581,7 @@ class LLMService: token_usage = {} res = self.llm.stream(self.chart_message) for chunk in res: - print(chunk) + SQLBotLogUtil.info(chunk) reasoning_content_chunk = '' if 'reasoning_content' in chunk.additional_kwargs: reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') @@ -691,7 +690,7 @@ class LLMService: Returns: Query results """ - print(f"Executing SQL on ds_id {self.ds.id}: {sql}") + SQLBotLogUtil.info(f"Executing SQL on ds_id {self.ds.id}: {sql}") return exec_sql(self.ds, sql) @@ -717,7 +716,7 @@ def execute_sql_with_db(db: SQLDatabase, sql: str) -> str: except Exception as e: error_msg = f"SQL execution failed: {str(e)}" - logging.error(error_msg) + SQLBotLogUtil.exception(error_msg) raise RuntimeError(error_msg) @@ -741,7 +740,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True): ds_res = llm_service.select_datasource() for chunk in ds_res: - print(chunk) + SQLBotLogUtil.info(chunk) if in_chat: yield orjson.dumps( {'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'), @@ -765,7 +764,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True): yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n' # filter sql - print(full_sql_text) + SQLBotLogUtil.info(full_sql_text) # todo row permission sql_json_str = extract_nested_json(full_sql_text) @@ -785,11 +784,11 @@ def run_task(llm_service: LLMService, in_chat: bool = True): raise Exception("SQL query is empty") sql_result = llm_service.generate_filter(data.get('sql'), data.get('tables')) # maybe no sql and tables - print(sql_result) + SQLBotLogUtil.info(sql_result) sql = llm_service.check_save_sql(res=sql_result) # sql = llm_service.check_save_sql(res=full_sql_text) - print(sql) + SQLBotLogUtil.info(sql) format_sql = sqlparse.format(sql, reindent=True) if in_chat: yield orjson.dumps({'content': format_sql, 'type': 'sql'}).decode() + '\n\n' @@ -815,9 +814,9 @@ def run_task(llm_service: LLMService, in_chat: bool = True): yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n' # filter chart - print(full_chart_text) + SQLBotLogUtil.info(full_chart_text) chart = llm_service.check_save_chart(res=full_chart_text) - print(chart) + SQLBotLogUtil.info(chart) if in_chat: yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n' else: @@ -856,7 +855,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True): if chart['type'] != 'table': yield '### generated chart picture\n\n' image_url = request_picture(llm_service.record.chat_id, llm_service.record.id, chart, result) - print(image_url) + SQLBotLogUtil.info(image_url) yield f'![{chart["type"]}]({image_url})' except Exception as e: traceback.print_exc() diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index cde525f5..c48c954f 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -15,7 +15,7 @@ from apps.db.db import get_engine, get_tables, get_fields, exec_sql from apps.db.engine import get_engine_config, get_engine_conn from apps.db.type import db_type_relation from common.core.deps import SessionDep, CurrentUser, Trans -from common.utils.utils import deepcopy_ignore_extra +from common.utils.utils import SQLBotLogUtil, deepcopy_ignore_extra from .table import get_tables_by_ds_id from ..crud.field import delete_field_by_ds_id, update_field from ..crud.table import delete_table_by_ds_id, update_table @@ -39,12 +39,12 @@ def check_status(session: SessionDep, ds: CoreDatasource, is_raise: bool = False conn = get_engine(ds) try: with conn.connect() as connection: - print("success") + SQLBotLogUtil.info("success") return True except Exception as e: + SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}") if is_raise: raise HTTPException(status_code=500, detail=f'Connect Failed: {e.args}') - print("Fail:", e) return False diff --git a/backend/apps/system/crud/user.py b/backend/apps/system/crud/user.py index c0549f63..ffad38a8 100644 --- a/backend/apps/system/crud/user.py +++ b/backend/apps/system/crud/user.py @@ -1,13 +1,13 @@ from typing import Optional from sqlmodel import Session, select, delete as sqlmodel_delete -import logging from apps.system.models.system_model import UserWsModel, WorkspaceModel from apps.system.schemas.auth import CacheName, CacheNamespace from apps.system.schemas.system_schema import BaseUserDTO, UserInfoDTO, UserWs from common.core.deps import SessionDep from common.core.sqlbot_cache import cache, clear_cache from common.utils.locale import I18n +from common.utils.utils import SQLBotLogUtil from ..models.user import UserModel from common.core.security import verify_md5pwd @@ -68,4 +68,4 @@ async def single_delete(session: SessionDep, id: int): #@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id") async def clean_user_cache(id: int): - logging.info(f"User cache for [{id}] has been cleaned") + SQLBotLogUtil.info(f"User cache for [{id}] has been cleaned") diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py index e298d3a7..365dfff1 100644 --- a/backend/apps/system/middleware/auth.py +++ b/backend/apps/system/middleware/auth.py @@ -1,6 +1,6 @@ from typing import Optional -from fastapi import HTTPException, Request +from fastapi import Request from fastapi.responses import JSONResponse import jwt from sqlmodel import Session @@ -13,6 +13,7 @@ from apps.system.schemas.system_schema import AssistantHeader, UserInfoDTO from common.core import security from common.core.config import settings from common.core.schemas import TokenPayload +from common.utils.utils import SQLBotLogUtil from common.utils.whitelist import whiteUtils from fastapi.security.utils import get_authorization_scheme_param class TokenMiddleware(BaseHTTPMiddleware): @@ -70,6 +71,7 @@ class TokenMiddleware(BaseHTTPMiddleware): ) """ return True, session_user except Exception as e: + SQLBotLogUtil.exception(f"Token validation error: {str(e)}", exc_info=True) return False, e @@ -95,4 +97,6 @@ class TokenMiddleware(BaseHTTPMiddleware): assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True)) return True, session_user, assistant_info except Exception as e: + SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}", exc_info=True) + # Return False and the exception message return False, e \ No newline at end of file diff --git a/backend/common/core/config.py b/backend/common/core/config.py index 99c7fbc7..f9026b09 100644 --- a/backend/common/core/config.py +++ b/backend/common/core/config.py @@ -60,6 +60,11 @@ class Settings(BaseSettings): TOKEN_KEY: str = "X-SQLBOT-TOKEN" DEFAULT_PWD: str = "SQLBot@123456" ASSISTANT_TOKEN_KEY: str = "X-SQLBOT-ASSISTANT-TOKEN" + + LOG_LEVEL: str = "INFO" # DEBUG, INFO, WARNING, ERROR + LOG_DIR: str = "logs" + LOG_FORMAT: str = "%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s" + SQL_DEBUG: bool = False @computed_field # type: ignore[prop-decorator] @property diff --git a/backend/common/core/response_middleware.py b/backend/common/core/response_middleware.py index 880117d3..59590ef8 100644 --- a/backend/common/core/response_middleware.py +++ b/backend/common/core/response_middleware.py @@ -1,10 +1,10 @@ import json -import logging from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import JSONResponse from starlette.exceptions import HTTPException from starlette.requests import Request from common.core.config import settings +from common.utils.utils import SQLBotLogUtil class ResponseMiddleware(BaseHTTPMiddleware): def __init__(self, app): super().__init__(app) @@ -48,7 +48,7 @@ class ResponseMiddleware(BaseHTTPMiddleware): } ) except Exception as e: - logging.error(f"Response processing error: {str(e)}", exc_info=True) + SQLBotLogUtil.error(f"Response processing error: {str(e)}", exc_info=True) return JSONResponse( status_code=500, content={ @@ -68,6 +68,7 @@ class ResponseMiddleware(BaseHTTPMiddleware): class exception_handler(): @staticmethod async def http_exception_handler(request: Request, exc: HTTPException): + SQLBotLogUtil.exception(f"HTTP Exception: {exc.detail}", exc_info=True) return JSONResponse( status_code=exc.status_code, content={ @@ -81,8 +82,7 @@ class exception_handler(): @staticmethod async def global_exception_handler(request: Request, exc: Exception): - logging.error(f"Error info: {str(exc)}", exc_info=True) - + SQLBotLogUtil.exception(f"Unhandled Exception: {str(exc)}", exc_info=True) return JSONResponse( status_code=500, content={ diff --git a/backend/common/core/sqlbot_cache.py b/backend/common/core/sqlbot_cache.py index 196a7275..ea83cdec 100644 --- a/backend/common/core/sqlbot_cache.py +++ b/backend/common/core/sqlbot_cache.py @@ -1,12 +1,13 @@ -import logging from fastapi_cache import FastAPICache from fastapi_cache.decorator import cache as original_cache from functools import partial, wraps -from typing import Optional, Set, Any, Dict, Tuple, Callable +from typing import Optional, Set, Any, Dict, Tuple from inspect import Parameter, signature from contextlib import asynccontextmanager import asyncio +from common.utils.utils import SQLBotLogUtil + # 锁管理 _cache_locks = {} @@ -65,7 +66,7 @@ def custom_key_builder( base_key += f"{value}:" except (IndexError, KeyError, AttributeError) as e: - logging.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}") + SQLBotLogUtil.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}") return base_key @@ -127,19 +128,16 @@ def cache( ) async with _get_cache_lock(cache_key): - logging.info(f"Using cache key: {cache_key}") - print(f"Using cache key: {cache_key}") + SQLBotLogUtil.info(f"Using cache key: {cache_key}") backend = FastAPICache.get_backend() cached_value = await backend.get(cache_key) if cached_value is not None: - logging.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}") - print(f"Cache hit for key: {cache_key}, the value is: {cached_value}") + SQLBotLogUtil.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}") return cached_value result = await func(*args, **kwargs) await backend.set(cache_key, result, expire) - logging.info(f"Cache miss for key: {cache_key}, result cached.") - print(f"Cache miss for key: {cache_key}, result cached.") + SQLBotLogUtil.info(f"Cache miss for key: {cache_key}, result cached.") return result return wrapper @@ -167,8 +165,7 @@ def clear_cache( if await FastAPICache.get_backend().get(cache_key): await FastAPICache.clear(key=cache_key) result = await func(*args, **kwargs) - logging.info(f"Clearing cache for key: {cache_key}") - print(f"Clearing cache for key: {cache_key}") + SQLBotLogUtil.info(f"Clearing cache for key: {cache_key}") return result return await func(*args, **kwargs) diff --git a/backend/common/utils/utils.py b/backend/common/utils/utils.py index f690b27f..bbc4c2a0 100644 --- a/backend/common/utils/utils.py +++ b/backend/common/utils/utils.py @@ -1,6 +1,10 @@ import hashlib +import inspect import logging from datetime import datetime, timedelta, timezone +from logging.handlers import RotatingFileHandler +from pathlib import Path +from common.core.config import settings from typing import Optional import jwt @@ -10,9 +14,6 @@ from jwt.exceptions import InvalidTokenError from common.core import security from common.core.config import settings -logging.basicConfig(level=logging.DEBUG) -logger = logging.getLogger(__name__) - def generate_password_reset_token(email: str) -> str: delta = timedelta(hours=settings.EMAIL_RESET_TOKEN_EXPIRE_HOURS) @@ -76,3 +77,142 @@ def extract_nested_json(text): def string_to_numeric_hash(text: str, bits: Optional[int] = 64) -> int: hash_bytes = hashlib.sha256(text.encode()).digest() return int.from_bytes(hash_bytes[:bits//8], byteorder='big') + + +def setup_logging(): + # 确保日志目录存在 + log_dir = Path(settings.LOG_DIR) + log_dir.mkdir(parents=True, exist_ok=True) + + # 日志格式 + formatter = logging.Formatter( + f'{settings.LOG_FORMAT}' + ) + + # 控制台日志 + console_handler = logging.StreamHandler() + console_handler.setLevel(settings.LOG_LEVEL) + console_handler.setFormatter(formatter) + + # 文件日志处理器 + file_handlers = { + 'debug': logging.DEBUG, + 'info': logging.INFO, + 'warn': logging.WARNING, + 'error': logging.ERROR + } + + # 主日志记录器 + root_logger = logging.getLogger() + root_logger.setLevel(logging.DEBUG) # 设置最低级别 + + # 添加控制台处理器 + root_logger.addHandler(console_handler) + + # 为每个级别创建文件处理器 + for level_name, level in file_handlers.items(): + file_path = log_dir / f"{level_name}.log" + handler = RotatingFileHandler( + file_path, + maxBytes=10 * 1024 * 1024, # 10MB + backupCount=5, + encoding='utf-8' + ) + handler.setLevel(level) + handler.setFormatter(formatter) + + # 添加过滤器只处理特定级别日志 + if level_name == 'debug': + handler.addFilter(lambda record: record.levelno == logging.DEBUG) + elif level_name == 'info': + handler.addFilter(lambda record: record.levelno == logging.INFO) + elif level_name == 'warn': + handler.addFilter(lambda record: record.levelno == logging.WARNING) + elif level_name == 'error': + handler.addFilter(lambda record: record.levelno >= logging.ERROR) + + root_logger.addHandler(handler) + + # SQL日志特殊处理 + if settings.LOG_LEVEL == "DEBUG" and settings.SQL_DEBUG: + sql_logger = logging.getLogger('sqlalchemy.engine') + sql_logger.setLevel(logging.DEBUG) + + sql_handler = RotatingFileHandler( + log_dir / "sql.log", + maxBytes=10 * 1024 * 1024, + backupCount=2, + encoding='utf-8' + ) + sql_handler.setFormatter(formatter) + sql_logger.addHandler(sql_handler) + +setup_logging() + + +class CallerLogger(logging.Logger): + def __init__(self, logger: logging.Logger): + self.logger = logger + super().__init__(logger.name, logger.level) + + def _log(self, level, msg, args, exc_info=None, extra=None, stacklevel=3): + if self.logger.isEnabledFor(level): + self.logger._log(level, msg, args, exc_info=exc_info, extra=extra, stacklevel=stacklevel) + +class SQLBotLogUtil: + + @staticmethod + def _get_logger() -> logging.Logger: + frame = inspect.currentframe() + try: + caller_frame = frame.f_back.f_back + module_name = caller_frame.f_globals.get('__name__', '__main__') + return CallerLogger(logging.getLogger(module_name)) + finally: + del frame + + + @staticmethod + def debug(msg: str, *args, **kwargs): + logger = SQLBotLogUtil._get_logger() + if logger.isEnabledFor(logging.DEBUG): + logger._log(logging.DEBUG, msg, args, **kwargs) + + @staticmethod + def info(msg: str, *args, **kwargs): + logger = SQLBotLogUtil._get_logger() + if logger.isEnabledFor(logging.INFO): + logger._log(logging.INFO, msg, args, **kwargs) + + @staticmethod + def warning(msg: str, *args, **kwargs): + logger = SQLBotLogUtil._get_logger() + if logger.isEnabledFor(logging.WARNING): + logger._log(logging.WARNING, msg, args, **kwargs) + + @staticmethod + def error(msg: str, *args, exc_info: Optional[bool] = None, **kwargs): + logger = SQLBotLogUtil._get_logger() + if logger.isEnabledFor(logging.ERROR): + logger._log( + logging.ERROR, + msg, + args, + exc_info=exc_info if exc_info is not None else True, + **kwargs + ) + + @staticmethod + def exception(msg: str, *args, **kwargs): + logger = SQLBotLogUtil._get_logger() + if logger.isEnabledFor(logging.ERROR): + logger._log(logging.ERROR, msg, args, exc_info=True, **kwargs) + + @staticmethod + def critical(msg: str, *args, **kwargs): + logger = SQLBotLogUtil._get_logger() + if logger.isEnabledFor(logging.CRITICAL): + logger._log(logging.CRITICAL, msg, args, **kwargs) + + + diff --git a/backend/common/utils/whitelist.py b/backend/common/utils/whitelist.py index 10fa0c02..225d94e1 100644 --- a/backend/common/utils/whitelist.py +++ b/backend/common/utils/whitelist.py @@ -2,6 +2,7 @@ import re from typing import List, Pattern from common.core.config import settings +from common.utils.utils import SQLBotLogUtil wlist = [ "/", "/docs", @@ -50,7 +51,7 @@ class WhitelistChecker: try: self._compiled_patterns.append(re.compile(regex_pattern)) except re.error: - print(f"Invalid regular expression pattern: {pattern}") + SQLBotLogUtil.error(f"Invalid regex pattern: {regex_pattern}") def is_whitelisted(self, path: str) -> bool: prefix = settings.API_V1_STR diff --git a/backend/main.py b/backend/main.py index 19ef6f33..54410660 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1,4 +1,3 @@ -import logging from fastapi.concurrency import asynccontextmanager import sentry_sdk from fastapi import FastAPI @@ -18,6 +17,8 @@ import sqlbot_xpack from fastapi_cache import FastAPICache from fastapi_cache.backends.inmemory import InMemoryBackend +from common.utils.utils import SQLBotLogUtil + def run_migrations(): alembic_cfg = Config("alembic.ini") command.upgrade(alembic_cfg, "head") @@ -27,10 +28,10 @@ def run_migrations(): async def lifespan(app: FastAPI): run_migrations() FastAPICache.init(InMemoryBackend()) - print("✅ FastAPICache 初始化完成") init_dynamic_cors(app) + SQLBotLogUtil.info("✅ SQLBot 初始化完成") yield - + SQLBotLogUtil.info("SQLBot 应用关闭") def custom_generate_unique_id(route: APIRoute) -> str: tag = route.tags[0] if route.tags and len(route.tags) > 0 else ""