klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(datasource): Support PostgreSQL

Browse Source
This commit is contained in:
junjun
2025-05-20 11:55:28 +08:00
parent e61ead2388
commit 279c4b78cc
7 changed files with 119 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/crud/datasource.py
+23 -20
View File
@@ -1,13 +1,18 @@
from sqlmodel import select
from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, EditObj
import datetime
from common.core.deps import SessionDep
from apps.db.db import get_session, get_tables, get_fields, exec_sql
from sqlalchemy import text, and_
from common.utils.utils import deepcopy_ignore_extra
import json
from typing import List
from ..crud.table import delete_table_by_ds_id, update_table
from sqlalchemy import text, and_
from sqlmodel import select
from apps.datasource.utils.utils import aes_decrypt
from apps.db.db import get_engine, get_tables, get_fields, exec_sql
from common.core.deps import SessionDep
from common.utils.utils import deepcopy_ignore_extra
from ..crud.field import delete_field_by_ds_id, update_field
from ..crud.table import delete_table_by_ds_id, update_table
from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, EditObj, \
DatasourceConf
def get_datasource_list(session: SessionDep) -> CoreDatasource:
@@ -17,17 +22,14 @@ def get_datasource_list(session: SessionDep) -> CoreDatasource:
def check_status(session: SessionDep, ds: CoreDatasource):
conn = get_session(ds)
conn = get_engine(ds)
try:
conn.execute(text("SELECT 1")).scalar()
print("success")
return True
with conn.connect() as connection:
print("success")
return True
except Exception as e:
print("Fail:", e)
return False
finally:
conn.close()
return False
def create_ds(session: SessionDep, create_ds: CreateDatasource):
@@ -175,13 +177,14 @@ def update_table_and_fields(session: SessionDep, data: EditObj):
def preview(session: SessionDep, id: int, data: EditObj):
ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
sql: str = ""
if ds.type == "mysql":
sql = f"""SELECT {", ".join([f.field_name for f in data.fields if f.checked])} FROM {data.table.table_name} LIMIT 100"""
sql = f"""SELECT `{"`, `".join([f.field_name for f in data.fields if f.checked])}` FROM `{data.table.table_name}` LIMIT 100"""
elif ds.type == "sqlServer":
sql = f"""
SELECT {", ".join([f.field_name for f in data.fields if f.checked])} FROM {data.table.table_name}
ORDER BY {data.fields[0].field_name}
OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY
"""
sql = f"""SELECT [{"], [".join([f.field_name for f in data.fields if f.checked])}] FROM [{conf.dbSchema}].[{data.table.table_name}]
ORDER BY [{data.fields[0].field_name}]
OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"""
elif ds.type == "pg":
sql = f"""SELECT "{'", "'.join([f.field_name for f in data.fields if f.checked])}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100"""
return exec_sql(ds, sql)
backend/apps/db/db.py
+57 -4
View File
@@ -1,8 +1,9 @@
import base64
import json
import urllib.parse
from typing import Any
from sqlalchemy import create_engine, text, Result
from sqlalchemy import create_engine, text, Result, Engine
from sqlalchemy.orm import sessionmaker
from apps.datasource.models.datasource import DatasourceConf, CoreDatasource, TableSchema, ColumnSchema
@@ -22,13 +23,28 @@ def get_uri(ds: CoreDatasource):
db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{urllib.parse.quote(conf.extraJdbc)}"
else:
db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
elif ds.type == "pg":
if conf.extraJdbc is not None and conf.extraJdbc != '':
db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{urllib.parse.quote(conf.extraJdbc)}"
else:
db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
else:
raise 'The datasource type not support.'
return db_url
def get_engine(ds: CoreDatasource) -> Engine:
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
if ds.type == "pg" and (conf.dbSchema is not None and conf.dbSchema != ""):
engine = create_engine(get_uri(ds), connect_args={"options": f"-c search_path={conf.dbSchema}"},
pool_timeout=30, pool_size=20, max_overflow=10)
else:
engine = create_engine(get_uri(ds), pool_timeout=30, pool_size=20, max_overflow=10)
return engine
def get_session(ds: CoreDatasource):
engine = create_engine(get_uri(ds))
engine = get_engine(ds)
session_maker = sessionmaker(bind=engine)
session = session_maker()
return session
@@ -66,6 +82,24 @@ def get_tables(ds: CoreDatasource):
t.TABLE_TYPE IN ('BASE TABLE', 'VIEW')
AND t.TABLE_SCHEMA = '{conf.dbSchema}';
"""
elif ds.type == "pg":
sql = """
SELECT
c.relname AS TABLE_NAME,
COALESCE(d.description, obj_description(c.oid)) AS TABLE_COMMENT
FROM
pg_class c
LEFT JOIN
pg_namespace n ON n.oid = c.relnamespace
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')
AND c.relname NOT LIKE 'pg_%'
AND c.relname NOT LIKE 'sql_%'
ORDER BY c.relname;
"""
result = session.execute(text(sql))
res = result.fetchall()
@@ -113,7 +147,26 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
WHERE
C.TABLE_SCHEMA = '{conf.dbSchema}'
"""
sql2 = f"AND C.TABLE_NAME = '{table_name}';" if table_name is not None and table_name != "" else ";"
sql2 = f" AND C.TABLE_NAME = '{table_name}';" if table_name is not None and table_name != "" else ";"
sql = sql1 + sql2
elif ds.type == "pg":
sql1 = """
SELECT
a.attname AS COLUMN_NAME,
pg_catalog.format_type(a.atttypid, a.atttypmod) AS DATA_TYPE,
col_description(c.oid, a.attnum) AS COLUMN_COMMENT
FROM
pg_catalog.pg_attribute a
JOIN
pg_catalog.pg_class c ON a.attrelid = c.oid
JOIN
pg_catalog.pg_namespace n ON n.oid = c.relnamespace
WHERE
n.nspname = current_schema()
AND a.attnum > 0
AND NOT a.attisdropped
"""
sql2 = f" AND c.relname = '{table_name}';" if table_name is not None and table_name != "" else ";"
sql = sql1 + sql2
result = session.execute(text(sql))
@@ -137,7 +190,7 @@ def exec_sql(ds: CoreDatasource, sql: str):
{columns[i]: value for i, value in enumerate(tuple_item)}
for tuple_item in res
]
return {"fields": columns, "data": result_list}
return {"fields": columns, "data": result_list, "sql": base64.b64encode(bytes(sql, 'utf-8'))}
finally:
if result is not None:
result.close()
backend/pyproject.toml
+2 -1
View File
@@ -31,7 +31,8 @@ dependencies = [
"llama_index>=0.12.35",
"pymssql (>=2.3.4,<3.0.0)",
"pandas (>=2.2.3,<3.0.0)",
"openpyxl (>=3.1.5,<4.0.0)"
"openpyxl (>=3.1.5,<4.0.0)",
"psycopg2 (>=2.9.10,<3.0.0)"
]
[[tool.uv.index]]
url = "https://pypi.tuna.tsinghua.edu.cn/simple"
frontend/src/assets/svg/ds/pg-ds.svg
+13
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 20 KiB

frontend/src/views/ds/form.vue
+2 -6
View File
@@ -65,7 +65,7 @@
<el-form-item label="Extra JDBC String">
<el-input v-model="config.extraJdbc" />
</el-form-item>
<el-form-item label="Schema" v-if="form.type === 'sqlServer'">
<el-form-item label="Schema" v-if="haveSchema.includes(form.type)">
<el-input v-model="config.dbSchema" />
<el-button link type="primary" :icon="Plus" v-if="false">Get Schema</el-button>
</el-form-item>
@@ -101,6 +101,7 @@ import { ElMessage } from 'element-plus'
import type { FormInstance, FormRules } from 'element-plus'
import { Plus } from '@element-plus/icons-vue'
import { useCache } from '@/utils/useCache'
import { dsType, haveSchema} from '@/views/ds/js/ds-type'
const { wsCache } = useCache()
const dsFormRef = ref<FormInstance>()
@@ -123,11 +124,6 @@ const rules = reactive<FormRules>({
})
const dialogVisible = ref<boolean>(false)
const dsType = [
{label:"Excel/CSV", value:"excel"},
{label:"MySQL", value:"mysql"},
{label:"SQL Server", value:"sqlServer"}
]
const form = ref<any>({
name:'',
description:'',
frontend/src/views/ds/index.vue
+4 -1
View File
@@ -28,11 +28,12 @@
<Icon>
<MysqlDs v-if="ds.type === 'mysql'"/>
<SQLServerDs v-else-if="ds.type === 'sqlServer'"/>
<PgDs v-else-if="ds.type === 'pg'"/>
</Icon>
</div>
<div class="connection-details">
<div class="connection-name">{{ ds.name }}</div>
<div class="connection-type">{{ ds.type }}</div>
<div class="connection-type">{{ dsRelation[ds.type] }}</div>
<div class="connection-host">{{ ds.description }}</div>
<div class="connection-last">{{ datetimeFormat(ds.create_time) }}</div>
</div>
@@ -54,6 +55,7 @@ import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg'
import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg'
import MysqlDs from '@/assets/svg/ds/mysql-ds.svg'
import SQLServerDs from '@/assets/svg/ds/sqlServer-ds.svg'
import PgDs from '@/assets/svg/ds/pg-ds.svg'
import { Search, List, CreditCard } from '@element-plus/icons-vue'
import { ref, onMounted } from 'vue'
import DsForm from './form.vue'
@@ -61,6 +63,7 @@ import { datasourceApi } from '@/api/datasource'
import { datetimeFormat } from '@/utils/utils'
import { ElMessageBox } from 'element-plus'
import { useRouter } from 'vue-router'
import { dsRelation } from '@/views/ds/js/ds-type'
const searchValue = ref<string>('')
const dsForm = ref()
frontend/src/views/ds/js/ds-type.ts
+18
View File
@@ -0,0 +1,18 @@
export const dsType = [
{label:"MySQL", value:"mysql"},
{label:"Microsoft SQL Server", value:"sqlServer"},
{label:"PostgreSQL", value:"pg"},
{label:"Excel/CSV", value:"excel"},
]
export const haveSchema = [
'sqlServer',
'pg'
]
export const dsRelation: any = {
"mysql": "MySQL",
"sqlServer": "Microsoft SQL Server",
"pg": "PostgreSQL",
"excel": "Excel/CSV",
}