opsxdb/server/internal/db/mssql.go

package db

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

	_ "github.com/microsoft/go-mssqldb"
)

func openMSSQL(c ConnReq) (*sql.DB, error) {
	dsn, err := mssqlDSN(c)
	if err != nil {
		return nil, err
	}
	return sql.Open("sqlserver", dsn)
}

// mssqlDSN 拼 sqlserver:// DSN;账号密码走 url.UserPassword 编码,
// 修复含空格 / '+' 等特殊字符密码连接失败的问题(url.QueryEscape 不适用 userinfo)。
func mssqlDSN(c ConnReq) (string, error) {
	host := c.Host
	if host == "" {
		host = "127.0.0.1"
	}
	port := c.Port
	if port == 0 {
		port = 1433
	}
	user := c.User
	if user == "" {
		user = "sa"
	}
	dbName := c.Database
	if dbName == "" {
		dbName = "master"
	}
	mode := EffectiveSSLMode(c)
	q := url.Values{}
	q.Set("database", dbName)
	q.Set("connection timeout", "5")
	switch mode {
	case "disable":
		q.Set("encrypt", "disable")
	case "skip-verify":
		q.Set("encrypt", "true")
		q.Set("TrustServerCertificate", "true")
		if ca := strings.TrimSpace(c.SSLCA); ca != "" {
			q.Set("certificate", ca)
		}
	case "require", "verify-ca":
		q.Set("encrypt", "true")
		q.Set("TrustServerCertificate", "false")
		ca := strings.TrimSpace(c.SSLCA)
		if mode == "verify-ca" && ca == "" {
			return "", fmt.Errorf("verify-ca 模式需要填写 CA 证书路径")
		}
		if ca != "" {
			// go-mssqldb msdsn certificate= builds RootCAs from PEM/DER
			q.Set("certificate", ca)
		}
		if host != "" {
			q.Set("hostnameincertificate", host)
		}
	default:
		return "", fmt.Errorf("未知 SSL 模式: %s", mode)
	}
	u := &url.URL{
		Scheme:   "sqlserver",
		User:     url.UserPassword(user, c.Password),
		Host:     fmt.Sprintf("%s:%d", host, port),
		RawQuery: q.Encode(),
	}
	return u.String(), nil
}

func listDatabasesMSSQL(ctx context.Context, db *sql.DB) ([]string, error) {
	rows, err := db.QueryContext(ctx, `SELECT name FROM sys.databases WHERE name NOT IN ('master','tempdb','model','msdb') ORDER BY name`)
	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 listTablesMSSQL(ctx context.Context, db *sql.DB, database string) ([]string, error) {
	_ = database
	rows, err := db.QueryContext(ctx, `
		SELECT TABLE_SCHEMA + '.' + TABLE_NAME
		FROM INFORMATION_SCHEMA.TABLES
		WHERE TABLE_TYPE = 'BASE TABLE'
		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 listColumnsMSSQL(ctx context.Context, dbh *sql.DB, _, table string) ([]ColumnInfo, error) {
	schema, t := splitSchemaTable(table)
	if schema == "" {
		schema = "dbo"
	}
	// 关联 sys.indexes / sys.index_columns 补主键标记(命中则 Key='PRI')
	rows, err := dbh.QueryContext(ctx, `
		SELECT c.COLUMN_NAME, c.DATA_TYPE,
			ISNULL(CAST(c.CHARACTER_MAXIMUM_LENGTH AS varchar), ''),
			c.IS_NULLABLE, c.COLUMN_DEFAULT,
			CASE WHEN pk.name IS NOT NULL THEN 'PRI' ELSE '' END
		FROM INFORMATION_SCHEMA.COLUMNS c
		LEFT JOIN sys.tables t ON t.name = c.TABLE_NAME AND SCHEMA_NAME(t.schema_id) = c.TABLE_SCHEMA
		LEFT JOIN sys.indexes i ON i.object_id = t.object_id AND i.is_primary_key = 1
		LEFT JOIN sys.index_columns ic ON ic.object_id = i.object_id AND ic.index_id = i.index_id
		LEFT JOIN sys.columns pk ON pk.object_id = ic.object_id AND pk.column_id = ic.column_id AND pk.name = c.COLUMN_NAME
		WHERE c.TABLE_SCHEMA = @p1 AND c.TABLE_NAME = @p2
		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 != "" && extra != "-1" {
			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()
}