load cron scripts from multiple sources

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.

Co-authored-by: Devon Warshaw <warshawd@vmware.com>
Signed-off-by: Matthew Conger-Eldeen <matthewco@vmware.com>

Author: m25n

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

@@ -412,6 +412,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,12 @@ 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",
+        "script_source.go",
+    ],
     importpath = "px.dev/pixie/src/vizier/services/query_broker/script_runner",
     visibility = ["//visibility:public"],
     deps = [
@@ -31,6 +36,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 +45,11 @@ 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/watch",
+        "@io_k8s_client_go//kubernetes/typed/core/v1:core",
+        "@io_k8s_client_go//rest",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
@@ -47,10 +58,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 +78,16 @@ 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_apimachinery//pkg/labels",
+        "@io_k8s_apimachinery//pkg/runtime",
+        "@io_k8s_apimachinery//pkg/watch",
+        "@io_k8s_client_go//kubernetes/fake",
+        "@io_k8s_client_go//testing",
         "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",