ai_site/platform/internal/schema/runner.go

package schema

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"net/url"
	"regexp"
	"strings"
	"unicode"
)

var dbNameRe = regexp.MustCompile(`^[a-z][a-z0-9_]{1,47}$`)

type Runner interface {
	ExecDDL(ctx context.Context, stmts []string) error
	EnsureDatabase(ctx context.Context, dbName string) error
}

type NoopRunner struct{}

func (NoopRunner) ExecDDL(context.Context, []string) error   { return nil }
func (NoopRunner) EnsureDatabase(context.Context, string) error { return nil }

type PostgresRunner struct {
	DB          *sql.DB
	AdminDSN    string // 用于 CREATE DATABASE（连到 postgres 库）
}

func (r *PostgresRunner) ExecDDL(ctx context.Context, stmts []string) error {
	tx, err := r.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	for _, s := range stmts {
		if _, err := tx.ExecContext(ctx, s); err != nil {
			return fmt.Errorf("ddl failed: %w\nsql: %s", err, s)
		}
	}
	return tx.Commit()
}

func (r *PostgresRunner) EnsureDatabase(ctx context.Context, dbName string) error {
	if !dbNameRe.MatchString(dbName) {
		return fmt.Errorf("invalid database name: %s", dbName)
	}
	admin := r.DB
	var err error
	if r.AdminDSN != "" {
		admin, err = sql.Open("postgres", r.AdminDSN)
		if err != nil {
			return err
		}
		defer admin.Close()
	}
	var exists bool
	if err := admin.QueryRowContext(ctx,
		`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname=$1)`, dbName,
	).Scan(&exists); err != nil {
		return err
	}
	if exists {
		return nil
	}
	// CREATE DATABASE 不能在事务中
	_, err = admin.ExecContext(ctx, fmt.Sprintf(`CREATE DATABASE %s`, quoteIdent(dbName)))
	return err
}

// DSNForDatabase 把原 DSN 的库名替换为目标库。
func DSNForDatabase(baseDSN, dbName string) (string, error) {
	if !dbNameRe.MatchString(dbName) {
		return "", fmt.Errorf("invalid database name")
	}
	u, err := url.Parse(baseDSN)
	if err != nil {
		return "", err
	}
	u.Path = "/" + dbName
	return u.String(), nil
}

func QuoteIdentExport(name string) string { return quoteIdent(name) }

func SanitizeDBName(tenantID int64, slug string) string {
	name := fmt.Sprintf("appdb_t%d_%s", tenantID, slug)
	name = strings.ToLower(name)
	if len(name) > 48 {
		name = name[:48]
	}
	return name
}

// ValidDBName 是否可作为 Postgres 物理库名（CREATE DATABASE）。
func ValidDBName(dbName string) bool {
	return dbNameRe.MatchString(strings.TrimSpace(dbName))
}

// SafePhysicalDBName 优先用 preferred；非法则用 fallback；仍非法则 user_fallback。
func SafePhysicalDBName(preferred, fallback string) string {
	p := strings.TrimSpace(preferred)
	if ValidDBName(p) {
		return p
	}
	f := strings.TrimSpace(fallback)
	if ValidDBName(f) {
		return f
	}
	return "appdb_default"
}

// SafeSyncDatabaseName Z35：智能体同步落点的 Postgres 物理库名。
// - 已是合法 ASCII → 原样
// - 中文/非法串：有 userID → user_{id}（昵称可变仍稳定）；否则对标签做稳定哈希 db{hex}
// - 展示中文请写 Binding.DisplayName / Binding.DatabaseName（可读名），勿把物理名当展示名
func SafeSyncDatabaseName(preferred string, agentID, userID int64) string {
	p := strings.TrimSpace(preferred)
	if ValidDBName(p) {
		return p
	}
	// 有登录用户：物理名跟用户稳定绑定；中文昵称只作展示
	if userID > 0 {
		n := fmt.Sprintf("user_%d", userID)
		if ValidDBName(n) {
			return n
		}
	}
	if p != "" {
		return HashLabelDBName(p, agentID)
	}
	if agentID > 0 {
		n := fmt.Sprintf("agent_%d", agentID)
		if ValidDBName(n) {
			return n
		}
	}
	return "appdb_default"
}

// HashLabelDBName 将任意标签（含中文）稳定映射为合法物理库名：db + 16 hex。
func HashLabelDBName(label string, agentID int64) string {
	label = strings.TrimSpace(label)
	sum := sha256.Sum256([]byte(fmt.Sprintf("ajz-db:%d:%s", agentID, label)))
	return "db" + hex.EncodeToString(sum[:8])
}

// IsDisplayLabel 是否更像「展示名」而非物理库名（含非 ASCII / 大写 / 连字符等）。
func IsDisplayLabel(s string) bool {
	s = strings.TrimSpace(s)
	if s == "" || ValidDBName(s) {
		return false
	}
	for _, r := range s {
		if r > unicode.MaxASCII || unicode.Is(unicode.Han, r) {
			return true
		}
	}
	return true
}