From 65917321ae6668f75c564bcb7e53ad625a36f599 Mon Sep 17 00:00:00 2001
From: Edbert Linardi <edbert.linardi@snowflake.com>
Date: Tue, 3 Oct 2023 12:34:07 -0700
Subject: [PATCH] Add healthcheck utils and fix exec utils (#333)

---
 services/exec/client/utils.go        | 21 +++++----------
 services/healthcheck/client/utils.go | 40 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 14 deletions(-)
 create mode 100644 services/healthcheck/client/utils.go

diff --git a/services/exec/client/utils.go b/services/exec/client/utils.go
index cfeb66e7..1f494547 100644
--- a/services/exec/client/utils.go
+++ b/services/exec/client/utils.go
@@ -39,19 +39,16 @@ func ExecRemoteCommand(ctx context.Context, conn *proxy.Conn, binary string, arg
 	if len(result) < 1 {
 		return nil, fmt.Errorf("ExecRemoteCommand error: received empty response")
 	}
-	return result[0].ExecResponse, nil
-}
-
-type ExecResponse struct {
-	*pb.ExecResponse
-	Index  int
-	Target string
+	if result[0].Error != nil {
+		return nil, fmt.Errorf("ExecRemoteCommand error: %v", result[0].Error)
+	}
+	return result[0].Resp, nil
 }
 
 // ExecRemoteCommand is a helper function for execing a command on one or remote hosts
 // using a proxy.Conn.
 // `binary` refers to the absolute path of the binary file on the remote host.
-func ExecRemoteCommandMany(ctx context.Context, conn *proxy.Conn, binary string, args ...string) ([]ExecResponse, error) {
+func ExecRemoteCommandMany(ctx context.Context, conn *proxy.Conn, binary string, args ...string) ([]*pb.RunManyResponse, error) {
 	c := pb.NewExecClientProxy(conn)
 	req := &pb.ExecRequest{
 		Command: binary,
@@ -61,13 +58,9 @@ func ExecRemoteCommandMany(ctx context.Context, conn *proxy.Conn, binary string,
 	if err != nil {
 		return nil, err
 	}
-	result := make([]ExecResponse, len(conn.Targets))
+	result := make([]*pb.RunManyResponse, len(conn.Targets))
 	for r := range respChan {
-		result[r.Index] = ExecResponse{
-			r.Resp,
-			r.Index,
-			r.Target,
-		}
+		result[r.Index] = r
 	}
 
 	return result, nil
diff --git a/services/healthcheck/client/utils.go b/services/healthcheck/client/utils.go
new file mode 100644
index 00000000..5ee04089
--- /dev/null
+++ b/services/healthcheck/client/utils.go
@@ -0,0 +1,40 @@
+/*
+Copyright (c) 2023 Snowflake Inc. All rights reserved.
+
+	Licensed under the Apache License, Version 2.0 (the
+	"License"); you may not use this file except in compliance
+	with the License.  You may obtain a copy of the License at
+
+	  http://www.apache.org/licenses/LICENSE-2.0
+
+	Unless required by applicable law or agreed to in writing,
+	software distributed under the License is distributed on an
+	"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+	KIND, either express or implied.  See the License for the
+	specific language governing permissions and limitations
+	under the License.
+*/
+package client
+
+import (
+	"context"
+
+	"github.com/Snowflake-Labs/sansshell/proxy/proxy"
+	pb "github.com/Snowflake-Labs/sansshell/services/healthcheck"
+	"google.golang.org/protobuf/types/known/emptypb"
+)
+
+func HealthcheckValidateMany(ctx context.Context, conn *proxy.Conn) ([]*pb.OkManyResponse, error) {
+	c := pb.NewHealthCheckClientProxy(conn)
+
+	respChan, err := c.OkOneMany(ctx, &emptypb.Empty{})
+	if err != nil {
+		return nil, err
+	}
+	results := make([]*pb.OkManyResponse, len(conn.Targets))
+	for r := range respChan {
+		results[r.Index] = r
+	}
+
+	return results, nil
+}