From 23f4ea3c0fd501b4c2c8b396f0ea3a37e6c0d8a7 Mon Sep 17 00:00:00 2001
From: Tanmay Rustagi <tanmay.rustagi@databricks.com>
Date: Fri, 31 Oct 2025 16:53:34 +0530
Subject: [PATCH] [Internal] Unified Terraform Provider: Add plan validations
 for wrong workspace_id

---
 .../products/sharing/data_shares_acc_test.go  | 29 +++++++++++++++++
 .../products/sharing/resource_share.go        | 31 +++++++++++++++++++
 .../sharing/resource_share_acc_test.go        | 29 +++++++++++++++++
 3 files changed, 89 insertions(+)

diff --git a/internal/providers/pluginfw/products/sharing/data_shares_acc_test.go b/internal/providers/pluginfw/products/sharing/data_shares_acc_test.go
index 01b726c082..5e565a0704 100644
--- a/internal/providers/pluginfw/products/sharing/data_shares_acc_test.go
+++ b/internal/providers/pluginfw/products/sharing/data_shares_acc_test.go
@@ -1,11 +1,13 @@
 package sharing_test
 
 import (
+	"context"
 	"fmt"
 	"regexp"
 	"strconv"
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/terraform-provider-databricks/internal/acceptance"
 	"github.com/hashicorp/terraform-plugin-testing/terraform"
 	"github.com/stretchr/testify/assert"
@@ -135,3 +137,30 @@ func TestAccSharesData_ProviderConfig_Mismatched(t *testing.T) {
 		),
 	})
 }
+
+func TestAccSharesData_ProviderConfig_MismatchReapply(t *testing.T) {
+	acceptance.LoadUcwsEnv(t)
+	ctx := context.Background()
+	w := databricks.Must(databricks.NewWorkspaceClient())
+	workspaceID, err := w.CurrentWorkspaceID(ctx)
+	require.NoError(t, err)
+	workspaceIDStr := strconv.FormatInt(workspaceID, 10)
+	acceptance.UnityWorkspaceLevel(t, acceptance.Step{
+		Template: preTestTemplateSchema + dataSourceSharesTemplate(`
+			provider_config = {
+				workspace_id = "123"
+			}
+		`),
+		ExpectError: regexp.MustCompile(
+			`(?s)failed to get workspace client.*workspace_id mismatch` +
+				`.*please check the workspace_id provided in ` +
+				`provider_config`,
+		),
+	}, acceptance.Step{
+		Template: preTestTemplateSchema + dataSourceSharesTemplate(fmt.Sprintf(`
+			provider_config = {
+				workspace_id = "%s"
+			}
+		`, workspaceIDStr)),
+	})
+}
diff --git a/internal/providers/pluginfw/products/sharing/resource_share.go b/internal/providers/pluginfw/products/sharing/resource_share.go
index 14f88aef87..6df6deef7d 100644
--- a/internal/providers/pluginfw/products/sharing/resource_share.go
+++ b/internal/providers/pluginfw/products/sharing/resource_share.go
@@ -182,6 +182,37 @@ func (d *ShareResource) Configure(ctx context.Context, req resource.ConfigureReq
 	}
 }
 
+func (r *ShareResource) ModifyPlan(ctx context.Context, req resource.ModifyPlanRequest, resp *resource.ModifyPlanResponse) {
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
+
+	var plan ShareInfoExtended
+	resp.Diagnostics.Append(req.Plan.Get(ctx, &plan)...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+	var planGoSDK sharing.ShareInfo
+	resp.Diagnostics.Append(converters.TfSdkToGoSdkStruct(ctx, plan, &planGoSDK)...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+	workspaceID, diags := tfschema.GetWorkspaceID_SdkV2(ctx, plan.ProviderConfig)
+	resp.Diagnostics.Append(diags...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+	// Validate the workspace ID by making sure we are able to get a workspace client
+	// for the new workspace ID.
+	_, clientDiags := r.Client.GetWorkspaceClientForUnifiedProviderWithDiagnostics(ctx, workspaceID)
+	resp.Diagnostics.Append(clientDiags...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+}
+
 func (r *ShareResource) Create(ctx context.Context, req resource.CreateRequest, resp *resource.CreateResponse) {
 	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 
diff --git a/internal/providers/pluginfw/products/sharing/resource_share_acc_test.go b/internal/providers/pluginfw/products/sharing/resource_share_acc_test.go
index e82d65f51a..79be07bb85 100644
--- a/internal/providers/pluginfw/products/sharing/resource_share_acc_test.go
+++ b/internal/providers/pluginfw/products/sharing/resource_share_acc_test.go
@@ -654,6 +654,35 @@ func TestAccShare_ProviderConfig_Mismatched(t *testing.T) {
 				`.*please check the workspace_id provided in ` +
 				`provider_config`,
 		),
+		PlanOnly: true,
+	})
+}
+
+func TestAccShare_ProviderConfig_MismatchedReapply(t *testing.T) {
+	acceptance.LoadUcwsEnv(t)
+	ctx := context.Background()
+	w := databricks.Must(databricks.NewWorkspaceClient())
+	workspaceID, err := w.CurrentWorkspaceID(ctx)
+	require.NoError(t, err)
+	workspaceIDStr := strconv.FormatInt(workspaceID, 10)
+	acceptance.UnityWorkspaceLevel(t, acceptance.Step{
+		Template: preTestTemplateSchema + shareTemplate(`
+			provider_config {
+				workspace_id = "123"
+			}
+		`),
+		ExpectError: regexp.MustCompile(
+			`(?s)failed to get workspace client.*workspace_id mismatch` +
+				`.*please check the workspace_id provided in ` +
+				`provider_config`,
+		),
+		PlanOnly: true,
+	}, acceptance.Step{
+		Template: preTestTemplateSchema + shareTemplate(fmt.Sprintf(`
+			provider_config {
+				workspace_id = "%s"
+			}
+		`, workspaceIDStr)),
 	})
 }