kra-new/internal/data/export.go

309 lines
10 KiB
Go

package data
import (
"context"
"errors"
"fmt"
"gorm.io/gorm"
"kra/internal/biz"
"regexp"
"sort"
"strconv"
"strings"
"time"
)
type exportTemplatePO struct {
ID uint `gorm:"primaryKey"`
CreatedAt time.Time
UpdatedAt time.Time
DeletedAt gorm.DeletedAt `gorm:"index"`
DBName string
Name string
DBTableName string `gorm:"column:table_name"`
TemplateID string `gorm:"uniqueIndex"`
TemplateInfo string `gorm:"type:text"`
SQL string `gorm:"type:text"`
ImportSQL string `gorm:"type:text"`
Limit *int
Order string
}
func (exportTemplatePO) TableName() string { return "sys_export_templates" }
type exportConditionPO struct {
ID uint `gorm:"primaryKey"`
TemplateID string `gorm:"index"`
From string
Column string
Operator string
}
func (exportConditionPO) TableName() string { return "sys_export_template_condition" }
type exportJoinPO struct {
ID uint `gorm:"primaryKey"`
TemplateID string `gorm:"index"`
Join string `gorm:"column:joins"`
Table string
On string `gorm:"column:on"`
}
func (exportJoinPO) TableName() string { return "sys_export_template_join" }
type exportRepo struct{ data *Data }
func NewExportRepo(data *Data) biz.ExportRepo { return &exportRepo{data: data} }
func exportFromPO(po exportTemplatePO, conditions []exportConditionPO, joins []exportJoinPO) *biz.ExportTemplate {
v := &biz.ExportTemplate{ID: po.ID, CreatedAt: po.CreatedAt, UpdatedAt: po.UpdatedAt, DBName: po.DBName, Name: po.Name, TableName: po.DBTableName, TemplateID: po.TemplateID, TemplateInfo: po.TemplateInfo, SQL: po.SQL, ImportSQL: po.ImportSQL, Limit: po.Limit, Order: po.Order}
for _, x := range conditions {
v.Conditions = append(v.Conditions, biz.ExportCondition{From: x.From, Column: x.Column, Operator: x.Operator})
}
for _, x := range joins {
v.Joins = append(v.Joins, biz.ExportJoin{Join: x.Join, Table: x.Table, On: x.On})
}
return v
}
func exportToPO(v *biz.ExportTemplate) exportTemplatePO {
return exportTemplatePO{ID: v.ID, DBName: v.DBName, Name: v.Name, DBTableName: v.TableName, TemplateID: v.TemplateID, TemplateInfo: v.TemplateInfo, SQL: v.SQL, ImportSQL: v.ImportSQL, Limit: v.Limit, Order: v.Order}
}
func (r *exportRepo) saveRelations(tx *gorm.DB, v *biz.ExportTemplate) error {
if err := tx.Where("template_id = ?", v.TemplateID).Delete(&exportConditionPO{}).Error; err != nil {
return err
}
if err := tx.Where("template_id = ?", v.TemplateID).Delete(&exportJoinPO{}).Error; err != nil {
return err
}
conditions := make([]exportConditionPO, 0, len(v.Conditions))
for _, x := range v.Conditions {
conditions = append(conditions, exportConditionPO{TemplateID: v.TemplateID, From: x.From, Column: x.Column, Operator: x.Operator})
}
if len(conditions) > 0 {
if err := tx.Create(&conditions).Error; err != nil {
return err
}
}
joins := make([]exportJoinPO, 0, len(v.Joins))
for _, x := range v.Joins {
joins = append(joins, exportJoinPO{TemplateID: v.TemplateID, Join: x.Join, Table: x.Table, On: x.On})
}
if len(joins) > 0 {
return tx.Create(&joins).Error
}
return nil
}
func (r *exportRepo) CreateExportTemplate(ctx context.Context, v *biz.ExportTemplate) error {
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
po := exportToPO(v)
if err := tx.Create(&po).Error; err != nil {
return err
}
v.ID = po.ID
return r.saveRelations(tx, v)
})
}
func (r *exportRepo) UpdateExportTemplate(ctx context.Context, v *biz.ExportTemplate) error {
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
po := exportToPO(v)
if err := tx.Model(&exportTemplatePO{}).Where("id = ?", v.ID).Updates(&po).Error; err != nil {
return err
}
return r.saveRelations(tx, v)
})
}
func (r *exportRepo) DeleteExportTemplates(ctx context.Context, ids []uint) error {
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
var tids []string
if err := tx.Model(&exportTemplatePO{}).Where("id IN ?", ids).Pluck("template_id", &tids).Error; err != nil {
return err
}
if len(tids) > 0 {
if err := tx.Where("template_id IN ?", tids).Delete(&exportConditionPO{}).Error; err != nil {
return err
}
if err := tx.Where("template_id IN ?", tids).Delete(&exportJoinPO{}).Error; err != nil {
return err
}
}
return tx.Delete(&exportTemplatePO{}, ids).Error
})
}
func (r *exportRepo) FindExportTemplate(ctx context.Context, id uint, tid string) (*biz.ExportTemplate, error) {
var po exportTemplatePO
db := r.data.gormDB.WithContext(ctx)
var err error
if id != 0 {
err = db.First(&po, id).Error
} else {
err = db.Where("template_id = ?", tid).First(&po).Error
}
if err != nil {
return nil, err
}
var conditions []exportConditionPO
var joins []exportJoinPO
if err = db.Where("template_id = ?", po.TemplateID).Find(&conditions).Error; err != nil {
return nil, err
}
if err = db.Where("template_id = ?", po.TemplateID).Find(&joins).Error; err != nil {
return nil, err
}
return exportFromPO(po, conditions, joins), nil
}
func (r *exportRepo) ListExportTemplates(ctx context.Context, page, size int, q *biz.ExportTemplate) ([]*biz.ExportTemplate, int64, error) {
if page < 1 {
page = 1
}
if size < 1 {
size = 10
}
db := r.data.gormDB.WithContext(ctx).Model(&exportTemplatePO{})
if q != nil {
if q.Name != "" {
db = db.Where("name LIKE ?", "%"+q.Name+"%")
}
if q.TableName != "" {
db = db.Where("table_name = ?", q.TableName)
}
if q.TemplateID != "" {
db = db.Where("template_id = ?", q.TemplateID)
}
}
var total int64
if err := db.Count(&total).Error; err != nil {
return nil, 0, err
}
var pos []exportTemplatePO
if err := db.Order("id desc").Offset((page - 1) * size).Limit(size).Find(&pos).Error; err != nil {
return nil, 0, err
}
out := make([]*biz.ExportTemplate, 0, len(pos))
for _, po := range pos {
v, err := r.FindExportTemplate(ctx, po.ID, "")
if err != nil {
return nil, 0, err
}
out = append(out, v)
}
return out, total, nil
}
var safeIdentifier = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_.]*$`)
var allowedOperator = map[string]bool{"=": true, "!=": true, ">": true, ">=": true, "<": true, "<=": true, "LIKE": true, "IN": true, "NOT IN": true, "BETWEEN": true}
func validateSelectSQL(sql string) error {
normalized := strings.ToLower(strings.TrimSpace(sql))
if !strings.HasPrefix(normalized, "select ") || strings.Contains(normalized, ";") {
return errors.New("仅允许单条 SELECT 查询")
}
return nil
}
func (r *exportRepo) QueryExport(ctx context.Context, t *biz.ExportTemplate, params map[string]string) ([]map[string]any, string, error) {
db := r.data.gormDB.WithContext(ctx)
var rows []map[string]any
if t.SQL != "" {
if err := validateSelectSQL(t.SQL); err != nil {
return nil, "", err
}
args := map[string]any{}
for k, v := range params {
args[k] = v
}
err := db.Raw(t.SQL, args).Scan(&rows).Error
return rows, t.SQL, err
}
if !safeIdentifier.MatchString(t.TableName) {
return nil, "", errors.New("表名不合法")
}
columns := make([]string, 0)
for key := range parseTemplateColumns(t.TemplateInfo) {
if !safeIdentifier.MatchString(key) {
return nil, "", errors.New("列名不合法")
}
columns = append(columns, key)
}
sort.Strings(columns)
if len(columns) == 0 {
return nil, "", errors.New("模板列为空")
}
query := db.Table(t.TableName).Select(strings.Join(columns, ","))
for _, j := range t.Joins {
joinType := strings.ToUpper(strings.TrimSpace(j.Join))
if joinType != "LEFT JOIN" && joinType != "RIGHT JOIN" && joinType != "INNER JOIN" && joinType != "JOIN" {
return nil, "", errors.New("关联类型不合法")
}
if !safeIdentifier.MatchString(j.Table) {
return nil, "", errors.New("关联表不合法")
}
query = query.Joins(joinType + " " + j.Table + " ON " + j.On)
}
for _, condition := range t.Conditions {
operator := strings.ToUpper(strings.TrimSpace(condition.Operator))
if !allowedOperator[operator] || !safeIdentifier.MatchString(condition.Column) {
return nil, "", errors.New("查询条件不合法")
}
value := params[condition.From]
if value == "" {
continue
}
switch operator {
case "LIKE":
query = query.Where(condition.Column+" LIKE ?", "%"+value+"%")
case "IN", "NOT IN":
query = query.Where(condition.Column+" "+operator+" ?", strings.Split(value, ","))
case "BETWEEN":
start, end := params["start"+condition.From], params["end"+condition.From]
if start != "" && end != "" {
query = query.Where(condition.Column+" BETWEEN ? AND ?", start, end)
}
default:
query = query.Where(condition.Column+" "+operator+" ?", value)
}
}
limit := 0
if raw := params["limit"]; raw != "" {
limit, _ = strconv.Atoi(raw)
} else if t.Limit != nil {
limit = *t.Limit
}
if limit > 0 {
query = query.Limit(limit)
}
if t.Order != "" {
parts := strings.Fields(t.Order)
if len(parts) < 1 || len(parts) > 2 || !safeIdentifier.MatchString(parts[0]) || (len(parts) == 2 && strings.ToUpper(parts[1]) != "ASC" && strings.ToUpper(parts[1]) != "DESC") {
return nil, "", errors.New("排序不合法")
}
query = query.Order(t.Order)
}
preview := query.ToSQL(func(tx *gorm.DB) *gorm.DB { return tx.Find(&[]map[string]any{}) })
err := query.Find(&rows).Error
return rows, preview, err
}
func parseTemplateColumns(raw string) map[string]string {
out := map[string]string{}
parts := strings.Split(strings.Trim(strings.TrimSpace(raw), "{}"), ",")
for _, part := range parts {
pair := strings.SplitN(part, ":", 2)
if len(pair) == 2 {
out[strings.Trim(strings.TrimSpace(pair[0]), `"`)] = strings.Trim(strings.TrimSpace(pair[1]), `"`)
}
}
return out
}
func (r *exportRepo) ImportExportRows(ctx context.Context, t *biz.ExportTemplate, rows []map[string]any) error {
sql := strings.TrimSpace(t.ImportSQL)
lower := strings.ToLower(sql)
if sql == "" || strings.Contains(sql, ";") || (!strings.HasPrefix(lower, "insert ") && !strings.HasPrefix(lower, "update ")) {
return errors.New("导入 SQL 仅允许单条 INSERT 或 UPDATE")
}
return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
for index, row := range rows {
if err := tx.Exec(sql, row).Error; err != nil {
return fmt.Errorf("第%d行导入失败: %w", index+2, err)
}
}
return nil
})
}