From b159812bb2ed56f6fee010f3d80cefdeb8f57956 Mon Sep 17 00:00:00 2001
From: Vihang Mehta <vihang@pixielabs.ai>
Date: Tue, 5 Oct 2021 10:52:17 -0700
Subject: [PATCH] Create RSA keys on VizierClient creation instead of every
 ExecuteScript

Summary: TSIA, creating RSA keys is expensive, let's cache them.

Test Plan: Run the go api examples

Reviewers: zasgar, philkuz

Reviewed By: zasgar

Signed-off-by: Vihang Mehta <vihang@pixielabs.ai>

Differential Revision: https://phab.corp.pixielabs.ai/D9908

GitOrigin-RevId: 5366ed29ad0fab18468cd3217e3abfdf6e55797d
---
 client.go | 18 ++++++++++++++----
 vizier.go | 22 +++++++---------------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/client.go b/client.go
index eaa4b9a..a69a633 100644
--- a/client.go
+++ b/client.go
@@ -30,6 +30,7 @@ import (
 	"google.golang.org/grpc/metadata"
 
 	"px.dev/pxapi/types"
+	"px.dev/pxapi/utils"
 	"px.dev/pxapi/proto/cloudpb"
 	"px.dev/pxapi/proto/vizierpb"
 )
@@ -148,12 +149,21 @@ func (c *Client) NewVizierClient(ctx context.Context, vizierID string) (*VizierC
 		vzConn = conn
 	}
 
+	var encOpts, decOpts *vizierpb.ExecuteScriptRequest_EncryptionOptions
+	if c.useEncryption {
+		encOpts, decOpts, err = utils.CreateEncryptionOptions()
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	// Now create the actual client.
 	vzClient := &VizierClient{
-		cloud:         c,
-		useEncryption: c.useEncryption,
-		vizierID:      vizierID,
-		vzClient:      vizierpb.NewVizierServiceClient(vzConn),
+		cloud:    c,
+		encOpts:  encOpts,
+		decOpts:  decOpts,
+		vizierID: vizierID,
+		vzClient: vizierpb.NewVizierServiceClient(vzConn),
 	}
 
 	return vzClient, nil
diff --git a/vizier.go b/vizier.go
index 191d985..a20e214 100644
--- a/vizier.go
+++ b/vizier.go
@@ -21,32 +21,24 @@ package pxapi
 import (
 	"context"
 
-	"px.dev/pxapi/utils"
 	"px.dev/pxapi/proto/vizierpb"
 )
 
 // VizierClient is the client for a single vizier.
 type VizierClient struct {
-	cloud         *Client
-	useEncryption bool
-	vizierID      string
-	vzClient      vizierpb.VizierServiceClient
+	cloud    *Client
+	vizierID string
+	vzClient vizierpb.VizierServiceClient
+	encOpts  *vizierpb.ExecuteScriptRequest_EncryptionOptions
+	decOpts  *vizierpb.ExecuteScriptRequest_EncryptionOptions
 }
 
 // ExecuteScript runs the script on vizier.
 func (v *VizierClient) ExecuteScript(ctx context.Context, pxl string, mux TableMuxer) (*ScriptResults, error) {
-	var encOpts, decOpts *vizierpb.ExecuteScriptRequest_EncryptionOptions
-	var err error
-	if v.useEncryption {
-		encOpts, decOpts, err = utils.CreateEncryptionOptions()
-		if err != nil {
-			return nil, err
-		}
-	}
 	req := &vizierpb.ExecuteScriptRequest{
 		ClusterID:         v.vizierID,
 		QueryStr:          pxl,
-		EncryptionOptions: encOpts,
+		EncryptionOptions: v.encOpts,
 	}
 	ctx, cancel := context.WithCancel(ctx)
 	res, err := v.vzClient.ExecuteScript(v.cloud.cloudCtxWithMD(ctx), req)
@@ -59,7 +51,7 @@ func (v *VizierClient) ExecuteScript(ctx context.Context, pxl string, mux TableM
 	sr.c = res
 	sr.cancel = cancel
 	sr.tm = mux
-	sr.decOpts = decOpts
+	sr.decOpts = v.decOpts
 
 	return sr, nil
 }