Adding grpc health service in the clientset (#228)

* Adding grpc health service in the clientset Signed-off-by: Prafulla Mahindrakar <[email protected]> * Added unit tests Signed-off-by: Prafulla Mahindrakar <[email protected]> Signed-off-by: Eduardo Apolinario <[email protected]>
flyteorg · Sep 13, 2023 · b657c46 · b657c46
1 parent 6aab097
commit b657c46
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/flyteidl/clients/go/admin/client.go b/flyteidl/clients/go/admin/client.go
@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
+
 	"sync"
 
 	"github.com/flyteorg/flyteidl/clients/go/admin/mocks"
@@ -17,6 +18,7 @@ import (
 	grpcPrometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/health/grpc_health_v1"
 )
 
 // IDE "Go Generate File". This will create a mocks/AdminServiceClient.go file
@@ -36,6 +38,7 @@ var (
 type Clientset struct {
 	adminServiceClient        service.AdminServiceClient
 	authMetadataServiceClient service.AuthMetadataServiceClient
+	healthServiceClient       grpc_health_v1.HealthClient
 	identityServiceClient     service.IdentityServiceClient
 }
 
@@ -49,6 +52,11 @@ func (c Clientset) AuthMetadataClient() service.AuthMetadataServiceClient {
 	return c.authMetadataServiceClient
 }
 
+// HealthServiceClient retrieves the grpc_health_v1.HealthClient
+func (c Clientset) HealthServiceClient() grpc_health_v1.HealthClient {
+	return c.healthServiceClient
+}
+
 func (c Clientset) IdentityClient() service.IdentityServiceClient {
 	return c.identityServiceClient
 }
@@ -192,6 +200,7 @@ func initializeClients(ctx context.Context, cfg *Config, tokenCache pkce.TokenCa
 	cs.adminServiceClient = NewAdminClient(ctx, adminConnection)
 	cs.authMetadataServiceClient = service.NewAuthMetadataServiceClient(adminConnection)
 	cs.identityServiceClient = service.NewIdentityServiceClient(adminConnection)
+	cs.healthServiceClient = grpc_health_v1.NewHealthClient(adminConnection)
 	return &cs, nil
 }
 

diff --git a/flyteidl/clients/go/admin/client_test.go b/flyteidl/clients/go/admin/client_test.go
@@ -72,6 +72,7 @@ func TestGetAdditionalAdminClientConfigOptions(t *testing.T) {
 		assert.NotNil(t, clientSet.AdminClient())
 		assert.NotNil(t, clientSet.AuthMetadataClient())
 		assert.NotNil(t, clientSet.IdentityClient())
+		assert.NotNil(t, clientSet.HealthServiceClient())
 	})
 
 	t.Run("legal-from-config", func(t *testing.T) {
@@ -80,6 +81,7 @@ func TestGetAdditionalAdminClientConfigOptions(t *testing.T) {
 		assert.NotNil(t, clientSet)
 		assert.NotNil(t, clientSet.AuthMetadataClient())
 		assert.NotNil(t, clientSet.AdminClient())
+		assert.NotNil(t, clientSet.HealthServiceClient())
 	})
 }