Load cron scripts from multiple sources (#1326)

Summary: For our usage of Pixie, it's convenient to exclude Pixie's
dependency on the control plane. This commit allows the query broker:
1. to load cron scripts from config maps in the same namespace
2. to turn on and off loading cron scripts from the cloud or config maps

Type of change: /kind feature

Changes: This commit implements the functionality described above.

* added a `scriptrunner.Source` type that encapsulates how cron scripts
are initialized/updated.
* moved existing logic for fetching cron scripts to
`scriptrunner.CloudSource`
* added `scriptrunner.ConfigMapSource`
* Modified `scriptrunner.ScriptRunner` to use `scriptrunner.Source`s.
* Added a configuration flag (`--cron_script_sources` or or
`PL_CRON_SCRIPT_SOURCES`) that allows a customer to choose which sources
cron scripts can be loaded from. By default, it only loads cron scripts
from the control plane. If you wanted to use both, for example, you
could set `PL_CRON_SCRIPT_SOURCES="cloud configmaps"`.
* Moved many of the fakes/mocks/test helpers to `helper_test.go`

Test Plan: All of these changes have been tested at the unit level in
addition to manual testing in all relevant configurations. We did our
manual testing on GKE, EKS, and Kind with Kubernetes version 1.25. We
don't anticipate other Kubernetes platforms or versions to be a problem
since the config map API is stable and consistent across those
dimensions. Moving the logic to load scripts from the control plane was
probably the riskiest change and might warrant some extra attention.

---------

Signed-off-by: Matthew Conger-Eldeen <matthewco@vmware.com>
Co-authored-by: Devon Warshaw <warshawd@vmware.com>
Co-authored-by: Joe Gee <jgee@vmware.com>
Co-authored-by: Glenn Oppegard <oppegard@vmware.com>

Author: 4 people

k8s/vizier/base/default_role.yaml

@@ -0,0 +1,12 @@
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: RoleBinding
+metadata:
+  name: pl-vizier-crd-binding
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: Role
+  name: pl-vizier-crd-role
+subjects:
+- kind: ServiceAccount
+  name: default

k8s/vizier/base/kustomization.yaml

@@ -21,6 +21,7 @@ patches:
     kind: StatefulSet
 resources:
 - ../bootstrap
+- default_role.yaml
 - kelvin_deployment.yaml
 - kelvin_service.yaml
 - metadata_role.yaml

k8s/vizier/base/query_broker_deployment.yaml

@@ -105,6 +105,7 @@ spec:
             - ALL
           seccompProfile:
             type: RuntimeDefault
+      serviceAccountName: query-broker-service-account
       securityContext:
         runAsUser: 10100
         runAsGroup: 10100

k8s/vizier/base/query_broker_role.yaml

@@ -1,13 +1,44 @@
 ---
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: query-broker-service-account
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+  creationTimestamp: null
+  name: pl-vizier-query-broker-role
+rules:
+- apiGroups:
+  - ""
+  resources:
+  - configmaps
+  verbs:
+  - get
+  - list
+  - watch
+---
 apiVersion: rbac.authorization.k8s.io/v1
 kind: RoleBinding
 metadata:
-  name: pl-vizier-crd-binding
+  name: pl-vizier-query-broker-crd-binding
 roleRef:
   apiGroup: rbac.authorization.k8s.io
   kind: Role
   name: pl-vizier-crd-role
 subjects:
 - kind: ServiceAccount
-  name: default
-  namespace: pl
+  name: query-broker-service-account
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: RoleBinding
+metadata:
+  name: pl-vizier-query-broker-binding
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: Role
+  name: pl-vizier-query-broker-role
+subjects:
+- kind: ServiceAccount
+  name: query-broker-service-account

src/utils/template_generator/vizier_yamls/vizier_yamls.go

@@ -418,6 +418,18 @@ func generateVzYAMLs(yamlMap map[string]string) ([]*yamls.YAMLFile, error) {
 			Placeholder:     "__PX_SUBJECT_NAMESPACE__",
 			TemplateValue:   nsTmpl,
 		},
+		{
+			TemplateMatcher: yamls.GenerateResourceNameMatcherFn("pl-vizier-query-broker-crd-binding"),
+			Patch:           `{ "subjects": [{ "name": "query-broker-service-account", "namespace": "__PX_SUBJECT_NAMESPACE__", "kind": "ServiceAccount" }] }`,
+			Placeholder:     "__PX_SUBJECT_NAMESPACE__",
+			TemplateValue:   nsTmpl,
+		},
+		{
+			TemplateMatcher: yamls.GenerateResourceNameMatcherFn("pl-vizier-query-broker-binding"),
+			Patch:           `{ "subjects": [{ "name": "query-broker-service-account", "namespace": "__PX_SUBJECT_NAMESPACE__", "kind": "ServiceAccount" }] }`,
+			Placeholder:     "__PX_SUBJECT_NAMESPACE__",
+			TemplateValue:   nsTmpl,
+		},
 		{
 			TemplateMatcher: yamls.GenerateResourceNameMatcherFn("pl-vizier-metadata-node-view-cluster-binding"),
 			Patch:           `{ "subjects": [{ "name": "metadata-service-account", "namespace": "__PX_SUBJECT_NAMESPACE__", "kind": "ServiceAccount" }] }`,

src/vizier/services/query_broker/controllers/query_executor.go

@@ -508,7 +508,7 @@ func (q *QueryExecutorImpl) prepareScript(ctx context.Context, resultCh chan<- *
 func (q *QueryExecutorImpl) runScript(ctx context.Context, resultCh chan<- *vizierpb.ExecuteScriptResponse, req *vizierpb.ExecuteScriptRequest) error {
 	defer close(resultCh)
 	q.startTime = time.Now()
-	log.WithField("query_id", q.queryID).Infof("Running script")
+	log.WithField("query_id", q.queryID).WithField("query_name", q.queryName).Infof("Running script")
 
 	if req.QueryID == "" {
 		if err := q.prepareScript(ctx, resultCh, req); err != nil {

src/vizier/services/query_broker/query_broker_server.go

@@ -59,6 +59,7 @@ func init() {
 	pflag.String("mds_service", "vizier-metadata-svc", "The metadata service name")
 	pflag.String("mds_port", "50400", "The querybroker service port")
 	pflag.String("pod_namespace", "pl", "The namespace this pod runs in.")
+	pflag.StringArray("cron_script_sources", scriptrunner.DefaultSources, "Where to find cron scripts (cloud, configmaps)")
 }
 
 // NewVizierServiceClient creates a new vz RPC client stub.
@@ -206,10 +207,14 @@ func main() {
 	defer ptProxy.Close()
 
 	// Start cron script runner.
-	sr, err := scriptrunner.New(natsConn, csClient, vzServiceClient, viper.GetString("jwt_signing_key"))
-	if err != nil {
-		log.WithError(err).Fatal("Failed to start script runner")
-	}
+	sources := scriptrunner.Sources(
+		natsConn,
+		csClient,
+		viper.GetString("jwt_signing_key"),
+		viper.GetString("pod_namespace"),
+		viper.GetStringSlice("cron_script_sources"),
+	)
+	sr := scriptrunner.New(csClient, vzServiceClient, viper.GetString("jwt_signing_key"), sources...)
 
 	// Load the scripts and start the background sync.
 	go func() {
@@ -218,6 +223,7 @@ func main() {
 			log.WithError(err).Error("Failed to sync cron scripts")
 		}
 	}()
+
 	s.Start()
 	s.StopOnInterrupt()
 }

src/vizier/services/query_broker/script_runner/BUILD.bazel

@@ -19,7 +19,13 @@ load("//bazel:pl_build_system.bzl", "pl_go_test")
 
 go_library(
     name = "script_runner",
-    srcs = ["script_runner.go"],
+    srcs = [
+        "cloud_source.go",
+        "config_map_source.go",
+        "script_runner.go",
+        "source.go",
+        "sources.go",
+    ],
     importpath = "px.dev/pixie/src/vizier/services/query_broker/script_runner",
     visibility = ["//visibility:public"],
     deps = [
@@ -31,6 +37,7 @@ go_library(
         "//src/shared/scripts",
         "//src/shared/services/utils",
         "//src/utils",
+        "//src/utils/shared/k8s",
         "//src/vizier/services/metadata/metadatapb:service_pl_go_proto",
         "//src/vizier/utils/messagebus",
         "@com_github_gofrs_uuid//:uuid",
@@ -39,6 +46,15 @@ go_library(
         "@com_github_nats_io_nats_go//:nats_go",
         "@com_github_sirupsen_logrus//:logrus",
         "@in_gopkg_yaml_v2//:yaml_v2",
+        "@io_k8s_api//core/v1:core",
+        "@io_k8s_apimachinery//pkg/apis/meta/v1:meta",
+        "@io_k8s_apimachinery//pkg/labels",
+        "@io_k8s_client_go//informers",
+        "@io_k8s_client_go//informers/core/v1:core",
+        "@io_k8s_client_go//kubernetes",
+        "@io_k8s_client_go//listers/core/v1:core",
+        "@io_k8s_client_go//rest",
+        "@io_k8s_client_go//tools/cache",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
@@ -47,10 +63,16 @@ go_library(
 
 pl_go_test(
     name = "script_runner_test",
-    srcs = ["script_runner_test.go"],
+    srcs = [
+        "cloud_source_test.go",
+        "config_map_source_test.go",
+        "helper_test.go",
+        "script_runner_test.go",
+    ],
     embed = [":script_runner"],
     deps = [
         "//src/api/proto/vizierpb:vizier_pl_go_proto",
+        "//src/api/proto/vizierpb/mock",
         "//src/carnot/planner/compilerpb:compiler_status_pl_go_proto",
         "//src/common/base/statuspb:status_pl_go_proto",
         "//src/shared/cvmsgspb:cvmsgs_pl_go_proto",
@@ -61,9 +83,12 @@ pl_go_test(
         "@com_github_gofrs_uuid//:uuid",
         "@com_github_gogo_protobuf//proto",
         "@com_github_gogo_protobuf//types",
+        "@com_github_golang_mock//gomock",
         "@com_github_nats_io_nats_go//:nats_go",
-        "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
+        "@io_k8s_api//core/v1:core",
+        "@io_k8s_apimachinery//pkg/apis/meta/v1:meta",
+        "@io_k8s_client_go//kubernetes/fake",
         "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",