From df613759951bde518c40152fb806662a70271e79 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Tue, 8 Oct 2024 10:14:38 +0200
Subject: [PATCH] Address CR comments

---
 .../mutator/python/python_diagnostics.go      |  1 +
 .../config/mutator/python/python_locations.go | 27 +++++++++++++++----
 .../mutator/python/python_locations_test.go   |  3 ---
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/bundle/config/mutator/python/python_diagnostics.go b/bundle/config/mutator/python/python_diagnostics.go
index 12822065..7a1e13b4 100644
--- a/bundle/config/mutator/python/python_diagnostics.go
+++ b/bundle/config/mutator/python/python_diagnostics.go
@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/libs/dyn"
 )
 
+// pythonDiagnostic is a single entry in diagnostics.json
 type pythonDiagnostic struct {
 	Severity pythonSeverity           `json:"severity"`
 	Summary  string                   `json:"summary"`
diff --git a/bundle/config/mutator/python/python_locations.go b/bundle/config/mutator/python/python_locations.go
index 97b7e170..eb6698b4 100644
--- a/bundle/config/mutator/python/python_locations.go
+++ b/bundle/config/mutator/python/python_locations.go
@@ -9,9 +9,28 @@ import (
 	"github.com/databricks/cli/libs/dyn"
 )
 
+// generatedFileName is used as the virtual file name for YAML generated by PyDABs.
+//
+// mergePythonLocations replaces dyn.Location with generatedFileName with locations loaded
+// from locations.json
 const generatedFileName = "__generated_by_pydabs__.yml"
 
 // pythonLocations is data structure for efficient location lookup for a given path
+//
+// Locations form a tree, and we assign locations of the closest ancestor to each dyn.Value based on its path.
+// We implement it as a trie (prefix tree) where keys are components of the path. With that, lookups are O(n)
+// where n is the number of components in the path.
+//
+// For example, with locations.json:
+//
+//	{"path": "resources.jobs.job_0", "file": "src/examples/job_0.py", "line": 3, "column": 5}
+//	{"path": "resources.jobs.job_0.tasks[0].task_key", "file": "src/examples/job_0.py", "line": 10, "column": 5}
+//	{"path": "resources.jobs.job_1", "file": "src/examples/job_1.py", "line": 5, "column": 7}
+//
+// - resources.jobs.job_0.tasks[0].task_key is located at job_0.py:10:5
+//
+//   - resources.jobs.job_0.tasks[0].email_notifications is located at job_0.py:3:5,
+//     because we use the location of the job as the most precise approximation.
 type pythonLocations struct {
 	// descendants referenced by index, e.g. '.foo'
 	keys map[string]*pythonLocations
@@ -26,6 +45,7 @@ type pythonLocations struct {
 	exists bool
 }
 
+// pythonLocationEntry is a single entry in locations.json
 type pythonLocationEntry struct {
 	Path   string `json:"path"`
 	File   string `json:"file"`
@@ -46,15 +66,12 @@ func mergePythonLocations(value dyn.Value, locations *pythonLocations) (dyn.Valu
 			newLocations = append(newLocations, newLocation)
 
 			for _, location := range value.Locations() {
+				// When loaded, dyn.Value created by PyDABs use the virtual file path as their location,
+				// we replace it with newLocation.
 				if filepath.Base(location.File) == generatedFileName {
 					continue
 				}
 
-				// don't add duplicates if location already exists
-				if location == newLocation {
-					continue
-				}
-
 				newLocations = append(newLocations, location)
 			}
 
diff --git a/bundle/config/mutator/python/python_locations_test.go b/bundle/config/mutator/python/python_locations_test.go
index 28cd33de..2dd2785d 100644
--- a/bundle/config/mutator/python/python_locations_test.go
+++ b/bundle/config/mutator/python/python_locations_test.go
@@ -20,7 +20,6 @@ func TestMergeLocations(t *testing.T) {
 		map[string]dyn.Value{
 			"foo": dyn.NewValue(
 				map[string]dyn.Value{
-					"bar": dyn.NewValue("baz", []dyn.Location{yamlLocation, pythonLocation}),
 					"baz": dyn.NewValue("baz", []dyn.Location{yamlLocation}),
 					"qux": dyn.NewValue("baz", []dyn.Location{generatedLocation, yamlLocation}),
 				},
@@ -35,8 +34,6 @@ func TestMergeLocations(t *testing.T) {
 		map[string]dyn.Value{
 			"foo": dyn.NewValue(
 				map[string]dyn.Value{
-					// pythonLocation was present before, we don't duplicate it, only move it to the beginning
-					"bar": dyn.NewValue("baz", []dyn.Location{pythonLocation, yamlLocation}),
 					// pythonLocation is appended to the beginning of the list if absent
 					"baz": dyn.NewValue("baz", []dyn.Location{pythonLocation, yamlLocation}),
 					// generatedLocation is replaced by pythonLocation