From 3f8036f2dfbf2624be3c20edc43393487579c8a2 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Tue, 21 May 2024 12:00:04 +0200
Subject: [PATCH] Fixed seg fault when specifying environment key for tasks
 (#1443)

## Changes
Fixed seg fault when specifying environment key for tasks
---
 bundle/artifacts/artifacts.go                    |  4 ++++
 bundle/libraries/libraries.go                    |  4 ++++
 bundle/libraries/match.go                        |  4 ++++
 bundle/tests/enviroment_key_test.go              | 11 +++++++++++
 bundle/tests/environment_key_only/databricks.yml | 16 ++++++++++++++++
 5 files changed, 39 insertions(+)
 create mode 100644 bundle/tests/environment_key_only/databricks.yml

diff --git a/bundle/artifacts/artifacts.go b/bundle/artifacts/artifacts.go
index 101b598dd5..470c329a19 100644
--- a/bundle/artifacts/artifacts.go
+++ b/bundle/artifacts/artifacts.go
@@ -150,6 +150,10 @@ func uploadArtifact(ctx context.Context, b *bundle.Bundle, a *config.Artifact, u
 
 			for i := range job.Environments {
 				env := &job.Environments[i]
+				if env.Spec == nil {
+					continue
+				}
+
 				for j := range env.Spec.Dependencies {
 					lib := env.Spec.Dependencies[j]
 					if isArtifactMatchLibrary(f, lib, b) {
diff --git a/bundle/libraries/libraries.go b/bundle/libraries/libraries.go
index a79adedbf2..84ead052b3 100644
--- a/bundle/libraries/libraries.go
+++ b/bundle/libraries/libraries.go
@@ -30,6 +30,10 @@ func FindAllEnvironments(b *bundle.Bundle) map[string]([]jobs.JobEnvironment) {
 
 func isEnvsWithLocalLibraries(envs []jobs.JobEnvironment) bool {
 	for _, e := range envs {
+		if e.Spec == nil {
+			continue
+		}
+
 		for _, l := range e.Spec.Dependencies {
 			if IsEnvironmentDependencyLocal(l) {
 				return true
diff --git a/bundle/libraries/match.go b/bundle/libraries/match.go
index 096cdf4a52..4feb4225d6 100644
--- a/bundle/libraries/match.go
+++ b/bundle/libraries/match.go
@@ -62,6 +62,10 @@ func validateTaskLibraries(libs []compute.Library, b *bundle.Bundle) error {
 
 func validateEnvironments(envs []jobs.JobEnvironment, b *bundle.Bundle) error {
 	for _, env := range envs {
+		if env.Spec == nil {
+			continue
+		}
+
 		for _, dep := range env.Spec.Dependencies {
 			matches, err := filepath.Glob(filepath.Join(b.RootPath, dep))
 			if err != nil {
diff --git a/bundle/tests/enviroment_key_test.go b/bundle/tests/enviroment_key_test.go
index 3e12ddb68c..aed3964db5 100644
--- a/bundle/tests/enviroment_key_test.go
+++ b/bundle/tests/enviroment_key_test.go
@@ -1,8 +1,11 @@
 package config_tests
 
 import (
+	"context"
 	"testing"
 
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/libraries"
 	"github.com/stretchr/testify/require"
 )
 
@@ -10,3 +13,11 @@ func TestEnvironmentKeySupported(t *testing.T) {
 	_, diags := loadTargetWithDiags("./python_wheel/environment_key", "default")
 	require.Empty(t, diags)
 }
+
+func TestEnvironmentKeyProvidedAndNoPanic(t *testing.T) {
+	b, diags := loadTargetWithDiags("./environment_key_only", "default")
+	require.Empty(t, diags)
+
+	diags = bundle.Apply(context.Background(), b, libraries.ValidateLocalLibrariesExist())
+	require.Empty(t, diags)
+}
diff --git a/bundle/tests/environment_key_only/databricks.yml b/bundle/tests/environment_key_only/databricks.yml
new file mode 100644
index 0000000000..caa34f8e35
--- /dev/null
+++ b/bundle/tests/environment_key_only/databricks.yml
@@ -0,0 +1,16 @@
+bundle:
+  name: environment_key_only
+
+resources:
+  jobs:
+    test_job:
+      name: "My Wheel Job"
+      tasks:
+        - task_key: TestTask
+          existing_cluster_id: "0717-132531-5opeqon1"
+          python_wheel_task:
+            package_name: "my_test_code"
+            entry_point: "run"
+          environment_key:  "test_env"
+      environments:
+        - environment_key:  "test_env"