Add job artifacts as a general available env var for JM and TM (#118)

spotify · Mar 21, 2022 · e3b082d · e3b082d
1 parent f743b25
commit e3b082d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 8 deletions.
diff --git a/pkg/flink/resources.go b/pkg/flink/resources.go
@@ -16,13 +16,14 @@ package flink
 
 import (
 	"bytes"
-	"k8s.io/apimachinery/pkg/types"
-	"k8s.io/apimachinery/pkg/util/json"
 	"net/url"
-	"sigs.k8s.io/controller-runtime/pkg/client"
 	"strings"
 	"text/template"
 
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/json"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
 	corev1 "k8s.io/api/core/v1"
 
 	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/utils"
@@ -34,8 +35,9 @@ import (
 )
 
 var (
-	containerTmpl       = template.New("container-template").Funcs(template.FuncMap{"join": strings.Join})
-	flinkPropertiesTmpl = template.New("flink-properties-template").Funcs(template.FuncMap{"join": strings.Join})
+	containerTmpl        = template.New("container-template").Funcs(template.FuncMap{"join": strings.Join})
+	flinkPropertiesTmpl  = template.New("flink-properties-template").Funcs(template.FuncMap{"join": strings.Join})
+	stagedJarsEnvVarName = "STAGED_JARS"
 )
 
 type ContainerTemplateData struct {
@@ -298,6 +300,10 @@ func NewFlinkCluster(config *Config, taskCtx FlinkTaskContext) (*flinkOp.FlinkCl
 		Kind:       KindFlinkCluster,
 		APIVersion: flinkOp.GroupVersion.String(),
 	}
+	cluster.Spec.EnvVars = []corev1.EnvVar{{
+		Name:  stagedJarsEnvVarName,
+		Value: strings.Join(taskCtx.Job.JarFiles, " "),
+	}}
 
 	cluster.updateFlinkProperties(config, taskCtx)
 

diff --git a/pkg/flink/resources_test.go b/pkg/flink/resources_test.go
@@ -15,11 +15,13 @@
 package flink
 
 import (
-	"k8s.io/apimachinery/pkg/types"
-	"k8s.io/apimachinery/pkg/util/json"
 	"reflect"
+	"strings"
 	"testing"
 
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/json"
+
 	flinkOp "github.com/spotify/flink-on-k8s-operator/apis/flinkcluster/v1beta1"
 	flinkIdl "github.com/spotify/flyte-flink-plugin/gen/pb-go/flyteidl-flink"
 	"gotest.tools/assert"
@@ -79,6 +81,13 @@ func TestBuildFlinkClusterSpecValid(t *testing.T) {
 
 	assert.Equal(t, string(job.CleanupPolicy.AfterJobSucceeds), flinkOp.CleanupActionDeleteCluster)
 	assert.Equal(t, cluster.Spec.FlinkProperties["metrics.reporter.promgateway.groupingKey"], "namespace=test-namespace;cluster=generated-name;execution_id=1")
+
+	assert.Assert(t, len(cluster.Spec.EnvVars) == 1, "EnvVars should contain only 1 env")
+	expectedEnvVars := []corev1.EnvVar{{
+		Name:  stagedJarsEnvVarName,
+		Value: strings.Join(jobIdl.JarFiles, " "),
+	}}
+	assert.Equal(t, cluster.Spec.EnvVars[0], expectedEnvVars[0])
 }
 
 func TestWithPersistentVolume(t *testing.T) {
@@ -294,7 +303,7 @@ func TestBuildAnnotationPatch(t *testing.T) {
 	assert.DeepEqual(
 		t,
 		jsonData["metadata"].(map[string]interface{})["annotations"].(map[string]interface{}),
-		map[string]interface{} {"testKey": "testValue"},
+		map[string]interface{}{"testKey": "testValue"},
 	)
 
 	assert.NilError(t, err)