Skip to content

Commit

Permalink
Add healthcheck utils and fix exec utils (Snowflake-Labs#333)
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-elinardi authored and stvnrhodes committed Oct 6, 2023
1 parent 4d04941 commit 6591732
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 14 deletions.
21 changes: 7 additions & 14 deletions services/exec/client/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down
40 changes: 40 additions & 0 deletions services/healthcheck/client/utils.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 6591732

Please sign in to comment.