xinmi/internal/application/service/vectorstore_healthcheck.go

package service

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	openSearchRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/opensearch"
	"github.com/Tencent/WeKnora/internal/errors"
	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/types"
	secutils "github.com/Tencent/WeKnora/internal/utils"
	"github.com/go-sql-driver/mysql"   // MySQL driver for database/sql, used by Doris connection test
	_ "github.com/jackc/pgx/v5/stdlib" // pgx driver for database/sql
	"github.com/qdrant/go-client/qdrant"
	"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
	"github.com/weaviate/weaviate-go-client/v5/weaviate"
	"github.com/weaviate/weaviate-go-client/v5/weaviate/auth"
	wgrpc "github.com/weaviate/weaviate-go-client/v5/weaviate/grpc"
	"google.golang.org/grpc"
)

const connectionTestTimeout = 10 * time.Second

// TestConnection tests connectivity to a vector database.
// Returns the detected server version on success (e.g., "7.10.1"), empty string if unknown.
func (s *vectorStoreService) TestConnection(
	ctx context.Context,
	engineType types.RetrieverEngineType,
	config types.ConnectionConfig,
) (string, error) {
	switch engineType {
	case types.ElasticsearchRetrieverEngineType:
		return testElasticsearchConnection(ctx, config)
	case types.PostgresRetrieverEngineType:
		return testPostgresConnection(ctx, config)
	case types.QdrantRetrieverEngineType:
		return testQdrantConnection(ctx, config)
	case types.MilvusRetrieverEngineType:
		return testMilvusConnection(ctx, config)
	case types.TencentVectorDBRetrieverEngineType:
		return testTencentVectorDBConnection(ctx, config)
	case types.WeaviateRetrieverEngineType:
		return testWeaviateConnection(ctx, config)
	case types.DorisRetrieverEngineType:
		return testDorisConnection(ctx, config)
	case types.OpenSearchRetrieverEngineType:
		return testOpenSearchConnection(ctx, config)
	case types.SQLiteRetrieverEngineType:
		// SQLite is file-based, no remote connection to test
		return "", nil
	default:
		return "", errors.NewBadRequestError(
			fmt.Sprintf("connection test not supported for engine type: %s", engineType))
	}
}

func testElasticsearchConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	// Use plain HTTP GET to the root endpoint instead of the go-elasticsearch SDK.
	// The v8 SDK's TypedClient performs a product check that rejects ES7 servers,
	// so we use a raw HTTP request to support both v7 and v8.
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, config.Addr, nil)
	if err != nil {
		return "", errors.NewBadRequestError("failed to create elasticsearch request: invalid address")
	}
	if config.Username != "" {
		req.SetBasicAuth(config.Username, config.Password)
	}

	clientCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
	clientCfg.Timeout = connectionTestTimeout
	client := secutils.NewSSRFSafeHTTPClient(clientCfg)
	client.CheckRedirect = func(*http.Request, []*http.Request) error {
		// A health probe must not follow redirects: a malicious or compromised
		// endpoint could return a 302 to an internal address, defeating the
		// SSRF check that only validated the original (user-supplied) Addr.
		return http.ErrUseLastResponse
	}
	resp, err := client.Do(req)
	if err != nil {
		logger.Warnf(ctx, "Elasticsearch connection test failed: %v", err)
		return "", errors.NewBadRequestError("failed to connect to elasticsearch: connection refused or authentication failed")
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		logger.Warnf(ctx, "Elasticsearch connection test returned status %d", resp.StatusCode)
		return "", errors.NewBadRequestError("failed to connect to elasticsearch: authentication failed or server error")
	}

	// Parse version from response: {"version": {"number": "7.10.1"}, ...}
	body, err := io.ReadAll(io.LimitReader(resp.Body, 4096))
	if err != nil {
		return "", nil // connected but version unknown
	}

	var esInfo struct {
		Version struct {
			Number string `json:"number"`
		} `json:"version"`
	}
	if err := json.Unmarshal(body, &esInfo); err != nil {
		return "", nil // connected but version unparseable
	}

	return esInfo.Version.Number, nil
}

func testPostgresConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout)
	defer cancel()

	if config.UseDefaultConnection {
		// Using the default app DB connection — always reachable if the app is running.
		// Cannot query version without a DB handle; return empty.
		return "", nil
	}

	db, err := sql.Open("pgx", config.Addr)
	if err != nil {
		return "", errors.NewBadRequestError("failed to create postgres connection: invalid configuration")
	}
	defer db.Close()

	if err := db.PingContext(testCtx); err != nil {
		logger.Warnf(ctx, "Postgres connection test failed: %v", err)
		return "", errors.NewBadRequestError("failed to connect to postgres: connection refused or authentication failed")
	}

	// Detect version
	var version string
	if err := db.QueryRowContext(testCtx, "SHOW server_version").Scan(&version); err != nil {
		logger.Warnf(ctx, "Postgres version detection failed: %v", err)
		return "", nil // connected but version unknown
	}

	return version, nil
}

func testQdrantConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout)
	defer cancel()

	port := config.Port
	if port == 0 {
		port = 6334
	}

	client, err := qdrant.NewClient(&qdrant.Config{
		Host:   config.Host,
		Port:   port,
		APIKey: config.APIKey,
		UseTLS: config.UseTLS,
		GrpcOptions: []grpc.DialOption{
			grpc.WithContextDialer(secutils.SSRFSafeGRPCDialer),
		},
	})
	if err != nil {
		return "", errors.NewBadRequestError("failed to create qdrant client: invalid configuration")
	}
	defer client.Close()

	result, err := client.HealthCheck(testCtx)
	if err != nil {
		logger.Warnf(ctx, "Qdrant connection test failed: %v", err)
		return "", errors.NewBadRequestError("failed to connect to qdrant: connection refused or authentication failed")
	}

	return result.GetVersion(), nil
}

func testMilvusConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	// Use TCP dial instead of the Milvus SDK to avoid protobuf namespace conflict
	// between milvus-proto and qdrant-client (both register "common.proto").
	// A TCP dial is sufficient for connectivity verification; the Milvus SDK client
	// creation in container.go (PR 3) will validate full gRPC connectivity.
	// Version detection is not possible with TCP dial alone.
	testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout)
	defer cancel()

	addr := config.Addr
	if addr == "" {
		addr = "localhost:19530"
	}

	conn, err := secutils.SSRFSafeDialContext(testCtx, "tcp", addr)
	if err != nil {
		logger.Warnf(ctx, "Milvus connection test failed: %v", err)
		return "", errors.NewBadRequestError("failed to connect to milvus: connection refused or server unreachable")
	}
	defer conn.Close()

	return "", nil
}

func testTencentVectorDBConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout)
	defer cancel()

	client, err := tcvectordb.NewRpcClient(config.Addr, config.Username, config.APIKey, &tcvectordb.ClientOption{
		ReadConsistency: tcvectordb.EventualConsistency,
		Timeout:         connectionTestTimeout,
		Transport: &secutils.SSRFValidatingRoundTripper{
			Base: secutils.NewSSRFSafeTransport(secutils.DefaultSSRFSafeHTTPClientConfig()),
		},
	})
	if err != nil {
		logger.Warnf(ctx, "Tencent VectorDB connection test failed: %v", err)
		return "", errors.NewBadRequestError("failed to connect to tencent vectordb: connection refused or authentication failed")
	}
	defer client.Close()

	if _, err := client.ListDatabase(testCtx); err != nil {
		logger.Warnf(ctx, "Tencent VectorDB list database failed: %v", err)
		return "", errors.NewBadRequestError("failed to connect to tencent vectordb: authentication failed or server error")
	}
	return "", nil
}

func testWeaviateConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout)
	defer cancel()

	host := config.Host
	if host == "" {
		host = "weaviate:8080"
	}
	grpcAddress := config.GrpcAddress
	if grpcAddress == "" {
		grpcAddress = "weaviate:50051"
	}
	scheme := config.Scheme
	if scheme == "" {
		scheme = "http"
	}

	weaviateCfg := weaviate.Config{
		Host:             host,
		ConnectionClient: secutils.NewSSRFSafeHTTPClient(secutils.DefaultSSRFSafeHTTPClientConfig()),
		GrpcConfig: &wgrpc.Config{
			Host: grpcAddress,
		},
		Scheme: scheme,
	}
	if config.APIKey != "" {
		weaviateCfg.AuthConfig = auth.ApiKey{Value: config.APIKey}
	}

	// Weaviate Go client v5 does not expose a Close() method;
	// it uses HTTP + gRPC transports that are managed internally.
	client, err := weaviate.NewClient(weaviateCfg)
	if err != nil {
		logger.Warnf(ctx, "Weaviate connection test failed: %v", err)
		return "", errors.NewBadRequestError("failed to create weaviate client: invalid configuration")
	}

	isReady, err := client.Misc().ReadyChecker().Do(testCtx)
	if err != nil || !isReady {
		logger.Warnf(ctx, "Weaviate connection test failed: ready=%v, err=%v", isReady, err)
		return "", errors.NewBadRequestError("failed to connect to weaviate: server not ready or authentication failed")
	}

	// Detect version via /v1/meta
	meta, err := client.Misc().MetaGetter().Do(testCtx)
	if err != nil || meta == nil {
		return "", nil // connected but version unknown
	}

	return meta.Version, nil
}

// testDorisConnection 通过 MySQL 协议（database/sql + go-sql-driver）
// Ping Doris FE 并查询 @@version。
//
// Doris 的 @@version 形如 "5.7.99 Doris-4.1.0"——前半段是 MySQL 协议
// 兼容性表达式，"Doris-" 之后才是真实版本号。统一只返回 "4.1.0" 这类
// 裸版本号，与 Postgres/ES 路径的格式保持一致。
func testDorisConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout)
	defer cancel()

	if config.Addr == "" {
		return "", errors.NewBadRequestError("failed to create doris connection: addr is required")
	}

	// Database 不强制要求；Ping 时无明确库则用 information_schema（任何 MySQL 兼容服务都有）。
	database := config.Database
	if database == "" {
		database = "information_schema"
	}

	// 用 mysql.Config.FormatDSN() 构造 DSN，避免用户名/密码中 `@` `:` `/`
	// 等特殊字符破坏字面量拼接（fmt.Sprintf 会跑偏，参考 issue #1234 类问题）。
	cfg := mysql.NewConfig()
	cfg.User = config.Username
	cfg.Passwd = config.Password
	secutils.RegisterMySQLSSRFDialer()
	cfg.Net = secutils.MySQLSSRFNetwork
	cfg.Addr = config.Addr
	cfg.DBName = database
	cfg.Timeout = 5 * time.Second
	db, err := sql.Open("mysql", cfg.FormatDSN())
	if err != nil {
		return "", errors.NewBadRequestError("failed to create doris connection: invalid configuration")
	}
	defer db.Close()

	if err := db.PingContext(testCtx); err != nil {
		logger.Warnf(ctx, "Doris connection test failed: %v", err)
		return "", errors.NewBadRequestError("failed to connect to doris: connection refused or authentication failed")
	}

	var version string
	if err := db.QueryRowContext(testCtx, "SELECT @@version").Scan(&version); err != nil {
		logger.Warnf(ctx, "Doris version detection failed: %v", err)
		return "", nil
	}
	if i := strings.Index(version, "Doris-"); i >= 0 {
		return strings.TrimSpace(version[i+len("Doris-"):]), nil
	}
	return version, nil
}

// testOpenSearchConnection verifies the cluster is reachable, runs a
// supported OpenSearch version, and has the k-NN plugin installed. The driver
// owns the probe logic; a generic message is returned on failure so cluster
// internals are not surfaced to the API caller.
func testOpenSearchConnection(ctx context.Context, config types.ConnectionConfig) (string, error) {
	if config.Addr == "" {
		return "", errors.NewBadRequestError("failed to create opensearch connection: addr is required")
	}
	testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout)
	defer cancel()
	if err := openSearchRepo.TestConnection(testCtx, &config); err != nil {
		logger.Warnf(ctx, "OpenSearch connection test failed: %v", err)
		return "", errors.NewBadRequestError(
			"failed to connect to opensearch: check address, credentials, version (>= 2.4), and that the k-NN plugin is installed")
	}
	// Version is detected during the probe but not surfaced here; lazy index
	// creation re-validates on first use.
	return "", nil
}