From 34adc4b30c8540eb2f3aad4f5d511574fe526146 Mon Sep 17 00:00:00 2001
From: Isabella Basso do Amaral <idoamara@redhat.com>
Date: Mon, 15 Jul 2024 12:09:26 -0300
Subject: [PATCH] Use grpc.NewClient instead of DialContext

as it's been deprecated: https://pkg.go.dev/google.golang.org/grpc#DialContext

Signed-off-by: Isabella Basso do Amaral <idoamara@redhat.com>
---
 cmd/proxy.go                               | 26 +++++++---------------
 docs/mr_go_library.md                      |  7 ++----
 internal/testutils/test_container_utils.go |  5 +----
 3 files changed, 11 insertions(+), 27 deletions(-)

diff --git a/cmd/proxy.go b/cmd/proxy.go
index f940ce56d..4650bca07 100644
--- a/cmd/proxy.go
+++ b/cmd/proxy.go
@@ -1,10 +1,8 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 	"net/http"
-	"time"
 
 	"github.com/golang/glog"
 	"github.com/kubeflow/model-registry/internal/mlmdtypes"
@@ -15,32 +13,24 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 )
 
-var (
-	// proxyCmd represents the proxy command
-	proxyCmd = &cobra.Command{
-		Use:   "proxy",
-		Short: "Starts the ml-metadata go OpenAPI proxy",
-		Long: `This command launches the ml-metadata go OpenAPI proxy server.
+// proxyCmd represents the proxy command
+var proxyCmd = &cobra.Command{
+	Use:   "proxy",
+	Short: "Starts the ml-metadata go OpenAPI proxy",
+	Long: `This command launches the ml-metadata go OpenAPI proxy server.
 
 The server connects to a mlmd CPP server. It supports options to customize the
 hostname and port where it listens.'`,
-		RunE: runProxyServer,
-	}
-)
+	RunE: runProxyServer,
+}
 
 func runProxyServer(cmd *cobra.Command, args []string) error {
 	glog.Infof("proxy server started at %s:%v", cfg.Hostname, cfg.Port)
 
-	ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Second*30)
-	defer cancel()
-
 	mlmdAddr := fmt.Sprintf("%s:%d", proxyCfg.MLMDHostname, proxyCfg.MLMDPort)
 	glog.Infof("connecting to MLMD server %s..", mlmdAddr)
-	conn, err := grpc.DialContext(
-		ctxTimeout,
+	conn, err := grpc.NewClient(
 		mlmdAddr,
-		grpc.WithReturnConnectionError(),
-		grpc.WithBlock(),
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 	)
 	if err != nil {
diff --git a/docs/mr_go_library.md b/docs/mr_go_library.md
index ed8e83b43..6702e902a 100644
--- a/docs/mr_go_library.md
+++ b/docs/mr_go_library.md
@@ -30,11 +30,8 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 )
 
-conn, err := grpc.DialContext(
-  context.Background(),
+conn, err := grpc.NewClient(
   "localhost:9090",
-  grpc.WithReturnConnectionError(),
-  grpc.WithBlock(), // optional
   grpc.WithTransportCredentials(insecure.NewCredentials()),
 )
 if err != nil {
@@ -43,7 +40,7 @@ if err != nil {
 defer conn.Close()
 ```
 
-> NOTE: check [grpc doc](https://pkg.go.dev/google.golang.org/grpc#DialContext) for more details.
+> NOTE: check [grpc doc](https://pkg.go.dev/google.golang.org/grpc#NewClient) for more details.
 
 Once the gRPC connection is setup, let's create the `ModelRegistryService`:
 
diff --git a/internal/testutils/test_container_utils.go b/internal/testutils/test_container_utils.go
index 2795b2a7a..edd74054f 100644
--- a/internal/testutils/test_container_utils.go
+++ b/internal/testutils/test_container_utils.go
@@ -121,11 +121,8 @@ func SetupMLMetadataTestContainer(t *testing.T) (*grpc.ClientConn, proto.Metadat
 	t.Log("MLMD test container running at: ", mlmdAddr)
 
 	// setup grpc connection
-	conn, err := grpc.DialContext(
-		context.Background(),
+	conn, err := grpc.NewClient(
 		mlmdAddr,
-		grpc.WithReturnConnectionError(),
-		grpc.WithBlock(),
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 	)
 	if err != nil {