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 }