From d7962860a011d1d74985fc4beb6508eb9302a55e Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Mon, 9 Feb 2026 14:47:08 +0100
Subject: [PATCH] refactor: combine GetVersion and RuntimeVersion

---
 bundle/config/mutator/configure_wsfs.go |  2 +-
 cmd/root/root.go                        |  2 +-
 libs/dbr/context.go                     | 16 ++++++++--------
 libs/dbr/context_test.go                | 21 ++++++++++++++++-----
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/bundle/config/mutator/configure_wsfs.go b/bundle/config/mutator/configure_wsfs.go
index 71b3fa4191..bc8c91f0ca 100644
--- a/bundle/config/mutator/configure_wsfs.go
+++ b/bundle/config/mutator/configure_wsfs.go
@@ -43,7 +43,7 @@ func (m *configureWSFS) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 	// Writing notebooks via FUSE is only supported for serverless client version 2+.
 	// Since we could only test v2.5, since the platform only allows selecting v2 (which is v2.5 internally),
 	// we restrict FUSE to only be used for v2.5+.
-	v := dbr.GetVersion(ctx)
+	v := dbr.RuntimeVersion(ctx)
 	if v.Type == dbr.ClusterTypeServerless && (v.Major > 2 || (v.Major == 2 && v.Minor >= 5)) {
 		return nil
 	}
diff --git a/cmd/root/root.go b/cmd/root/root.go
index 85ef4dc757..028103ad4f 100644
--- a/cmd/root/root.go
+++ b/cmd/root/root.go
@@ -190,7 +190,7 @@ Stack Trace:
 		Version:         build.GetInfo().Version,
 		Command:         commandStr,
 		OperatingSystem: runtime.GOOS,
-		DbrVersion:      dbr.RuntimeVersion(ctx),
+		DbrVersion:      dbr.RuntimeVersion(ctx).String(),
 		ExecutionTimeMs: time.Since(startTime).Milliseconds(),
 		ExitCode:        int64(exitCode),
 	})
diff --git a/libs/dbr/context.go b/libs/dbr/context.go
index dbf469e617..8aedc4a07d 100644
--- a/libs/dbr/context.go
+++ b/libs/dbr/context.go
@@ -47,6 +47,10 @@ type Version struct {
 	Raw   string
 }
 
+func (v Version) String() string {
+	return v.Raw
+}
+
 // ParseVersion parses a DBR version string and returns structured version info.
 // Examples:
 //   - "16.3" -> Interactive, Major=16, Minor=3
@@ -122,17 +126,13 @@ func RunsOnRuntime(ctx context.Context) bool {
 	return v.(Environment).IsDbr
 }
 
-func RuntimeVersion(ctx context.Context) string {
+// RuntimeVersion returns the parsed runtime version from the context.
+// It expects a context returned by [DetectRuntime] or [MockRuntime].
+func RuntimeVersion(ctx context.Context) Version {
 	v := ctx.Value(dbrKey)
 	if v == nil {
 		panic("dbr.RuntimeVersion called without calling dbr.DetectRuntime first")
 	}
 
-	return v.(Environment).Version
-}
-
-// GetVersion returns the parsed runtime version from the context.
-// It expects a context returned by [DetectRuntime] or [MockRuntime].
-func GetVersion(ctx context.Context) Version {
-	return ParseVersion(RuntimeVersion(ctx))
+	return ParseVersion(v.(Environment).Version)
 }
diff --git a/libs/dbr/context_test.go b/libs/dbr/context_test.go
index a72202b9a6..79d6edbcdf 100644
--- a/libs/dbr/context_test.go
+++ b/libs/dbr/context_test.go
@@ -81,8 +81,8 @@ func TestContext_RunsOnRuntimeWithMock(t *testing.T) {
 
 func TestContext_RuntimeVersionWithMock(t *testing.T) {
 	ctx := context.Background()
-	assert.Equal(t, "15.4", RuntimeVersion(MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"})))
-	assert.Empty(t, RuntimeVersion(MockRuntime(ctx, Environment{})))
+	assert.Equal(t, "15.4", RuntimeVersion(MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"})).String())
+	assert.Empty(t, RuntimeVersion(MockRuntime(ctx, Environment{})).String())
 }
 
 func TestParseVersion_Serverless(t *testing.T) {
@@ -153,19 +153,30 @@ func TestClusterType_String(t *testing.T) {
 	assert.Equal(t, "unknown", ClusterTypeUnknown.String())
 }
 
-func TestContext_GetVersion(t *testing.T) {
+func TestVersion_String(t *testing.T) {
+	v := ParseVersion("16.3")
+	assert.Equal(t, "16.3", v.String())
+
+	v = ParseVersion("client.4.9")
+	assert.Equal(t, "client.4.9", v.String())
+
+	v = ParseVersion("")
+	assert.Equal(t, "", v.String())
+}
+
+func TestContext_RuntimeVersionParsed(t *testing.T) {
 	ctx := context.Background()
 
 	// Test serverless version
 	serverlessCtx := MockRuntime(ctx, Environment{IsDbr: true, Version: "client.4.9"})
-	v := GetVersion(serverlessCtx)
+	v := RuntimeVersion(serverlessCtx)
 	assert.Equal(t, ClusterTypeServerless, v.Type)
 	assert.Equal(t, 4, v.Major)
 	assert.Equal(t, 9, v.Minor)
 
 	// Test interactive version
 	interactiveCtx := MockRuntime(ctx, Environment{IsDbr: true, Version: "17.3"})
-	v = GetVersion(interactiveCtx)
+	v = RuntimeVersion(interactiveCtx)
 	assert.Equal(t, ClusterTypeInteractive, v.Type)
 	assert.Equal(t, 17, v.Major)
 	assert.Equal(t, 3, v.Minor)