ai_site/platform/internal/dbsync/ensure_table.go

package dbsync

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
)

// EnsureTableFromRow 表不存在时按行字段自动建表（TEXT 列 + PK），满足 Z4「按 Binding 接受任意表」。
func EnsureTableFromRow(ctx context.Context, db *sql.DB, driver Driver, table, pkCol string, row map[string]any) error {
	table = strings.TrimSpace(table)
	pkCol = strings.TrimSpace(pkCol)
	if table == "" || pkCol == "" {
		return fmt.Errorf("table and pk required")
	}
	if row == nil {
		row = map[string]any{pkCol: ""}
	}
	cols := make([]string, 0, len(row))
	seen := map[string]struct{}{}
	if _, ok := row[pkCol]; !ok {
		cols = append(cols, pkCol)
		seen[pkCol] = struct{}{}
	}
	for k := range row {
		k = strings.TrimSpace(k)
		if k == "" {
			continue
		}
		if _, ok := seen[k]; ok {
			continue
		}
		seen[k] = struct{}{}
		cols = append(cols, k)
	}
	return EnsureTableFromColumns(ctx, db, driver, table, pkCol, cols)
}

// EnsureTableFromColumns 按列名建空表（全部 TEXT，指定 PK）。
// 已存在时补齐缺失列（ADD COLUMN），解决「先按 _row_id 建表、后 push 注入 id」导致的缺列 503。
func EnsureTableFromColumns(ctx context.Context, db *sql.DB, driver Driver, table, pkCol string, columns []string) error {
	table = strings.TrimSpace(table)
	pkCol = strings.TrimSpace(pkCol)
	if table == "" {
		return fmt.Errorf("table required")
	}
	if strings.HasPrefix(table, "_ajz_") {
		return fmt.Errorf("sync system table not allowed: %s", table)
	}
	if pkCol == "" {
		pkCol = "id"
	}
	exists, err := tableExists(ctx, db, driver, table)
	if err != nil {
		return err
	}
	want := make([]string, 0, len(columns)+1)
	seen := map[string]struct{}{}
	add := func(c string) {
		c = strings.TrimSpace(c)
		if c == "" {
			return
		}
		if _, ok := seen[c]; ok {
			return
		}
		seen[c] = struct{}{}
		want = append(want, c)
	}
	add(pkCol)
	for _, c := range columns {
		add(c)
	}
	if len(want) == 0 {
		return fmt.Errorf("columns required")
	}
	if exists {
		have, err := listColumns(ctx, db, driver, table)
		if err != nil {
			return err
		}
		haveSet := map[string]struct{}{}
		for _, c := range have {
			haveSet[c] = struct{}{}
		}
		for _, c := range want {
			if _, ok := haveSet[c]; ok {
				continue
			}
			ddl := fmt.Sprintf(`ALTER TABLE %s ADD COLUMN %s TEXT`, quoteIdent(driver, table), quoteIdent(driver, c))
			if _, err := db.ExecContext(ctx, ddl); err != nil {
				// SQLite 无 IF NOT EXISTS；并发下可能已存在
				if !strings.Contains(strings.ToLower(err.Error()), "duplicate column") {
					return err
				}
			}
		}
		return nil
	}
	defs := make([]string, 0, len(want))
	defs = append(defs, fmt.Sprintf("%s TEXT PRIMARY KEY", quoteIdent(driver, pkCol)))
	for _, c := range want {
		if c == pkCol {
			continue
		}
		defs = append(defs, fmt.Sprintf("%s TEXT", quoteIdent(driver, c)))
	}
	ddl := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (%s)`, quoteIdent(driver, table), strings.Join(defs, ", "))
	_, err = db.ExecContext(ctx, ddl)
	return err
}

func tableExists(ctx context.Context, db *sql.DB, driver Driver, table string) (bool, error) {
	var q string
	switch driver {
	case DriverSQLite:
		q = `SELECT 1 FROM sqlite_master WHERE type='table' AND name=? LIMIT 1`
	case DriverMySQL:
		q = `SELECT 1 FROM information_schema.tables WHERE table_schema=DATABASE() AND table_name=? LIMIT 1`
	case DriverPostgres:
		q = `SELECT 1 FROM information_schema.tables WHERE table_schema=current_schema() AND table_name=$1 LIMIT 1`
	default:
		return false, fmt.Errorf("unsupported driver")
	}
	var n int
	err := db.QueryRowContext(ctx, q, table).Scan(&n)
	if err == sql.ErrNoRows {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}