mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat: M-Schema
This commit is contained in:
fit2cloud-chenyw
+2
-1
@@ -182,4 +182,5 @@ cython_debug/
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
.DS_Store
|
||||
.DS_Store
|
||||
test.py
|
||||
@@ -0,0 +1,276 @@
|
||||
import os
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
from sqlalchemy import inspect, text
|
||||
from sqlmodel import SQLModel, create_engine, Session
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
from langchain.vectorstores import PGVector
|
||||
from langchain.schema import Document
|
||||
import logging
|
||||
from apps.chat.schemas.chat_base_schema import LLMConfig
|
||||
from common.core.config import settings
|
||||
from common.core.db import engine
|
||||
|
||||
class SchemaEmbeddingManager:
|
||||
def __init__(self, config: LLMConfig):
|
||||
"""Initialize SchemaEmbeddingManager
|
||||
|
||||
Args:
|
||||
db_uri: Database connection URI in format postgresql://user:password@host:port/database
|
||||
embedding_model: Embedding model to use
|
||||
"""
|
||||
# self.db_uri = db_uri
|
||||
self.embedding_model = OpenAIEmbeddings(
|
||||
model=config.model_name,
|
||||
openai_api_base=config.api_base_url,
|
||||
api_key=config.api_key
|
||||
)
|
||||
# self.engine = self._create_engine()
|
||||
self.engine = engine
|
||||
self._setup_vector_extension()
|
||||
|
||||
|
||||
def _setup_vector_extension(self):
|
||||
"""Set up PgVector extension using SQLModel"""
|
||||
try:
|
||||
with Session(self.engine) as session:
|
||||
# Enable vector extension
|
||||
session.exec(text("CREATE EXTENSION IF NOT EXISTS vector;"))
|
||||
session.commit()
|
||||
|
||||
# Create vector type if needed
|
||||
session.exec(text("""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'vector') THEN
|
||||
CREATE TYPE vector AS (x float8[]);
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
"""))
|
||||
session.commit()
|
||||
logging.info("Vector extension setup completed successfully")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to setup vector extension: {str(e)}")
|
||||
raise
|
||||
|
||||
def get_session(self) -> Session:
|
||||
"""Get a new database session"""
|
||||
return Session(self.engine)
|
||||
|
||||
def extract_schema_info(self, schema_name: str = "public") -> List[Dict]:
|
||||
"""
|
||||
Extract database schema information and structure it
|
||||
|
||||
Args:
|
||||
schema_name: Schema name to extract, defaults to public
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing schema information
|
||||
"""
|
||||
metadata = SQLModel.metadata
|
||||
metadata.reflect(bind=self.engine, schema=schema_name)
|
||||
|
||||
schema_info = []
|
||||
inspector = inspect(self.engine)
|
||||
|
||||
for table_name in inspector.get_table_names(schema=schema_name):
|
||||
table = metadata.tables[f"{schema_name}.{table_name}"]
|
||||
table_comment = inspector.get_table_comment(table_name, schema=schema_name)
|
||||
|
||||
# Table-level information
|
||||
table_info = {
|
||||
"type": "table",
|
||||
"schema": schema_name,
|
||||
"name": table_name,
|
||||
"description": table_comment.get('text') if table_comment else f"Table {table_name} in schema {schema_name}",
|
||||
"columns": []
|
||||
}
|
||||
|
||||
# Column-level information
|
||||
for column in inspector.get_columns(table_name, schema=schema_name):
|
||||
col_name = column['name']
|
||||
# Get column comment directly from column info
|
||||
col_comment = column.get('comment', f"Column {col_name} in table {table_name}")
|
||||
|
||||
column_info = {
|
||||
"type": "column",
|
||||
"schema": schema_name,
|
||||
"table": table_name,
|
||||
"name": col_name,
|
||||
"data_type": str(column['type']),
|
||||
"description": col_comment,
|
||||
"is_primary_key": col_name in [pk['name'] for pk in inspector.get_pk_constraint(table_name, schema=schema_name).get('constrained_columns', [])],
|
||||
"foreign_key": None # Will be updated below
|
||||
}
|
||||
|
||||
# Get foreign key information
|
||||
for fk in inspector.get_foreign_keys(table_name, schema=schema_name):
|
||||
if col_name in fk['constrained_columns']:
|
||||
column_info["foreign_key"] = f"{fk['referred_table']}.{fk['referred_columns'][0]}"
|
||||
break
|
||||
|
||||
table_info["columns"].append(column_info)
|
||||
|
||||
schema_info.append(table_info)
|
||||
|
||||
return schema_info
|
||||
|
||||
def generate_schema_documents(self, schema_info: List[Dict]) -> List[Document]:
|
||||
"""
|
||||
Convert schema information into LangChain Document objects
|
||||
|
||||
Args:
|
||||
schema_info: Schema information from extract_schema_info
|
||||
|
||||
Returns:
|
||||
List of LangChain Document objects
|
||||
"""
|
||||
documents = []
|
||||
|
||||
for table in schema_info:
|
||||
# Table-level document
|
||||
table_content = f"Table {table['schema']}.{table['name']}: {table['description']}. "
|
||||
table_content += f"Contains columns: {', '.join(col['name'] for col in table['columns'])}"
|
||||
|
||||
table_doc = Document(
|
||||
page_content=table_content,
|
||||
metadata={
|
||||
"type": "table",
|
||||
"schema": table["schema"],
|
||||
"name": table["name"],
|
||||
"description": table["description"],
|
||||
"full_name": f"{table['schema']}.{table['name']}"
|
||||
}
|
||||
)
|
||||
documents.append(table_doc)
|
||||
|
||||
# Column-level documents
|
||||
for column in table["columns"]:
|
||||
column_content = f"Column {table['schema']}.{table['name']}.{column['name']}: {column['description']}. "
|
||||
column_content += f"Data type: {column['data_type']}. "
|
||||
if column["is_primary_key"]:
|
||||
column_content += "This is a primary key. "
|
||||
if column["foreign_key"]:
|
||||
column_content += f"Foreign key to {column['foreign_key']}."
|
||||
|
||||
column_doc = Document(
|
||||
page_content=column_content,
|
||||
metadata={
|
||||
"type": "column",
|
||||
"schema": table["schema"],
|
||||
"table": table["name"],
|
||||
"name": column["name"],
|
||||
"description": column["description"],
|
||||
"data_type": column["data_type"],
|
||||
"full_name": f"{table['schema']}.{table['name']}.{column['name']}"
|
||||
}
|
||||
)
|
||||
documents.append(column_doc)
|
||||
|
||||
return documents
|
||||
|
||||
def store_embeddings(self, documents: List[Document], collection_name: str = "schema_embeddings"):
|
||||
"""Store schema document embeddings in PgVector
|
||||
|
||||
Args:
|
||||
documents: Documents from generate_schema_documents
|
||||
collection_name: Vector collection name
|
||||
"""
|
||||
try:
|
||||
PGVector.from_documents(
|
||||
embedding=self.embedding_model,
|
||||
documents=documents,
|
||||
collection_name=collection_name,
|
||||
connection_string=self.db_uri,
|
||||
pre_delete_collection=True # Optionally recreate collection
|
||||
)
|
||||
logging.info(f"Successfully stored {len(documents)} embeddings in collection {collection_name}")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to store embeddings: {str(e)}")
|
||||
raise
|
||||
|
||||
def get_relevant_schema(
|
||||
self,
|
||||
query: str,
|
||||
collection_name: str = "schema_embeddings",
|
||||
top_k: int = 5
|
||||
) -> Tuple[List[Dict], 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:
|
||||
Tuple containing (relevant_tables_list, relevant_columns_list)
|
||||
"""
|
||||
try:
|
||||
store = PGVector(
|
||||
collection_name=collection_name,
|
||||
connection_string=self.db_uri,
|
||||
embedding_function=self.embedding_model,
|
||||
)
|
||||
|
||||
# Execute similarity search
|
||||
docs = store.similarity_search_with_score(query, k=top_k)
|
||||
|
||||
relevant_tables = []
|
||||
relevant_columns = []
|
||||
seen_tables = set()
|
||||
|
||||
for doc, score in docs:
|
||||
metadata = doc.metadata
|
||||
item = {
|
||||
"name": metadata["full_name"],
|
||||
"description": metadata["description"],
|
||||
"type": metadata["type"],
|
||||
"score": float(score)
|
||||
}
|
||||
|
||||
if metadata["type"] == "table" and metadata["full_name"] not in seen_tables:
|
||||
relevant_tables.append(item)
|
||||
seen_tables.add(metadata["full_name"])
|
||||
elif metadata["type"] == "column":
|
||||
relevant_columns.append(item)
|
||||
# Ensure associated table is also included
|
||||
table_name = f"{metadata['schema']}.{metadata['table']}"
|
||||
if table_name not in seen_tables:
|
||||
relevant_tables.append({
|
||||
"name": table_name,
|
||||
"description": f"Table containing column {metadata['name']}",
|
||||
"type": "table",
|
||||
"score": float(score) * 0.9 # Slightly lower table score
|
||||
})
|
||||
seen_tables.add(table_name)
|
||||
|
||||
return relevant_tables, relevant_columns
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error during schema search: {str(e)}")
|
||||
raise
|
||||
|
||||
def generate_schema_context(self, relevant_tables: List[Dict], relevant_columns: List[Dict]) -> str:
|
||||
"""
|
||||
Generate concise schema context for LLM usage
|
||||
|
||||
Args:
|
||||
relevant_tables: List of relevant tables
|
||||
relevant_columns: List of relevant columns
|
||||
|
||||
Returns:
|
||||
Formatted schema context string
|
||||
"""
|
||||
context = "## Relevant Database Schema Information:\n\n"
|
||||
|
||||
if relevant_tables:
|
||||
context += "### Tables:\n"
|
||||
for table in sorted(relevant_tables, key=lambda x: -x["score"]):
|
||||
context += f"- {table['name']}: {table['description']} (relevance score: {table['score']:.2f})\n"
|
||||
|
||||
if relevant_columns:
|
||||
context += "\n### Columns:\n"
|
||||
for column in sorted(relevant_columns, key=lambda x: -x["score"]):
|
||||
context += f"- {column['name']}: {column['description']} (relevance score: {column['score']:.2f})\n"
|
||||
|
||||
return context
|
||||
@@ -2,6 +2,7 @@ from langchain_community.utilities import SQLDatabase
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from apps.chat.schemas.chat_base_schema import LLMConfig, LLMFactory
|
||||
from apps.chat.schemas.schema_engine import SchemaEngine
|
||||
from apps.datasource.models.datasource import CoreDatasource
|
||||
from apps.db.db import exec_sql, get_uri
|
||||
from common.core.config import settings
|
||||
@@ -103,9 +104,13 @@ class AgentService:
|
||||
async def async_generate(self, question: str) -> AsyncGenerator[str, None]:
|
||||
|
||||
chain = self.prompt | self.agent_executor
|
||||
schema = self.db.get_table_info()
|
||||
# schema = self.db.get_table_info()
|
||||
|
||||
async for chunk in chain.astream({"schema": schema, "question": question}):
|
||||
schema_engine = SchemaEngine(engine=self.db._engine)
|
||||
mschema = schema_engine.mschema
|
||||
mschema_str = mschema.to_mschema()
|
||||
|
||||
async for chunk in chain.astream({"schema": mschema_str, "question": question}):
|
||||
if not isinstance(chunk, dict):
|
||||
continue
|
||||
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
from apps.chat.schemas.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.
|
||||
|
||||
Args:
|
||||
selected_tables: None means select all tables
|
||||
selected_columns: None means select all columns, format: ['table_name.column_name']
|
||||
example_num: Number of examples to show for each field
|
||||
show_type_detail: Whether to show detailed type information
|
||||
|
||||
Returns:
|
||||
Formatted schema string
|
||||
"""
|
||||
output = []
|
||||
|
||||
output.append(f"【DB_ID】 {self.db_id}")
|
||||
output.append(f"【Schema】")
|
||||
|
||||
# Process selected tables and columns
|
||||
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]
|
||||
|
||||
# Process each table
|
||||
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))
|
||||
|
||||
# Add foreign key information, skip when table_type is 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):
|
||||
output.append(f"{fk[0]}.{fk[1]}={fk[3]}.{fk[4]}")
|
||||
|
||||
return '\n'.join(output)
|
||||
|
||||
def dump(self):
|
||||
"""
|
||||
Dump schema information to dictionary
|
||||
|
||||
Returns:
|
||||
Dictionary containing schema information
|
||||
"""
|
||||
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", [])
|
||||
@@ -0,0 +1,126 @@
|
||||
import json, os
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
|
||||
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, select, text
|
||||
from sqlalchemy.engine import Engine
|
||||
from llama_index.core import SQLDatabase
|
||||
from apps.chat.schemas.utils import read_json, write_json, save_raw_text, examples_to_str
|
||||
from apps.chat.schemas.m_schema import MSchema
|
||||
|
||||
|
||||
class SchemaEngine(SQLDatabase):
|
||||
def __init__(self, engine: Engine, schema: Optional[str] = None, metadata: Optional[MetaData] = 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] = ''):
|
||||
super().__init__(engine, schema, metadata, ignore_tables, include_tables, sample_rows_in_table_info,
|
||||
indexes_in_table_info, custom_table_info, view_support, max_string_length)
|
||||
|
||||
self._db_name = db_name
|
||||
# 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
|
||||
else:
|
||||
all_tables = []
|
||||
# Iterate through all available schemas
|
||||
for s in self.get_schema_names():
|
||||
tables = self._inspector.get_table_names(schema=s)
|
||||
all_tables.extend(tables)
|
||||
for table in tables:
|
||||
self._tables_schemas[table] = s
|
||||
self._usable_tables = all_tables """
|
||||
d_schema = self.default_schema_name()
|
||||
for table_name in self._usable_tables:
|
||||
self._tables_schemas[table_name] = d_schema
|
||||
|
||||
self._dialect = engine.dialect.name
|
||||
if mschema is not None:
|
||||
self._mschema = mschema
|
||||
else:
|
||||
self._mschema = MSchema(db_id=db_name)
|
||||
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] )['constrained_columns']
|
||||
|
||||
def get_table_comment(self, table_name: str):
|
||||
try:
|
||||
return self._inspector.get_table_comment(table_name, self._tables_schemas[table_name])['text']
|
||||
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_schema_names()
|
||||
|
||||
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, table_name: str, column_name: str, max_num: int = 5):
|
||||
table = Table(table_name, self.metadata_obj, autoload_with=self._engine, schema=self._tables_schemas[table_name])
|
||||
# Construct SELECT DISTINCT query
|
||||
query = select(table.c[column_name]).distinct().limit(max_num)
|
||||
values = []
|
||||
with self._engine.connect() as connection:
|
||||
result = connection.execute(query)
|
||||
distinct_values = result.fetchall()
|
||||
for value in distinct_values:
|
||||
if value[0] is not None and value[0] != '':
|
||||
values.append(value[0])
|
||||
return values
|
||||
|
||||
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_name, schema=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)
|
||||
if default is not None:
|
||||
default = f'{default}'
|
||||
|
||||
try:
|
||||
examples = self.fectch_distinct_values(table_name, field_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=field['nullable'], default=default, autoincrement=autoincrement,
|
||||
comment=field_comment, examples=examples
|
||||
)
|
||||
@@ -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]
|
||||
Vendored
+4
@@ -11,7 +11,9 @@ declare module 'vue' {
|
||||
ElAvatar: typeof import('element-plus/es')['ElAvatar']
|
||||
ElButton: typeof import('element-plus/es')['ElButton']
|
||||
ElCheckbox: typeof import('element-plus/es')['ElCheckbox']
|
||||
ElCol: typeof import('element-plus/es')['ElCol']
|
||||
ElDialog: typeof import('element-plus/es')['ElDialog']
|
||||
ElDrawer: typeof import('element-plus/es')['ElDrawer']
|
||||
ElDropdown: typeof import('element-plus/es')['ElDropdown']
|
||||
ElDropdownItem: typeof import('element-plus/es')['ElDropdownItem']
|
||||
ElDropdownMenu: typeof import('element-plus/es')['ElDropdownMenu']
|
||||
@@ -23,10 +25,12 @@ declare module 'vue' {
|
||||
ElMenuItem: typeof import('element-plus/es')['ElMenuItem']
|
||||
ElOption: typeof import('element-plus/es')['ElOption']
|
||||
ElPagination: typeof import('element-plus/es')['ElPagination']
|
||||
ElRow: typeof import('element-plus/es')['ElRow']
|
||||
ElSelect: typeof import('element-plus/es')['ElSelect']
|
||||
ElSlider: typeof import('element-plus/es')['ElSlider']
|
||||
ElTable: typeof import('element-plus/es')['ElTable']
|
||||
ElTableColumn: typeof import('element-plus/es')['ElTableColumn']
|
||||
ElTableV2: typeof import('element-plus/es')['ElTableV2']
|
||||
ElTabPane: typeof import('element-plus/es')['ElTabPane']
|
||||
ElTabs: typeof import('element-plus/es')['ElTabs']
|
||||
ElTooltip: typeof import('element-plus/es')['ElTooltip']
|
||||
|
||||
Reference in New Issue
Block a user