chore: 恢复旧布局基线（迁移改动丢失后重建，apps/packages 布局暂存）
This commit is contained in:
logan committed
1 parent
a9510235bb
commit
456fa6954b
203 files changed
+18787
-48
No files matched your search
@@ -0,0 +1,252 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/csv"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultExportLimit = 100000
|
||||
MaxExportLimit = 500000
|
||||
exportFlushRows = 512
|
||||
)
|
||||
|
||||
// ExportAllowedQuery 导出专用的只读校验:首个关键词须为 SELECT/WITH/SHOW/DESCRIBE/EXPLAIN。
|
||||
func ExportAllowedQuery(sqlText string) bool {
|
||||
s := stripSQLLead(strings.TrimSpace(sqlText))
|
||||
if s == "" {
|
||||
return false
|
||||
}
|
||||
up := strings.ToUpper(s)
|
||||
for _, p := range []string{"SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN"} {
|
||||
if strings.HasPrefix(up, p+" ") || up == p || strings.HasPrefix(up, p+"\t") || strings.HasPrefix(up, p+"\n") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// SelectAllSQL 表模式导出的取数 SQL:按 family 引用(支持 schema.table 多段)。
|
||||
func SelectAllSQL(driver, table string) string {
|
||||
return "SELECT * FROM " + quoteTableNameFor(driver, table)
|
||||
}
|
||||
|
||||
// ExportQuery 把查询结果按 c.Format(csv/json/insert)流式编码写入 w,避免全量载入内存。
|
||||
// limit <= 0 用默认上限 100000,硬顶 MaxExportLimit。行值经 normalize 保真(同 ExecSQL)。
|
||||
func ExportQuery(ctx context.Context, c ConnReq, w io.Writer) error {
|
||||
ctx = reqCtx(ctx)
|
||||
format := strings.ToLower(strings.TrimSpace(c.Format))
|
||||
sqlText := strings.TrimSpace(c.SQL)
|
||||
if sqlText == "" {
|
||||
return fmt.Errorf("缺少导出 SQL")
|
||||
}
|
||||
limit := c.Limit
|
||||
if limit <= 0 {
|
||||
limit = DefaultExportLimit
|
||||
}
|
||||
if limit > MaxExportLimit {
|
||||
limit = MaxExportLimit
|
||||
}
|
||||
return withDB(c, func(db *sql.DB) error {
|
||||
rows, err := db.QueryContext(ctx, sqlText)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
bw := bufio.NewWriterSize(w, 64*1024)
|
||||
enc := newExportEncoder(bw, format, c.Driver, exportTargetTable(c), cols)
|
||||
if err := enc.begin(); err != nil {
|
||||
return err
|
||||
}
|
||||
n := 0
|
||||
for rows.Next() {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
raw := make([]interface{}, len(cols))
|
||||
ptrs := make([]interface{}, len(cols))
|
||||
for i := range raw {
|
||||
ptrs[i] = &raw[i]
|
||||
}
|
||||
if err := rows.Scan(ptrs...); err != nil {
|
||||
return err
|
||||
}
|
||||
vals := make([]interface{}, len(cols))
|
||||
for i := range raw {
|
||||
vals[i] = normalize(raw[i])
|
||||
}
|
||||
if err := enc.row(vals); err != nil {
|
||||
return err
|
||||
}
|
||||
n++
|
||||
if n%exportFlushRows == 0 {
|
||||
if err := bw.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if n >= limit {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := enc.end(); err != nil {
|
||||
return err
|
||||
}
|
||||
return bw.Flush()
|
||||
})
|
||||
}
|
||||
|
||||
// exportTargetTable insert 格式的目标表名:显式 table 优先,否则用通用名。
|
||||
func exportTargetTable(c ConnReq) string {
|
||||
if strings.TrimSpace(c.Table) != "" {
|
||||
return c.Table
|
||||
}
|
||||
return "result"
|
||||
}
|
||||
|
||||
// exportEncoder 按格式把行流式编码进 buffered writer。
|
||||
type exportEncoder struct {
|
||||
bw *bufio.Writer
|
||||
format string
|
||||
driver string
|
||||
table string
|
||||
cols []string
|
||||
csvW *csv.Writer
|
||||
// jsonPending 上一行 JSON,延迟一行写出,使行间逗号紧跟行尾而非独占一行
|
||||
jsonPending string
|
||||
n int
|
||||
}
|
||||
|
||||
func newExportEncoder(bw *bufio.Writer, format, driver, table string, cols []string) *exportEncoder {
|
||||
return &exportEncoder{bw: bw, format: format, driver: driver, table: table, cols: cols}
|
||||
}
|
||||
|
||||
func (e *exportEncoder) begin() error {
|
||||
switch e.format {
|
||||
case "csv":
|
||||
e.csvW = csv.NewWriter(e.bw)
|
||||
return e.csvW.Write(e.cols)
|
||||
case "json":
|
||||
_, err := e.bw.WriteString("[\n")
|
||||
return err
|
||||
case "insert":
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("不支持的导出格式: %s", e.format)
|
||||
}
|
||||
}
|
||||
|
||||
func (e *exportEncoder) row(vals []interface{}) error {
|
||||
e.n++
|
||||
switch e.format {
|
||||
case "csv":
|
||||
rec := make([]string, len(vals))
|
||||
for i, v := range vals {
|
||||
rec[i] = csvCell(v)
|
||||
}
|
||||
return e.csvW.Write(rec)
|
||||
case "json":
|
||||
m := make(map[string]interface{}, len(e.cols))
|
||||
for i, col := range e.cols {
|
||||
m[col] = vals[i]
|
||||
}
|
||||
var rb bytes.Buffer
|
||||
if err := json.NewEncoder(&rb).Encode(m); err != nil {
|
||||
return err
|
||||
}
|
||||
line := strings.TrimSuffix(rb.String(), "\n")
|
||||
if e.jsonPending != "" {
|
||||
if _, err := e.bw.WriteString(e.jsonPending + ",\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
e.jsonPending = line
|
||||
return nil
|
||||
case "insert":
|
||||
var b strings.Builder
|
||||
b.WriteString("INSERT INTO ")
|
||||
b.WriteString(quoteTableNameFor(e.driver, e.table))
|
||||
b.WriteString(" (")
|
||||
for i, col := range e.cols {
|
||||
if i > 0 {
|
||||
b.WriteString(", ")
|
||||
}
|
||||
b.WriteString(quoteIdentFor(e.driver, col))
|
||||
}
|
||||
b.WriteString(") VALUES (")
|
||||
for i, v := range vals {
|
||||
if i > 0 {
|
||||
b.WriteString(", ")
|
||||
}
|
||||
b.WriteString(sqlLiteral(v))
|
||||
}
|
||||
b.WriteString(");\n")
|
||||
_, err := e.bw.WriteString(b.String())
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *exportEncoder) end() error {
|
||||
switch e.format {
|
||||
case "csv":
|
||||
e.csvW.Flush()
|
||||
return e.csvW.Error()
|
||||
case "json":
|
||||
if e.jsonPending != "" {
|
||||
if _, err := e.bw.WriteString(e.jsonPending + "\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
e.jsonPending = ""
|
||||
}
|
||||
_, err := e.bw.WriteString("]\n")
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func csvCell(v interface{}) string {
|
||||
switch t := v.(type) {
|
||||
case nil:
|
||||
return ""
|
||||
case string:
|
||||
return t
|
||||
default:
|
||||
return fmt.Sprintf("%v", t)
|
||||
}
|
||||
}
|
||||
|
||||
// sqlLiteral insert 字面量:NULL/数字原样,字符串单引号 doubling。
|
||||
func sqlLiteral(v interface{}) string {
|
||||
switch t := v.(type) {
|
||||
case nil:
|
||||
return "NULL"
|
||||
case string:
|
||||
return "'" + strings.ReplaceAll(t, "'", "''") + "'"
|
||||
case bool:
|
||||
if t {
|
||||
return "1"
|
||||
}
|
||||
return "0"
|
||||
case time.Time:
|
||||
return "'" + t.Format(time.RFC3339Nano) + "'"
|
||||
case []byte:
|
||||
return "'" + strings.ReplaceAll(string(t), "'", "''") + "'"
|
||||
default:
|
||||
return fmt.Sprintf("%v", t)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user