opsxdb/server/internal/db/postgres.go

package db

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

	_ "github.com/jackc/pgx/v5/stdlib"
)

func openPostgres(c ConnReq) (*sql.DB, error) {
	host := c.Host
	if host == "" {
		host = "127.0.0.1"
	}
	port := c.Port
	if port == 0 {
		port = 5432
	}
	user := c.User
	if user == "" {
		user = "postgres"
	}
	dbName := c.Database
	if dbName == "" {
		dbName = "postgres"
	}
	mode := EffectiveSSLMode(c)
	sslmode := "disable"
	switch mode {
	case "disable":
		sslmode = "disable"
	case "require", "skip-verify":
		// libpq require ≈ encrypt without CA/hostname verify
		sslmode = "require"
	case "verify-ca":
		sslmode = "verify-ca"
		if strings.TrimSpace(c.SSLCA) == "" {
			if strings.TrimSpace(c.SSLCAContent) != "" {
				return nil, fmt.Errorf("verify-ca 模式需要 CA 证书路径(pg 驱动暂不支持证书内容,请改用 sslCA)")
			}
			return nil, fmt.Errorf("verify-ca 模式需要填写 CA 证书路径")
		}
	default:
		return nil, fmt.Errorf("未知 SSL 模式: %s", mode)
	}

	var b strings.Builder
	fmt.Fprintf(&b, "host=%s port=%d user=%s password=%s dbname=%s sslmode=%s connect_timeout=5",
		pqEscape(host), port, pqEscape(user), pqEscape(c.Password), pqEscape(dbName), sslmode)
	if ca := strings.TrimSpace(c.SSLCA); ca != "" {
		fmt.Fprintf(&b, " sslrootcert=%s", pqEscape(ca))
	}
	if cert := strings.TrimSpace(c.SSLCert); cert != "" {
		fmt.Fprintf(&b, " sslcert=%s", pqEscape(cert))
	}
	if key := strings.TrimSpace(c.SSLKey); key != "" {
		fmt.Fprintf(&b, " sslkey=%s", pqEscape(key))
	}
	return sql.Open("pgx", b.String())
}

func pqEscape(v string) string {
	if v == "" {
		return "''"
	}
	needs := strings.ContainsAny(v, " '\\")
	if !needs {
		return v
	}
	return "'" + strings.ReplaceAll(strings.ReplaceAll(v, `\`, `\\`), `'`, `\'`) + "'"
}

func listDatabasesPostgres(ctx context.Context, db *sql.DB) ([]string, error) {
	rows, err := db.QueryContext(ctx, `SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY 1`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		out = append(out, name)
	}
	return out, rows.Err()
}

func listTablesPostgres(ctx context.Context, db *sql.DB, database string) ([]string, error) {
	// 当前连接库下的 public 表；切库需重新 Open
	_ = database
	rows, err := db.QueryContext(ctx, `
		SELECT CASE WHEN schemaname = 'public' THEN tablename ELSE schemaname || '.' || tablename END
		FROM pg_catalog.pg_tables
		WHERE schemaname NOT IN ('pg_catalog','information_schema')
		ORDER BY 1`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		out = append(out, name)
	}
	return out, rows.Err()
}

func listColumnsPostgres(ctx context.Context, dbh *sql.DB, _, table string) ([]ColumnInfo, error) {
	schema, t := splitSchemaTable(table)
	if schema == "" {
		schema = "public"
	}
	// to_regclass 解析失败(表不存在/跨 schema)返回 NULL 而非报错,主键标记自然落空
	rows, err := dbh.QueryContext(ctx, `
		SELECT c.column_name, c.data_type,
			COALESCE(c.character_maximum_length::text, c.numeric_precision::text, ''),
			c.is_nullable, c.column_default,
			CASE WHEN a.attname IS NOT NULL THEN 'PRI' ELSE '' END
		FROM information_schema.columns c
		LEFT JOIN pg_index i ON i.indrelid = to_regclass(quote_ident($1)||'.'||quote_ident($2)) AND i.indisprimary
		LEFT JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attname = c.column_name AND a.attnum = ANY(i.indkey)
		WHERE c.table_schema = $1 AND c.table_name = $2
		ORDER BY c.ordinal_position`, schema, t)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []ColumnInfo
	for rows.Next() {
		var name, typ, extra, nullable, key string
		var def sql.NullString
		if err := rows.Scan(&name, &typ, &extra, &nullable, &def, &key); err != nil {
			return nil, err
		}
		if extra != "" {
			typ = typ + "(" + extra + ")"
		}
		col := ColumnInfo{Name: name, Type: typ, Nullable: nullable == "YES", Key: key}
		if def.Valid {
			col.Default = def.String
		}
		out = append(out, col)
	}
	if out == nil {
		out = []ColumnInfo{}
	}
	return out, rows.Err()
}