Add support for connecting standalone PEM in pxapi. (#1762)

Summary: Add support for connecting standalone PEM in pxapi.

Type of change: /kind feature

Test Plan: Added an example program to connect standalone PEM.

---------

Signed-off-by: RagalahariP <[email protected]>
Signed-off-by: Pete Stevenson <[email protected]>
Co-authored-by: RagalahariP <[email protected]>
GitOrigin-RevId: ff140e45d3102331300a1456e4272391a9c31503

BUILD.bazel

@@ -38,6 +38,7 @@ go_library(
         "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//credentials",
+        "@org_golang_google_grpc//credentials/insecure",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
     ],

client.go

@@ -25,6 +25,7 @@ import (
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/metadata"
 
 	"px.dev/pxapi/types"
@@ -53,15 +54,15 @@ type TableMuxer interface {
 	AcceptTable(ctx context.Context, metadata types.TableMetadata) (TableRecordHandler, error)
 }
 
-// Client is the base client to use pixie cloud + vizier.
+// Client is the base client to use either pixie cloud + vizier or standalone pem + vizier.
 type Client struct {
 	apiKey     string
 	bearerAuth string
-
-	cloudAddr string
+	vzAddr     string
 
 	useEncryption          bool
 	disableTLSVerification bool
+	insecureDirect         bool
 
 	grpcConn *grpc.ClientConn
 	cmClient cloudpb.VizierClusterInfoClient
@@ -71,8 +72,10 @@ type Client struct {
 // NewClient creates a new Pixie API Client.
 func NewClient(ctx context.Context, opts ...ClientOption) (*Client, error) {
 	c := &Client{
-		cloudAddr:     defaultCloudAddr,
-		useEncryption: true,
+		vzAddr:                 defaultCloudAddr,
+		useEncryption:          true,
+		insecureDirect:         false,
+		disableTLSVerification: false,
 	}
 
 	for _, opt := range opts {
@@ -89,7 +92,11 @@ func (c *Client) init(ctx context.Context) error {
 	tlsConfig := &tls.Config{InsecureSkipVerify: c.disableTLSVerification}
 	creds := credentials.NewTLS(tlsConfig)
 
-	conn, err := grpc.Dial(c.cloudAddr, grpc.WithTransportCredentials(creds))
+	if c.insecureDirect {
+		creds = insecure.NewCredentials()
+	}
+
+	conn, err := grpc.Dial(c.vzAddr, grpc.WithTransportCredentials(creds))
 	if err != nil {
 		return err
 	}

examples/standalone_pem_example/BUILD.bazel

@@ -0,0 +1,36 @@
+# Copyright 2018- The Pixie Authors.
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("//bazel:pl_build_system.bzl", "pl_go_binary")
+
+go_library(
+    name = "standalone_pem_example_lib",
+    srcs = ["example.go"],
+    importpath = "px.dev/pxapi/examples/standalone_pem_example",
+    visibility = ["//visibility:private"],
+    deps = [
+        "//src/api/go/pxapi",
+        "//src/api/go/pxapi/errdefs",
+        "//src/api/go/pxapi/types",
+    ],
+)
+
+pl_go_binary(
+    name = "standalone_pem_example",
+    embed = [":standalone_pem_example_lib"],
+    visibility = ["//src:__subpackages__"],
+)

examples/standalone_pem_example/example.go

@@ -0,0 +1,111 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	"px.dev/pxapi"
+	"px.dev/pxapi/errdefs"
+	"px.dev/pxapi/types"
+)
+
+// Define PxL script with one table output.
+var (
+	pxl = `
+import px
+
+# Look at the http_events.
+df = px.DataFrame(table='http_events')
+
+# Grab the command line from the metadata.
+df.cmdline = px.upid_to_cmdline(df.upid)
+
+# Limit to the first 10.
+df = df.head(10)
+
+px.display(df)`
+)
+
+func main() {
+	// Create a Pixie client with local standalonePEM listening address
+	ctx := context.Background()
+	client, err := pxapi.NewClient(
+		ctx,
+		pxapi.WithDirectAddr("127.0.0.1:12345"),
+		pxapi.WithDirectCredsInsecure(),
+	)
+	if err != nil {
+		panic(err)
+	}
+	// Create a connection to the host.
+	hostID := "localhost"
+	vz, err := client.NewVizierClient(ctx, hostID)
+	if err != nil {
+		panic(err)
+	}
+	// Create TableMuxer to accept results table.
+	tm := &tableMux{}
+	// Execute the PxL script.
+	resultSet, err := vz.ExecuteScript(ctx, pxl, tm)
+	if err != nil && err != io.EOF {
+		panic(err)
+	}
+	// Receive the PxL script results.
+	defer resultSet.Close()
+	if err := resultSet.Stream(); err != nil {
+		if errdefs.IsCompilationError(err) {
+			fmt.Printf("Got compiler error: \n %s\n", err.Error())
+		} else {
+			fmt.Printf("Got error : %+v, while streaming\n", err)
+		}
+	}
+	// Get the execution stats for the script execution.
+	stats := resultSet.Stats()
+	fmt.Printf("Execution Time: %v\n", stats.ExecutionTime)
+	fmt.Printf("Bytes received: %v\n", stats.TotalBytes)
+}
+
+// Satisfies the TableRecordHandler interface.
+type tablePrinter struct{}
+
+func (t *tablePrinter) HandleInit(ctx context.Context, metadata types.TableMetadata) error {
+	return nil
+}
+func (t *tablePrinter) HandleRecord(ctx context.Context, r *types.Record) error {
+	for _, d := range r.Data {
+		fmt.Printf("%s ", d.String())
+	}
+	fmt.Printf("\n")
+	return nil
+}
+
+func (t *tablePrinter) HandleDone(ctx context.Context) error {
+	return nil
+}
+
+// Satisfies the TableMuxer interface.
+type tableMux struct {
+}
+
+func (s *tableMux) AcceptTable(ctx context.Context, metadata types.TableMetadata) (pxapi.TableRecordHandler, error) {
+	return &tablePrinter{}, nil
+}

opts.go

@@ -30,7 +30,7 @@ type ClientOption func(client *Client)
 // WithCloudAddr is the option to specify cloud address to use.
 func WithCloudAddr(cloudAddr string) ClientOption {
 	return func(c *Client) {
-		c.cloudAddr = cloudAddr
+		c.vzAddr = cloudAddr
 	}
 }
 
@@ -43,9 +43,15 @@ func WithDisableTLSVerification(cloudAddr string) ClientOption {
 
 		if !tlsDisabled && isInternal {
 			log.Fatalf("The `PX_DISABLE_TLS` environment variable must be set to \"1\" when making cloud connections that do not support TLS.\n")
-		} else {
-			c.disableTLSVerification = insecureSkipVerify
 		}
+		c.disableTLSVerification = insecureSkipVerify
+	}
+}
+
+// WithDirectAddr is the option to specify direct address to use for data from standalone pem.
+func WithDirectAddr(directAddr string) ClientOption {
+	return func(c *Client) {
+		c.vzAddr = directAddr
 	}
 }
 
@@ -69,3 +75,10 @@ func WithE2EEncryption(enabled bool) ClientOption {
 		c.useEncryption = enabled
 	}
 }
+
+// WithDirectCredsInsecure is the option to setup insecure credentials for direct connections.
+func WithDirectCredsInsecure() ClientOption {
+	return func(c *Client) {
+		c.insecureDirect = true
+	}
+}