[Internal] Support adding context in resources and data sources (#4085)

## Changes
<!-- Summary of your changes that are easy to understand -->
Set `resource` and `data` context for each resource in their respective
CRUD operations. Two methods are defined to make setting resources /
datasources as default easier: `GetDatabricksProductionName` and
`GetDatabricksStagingName`.

Note: Library update isn't supported and hence we shouldn't inject
useragent there.

Every resource has:

`````
        ctx = pluginfwcommon.SetResourceNameInContext(ctx, resourceName)
	w, diags := r.Client.GetWorkspaceClient()
	resp.Diagnostics.Append(diags...)
	if resp.Diagnostics.HasError() {
		return
	}
`````
We could avoid the repetition by creating a common base struct but since
this is going to be autogenerated soon, I think it's fine to have these
as is due to return statements / clarity of reading the code.


## Tests
<!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->
Unit tests
- [ ] `make test` run locally
- [ ] relevant change in `docs/` folder
- [ ] covered with integration tests in `internal/acceptance`
- [ ] relevant acceptance tests are passing
- [ ] using Go SDK

internal/providers/pluginfw/common/common.go

@@ -45,3 +45,13 @@ func ConfigureResource(req resource.ConfigureRequest, resp *resource.ConfigureRe
 	}
 	return client
 }
+
+// GetDatabricksStagingName returns the resource name for a given resource with _pluginframework suffix.
+// Once a migrated resource is ready to be used as default, the Metadata method for that resource should be updated to use GetDatabricksProductionName.
+func GetDatabricksStagingName(name string) string {
+	return fmt.Sprintf("databricks_%s_pluginframework", name)
+}
+
+func GetDatabricksProductionName(name string) string {
+	return fmt.Sprintf("databricks_%s", name)
+}

internal/providers/pluginfw/common/common_test.go

@@ -0,0 +1,21 @@
+package common
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetDatabricksStagingName(t *testing.T) {
+	resourceName := "test"
+	expected := "databricks_test_pluginframework"
+	result := GetDatabricksStagingName(resourceName)
+	assert.Equal(t, expected, result, "GetDatabricksStagingName should return the expected staging name")
+}
+
+func TestGetDatabricksProductionName(t *testing.T) {
+	resourceName := "test"
+	expected := "databricks_test"
+	result := GetDatabricksProductionName(resourceName)
+	assert.Equal(t, expected, result, "GetDatabricksProductionName should return the expected production name")
+}

internal/providers/pluginfw/context/context.go

@@ -0,0 +1,15 @@
+package common
+
+import (
+	"context"
+
+	"github.com/databricks/databricks-sdk-go/useragent"
+)
+
+func SetResourceNameInContext(ctx context.Context, resourceName string) context.Context {
+	return useragent.InContext(ctx, "resource", resourceName)
+}
+
+func SetDataSourceNameInContext(ctx context.Context, dataSourceName string) context.Context {
+	return useragent.InContext(ctx, "data", dataSourceName)
+}

internal/providers/pluginfw/context/context_test.go

@@ -0,0 +1,27 @@
+package common
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSetResourceNameInContext(t *testing.T) {
+	ctx := context.Background()
+	resourceKey := "resource"
+	resourceName := "test-resource"
+	actualContext := SetResourceNameInContext(ctx, resourceName)
+	expectedContext := useragent.InContext(ctx, resourceKey, resourceName)
+	assert.Equal(t, expectedContext, actualContext)
+}
+
+func TestSetDataSourceNameInContext(t *testing.T) {
+	ctx := context.Background()
+	dataSourceKey := "data"
+	dataSourceName := "test-datasource"
+	actualContext := SetDataSourceNameInContext(ctx, dataSourceName)
+	expectedContext := useragent.InContext(ctx, dataSourceKey, dataSourceName)
+	assert.Equal(t, expectedContext, actualContext)
+}

internal/providers/pluginfw/resources/cluster/data_cluster.go

@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/terraform-provider-databricks/common"
 	pluginfwcommon "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/common"
+	pluginfwcontext "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/context"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/converters"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/tfschema"
 	"github.com/databricks/terraform-provider-databricks/internal/service/compute_tf"
@@ -18,6 +19,8 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework/types"
 )
 
+const dataSourceName = "cluster"
+
 func DataSourceCluster() datasource.DataSource {
 	return &ClusterDataSource{}
 }
@@ -35,7 +38,7 @@ type ClusterInfo struct {
 }
 
 func (d *ClusterDataSource) Metadata(ctx context.Context, req datasource.MetadataRequest, resp *datasource.MetadataResponse) {
-	resp.TypeName = "databricks_cluster_pluginframework"
+	resp.TypeName = pluginfwcommon.GetDatabricksStagingName(dataSourceName)
 }
 
 func (d *ClusterDataSource) Schema(ctx context.Context, req datasource.SchemaRequest, resp *datasource.SchemaResponse) {
@@ -67,6 +70,7 @@ func validateClustersList(ctx context.Context, clusters []compute_tf.ClusterDeta
 }
 
 func (d *ClusterDataSource) Read(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse) {
+	ctx = pluginfwcontext.SetDataSourceNameInContext(ctx, dataSourceName)
 	w, diags := d.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {

internal/providers/pluginfw/resources/library/resource_library.go

@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/terraform-provider-databricks/clusters"
 	"github.com/databricks/terraform-provider-databricks/common"
 	pluginfwcommon "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/common"
+	pluginfwcontext "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/context"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/converters"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/tfschema"
 	"github.com/databricks/terraform-provider-databricks/internal/service/compute_tf"
@@ -25,6 +26,7 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 )
 
+const resourceName = "library"
 const libraryDefaultInstallationTimeout = 15 * time.Minute
 
 var _ resource.ResourceWithConfigure = &LibraryResource{}
@@ -67,7 +69,7 @@ type LibraryResource struct {
 }
 
 func (r *LibraryResource) Metadata(ctx context.Context, req resource.MetadataRequest, resp *resource.MetadataResponse) {
-	resp.TypeName = "databricks_library_pluginframework"
+	resp.TypeName = pluginfwcommon.GetDatabricksStagingName(resourceName)
 }
 
 func (r *LibraryResource) Schema(ctx context.Context, req resource.SchemaRequest, resp *resource.SchemaResponse) {
@@ -102,6 +104,7 @@ func (r *LibraryResource) Configure(ctx context.Context, req resource.ConfigureR
 }
 
 func (r *LibraryResource) Create(ctx context.Context, req resource.CreateRequest, resp *resource.CreateResponse) {
+	ctx = pluginfwcontext.SetResourceNameInContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -144,6 +147,7 @@ func (r *LibraryResource) Create(ctx context.Context, req resource.CreateRequest
 }
 
 func (r *LibraryResource) Read(ctx context.Context, req resource.ReadRequest, resp *resource.ReadResponse) {
+	ctx = pluginfwcontext.SetResourceNameInContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -181,6 +185,7 @@ func (r *LibraryResource) Update(ctx context.Context, req resource.UpdateRequest
 }
 
 func (r *LibraryResource) Delete(ctx context.Context, req resource.DeleteRequest, resp *resource.DeleteResponse) {
+	ctx = pluginfwcontext.SetResourceNameInContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {

internal/providers/pluginfw/resources/qualitymonitor/resource_quality_monitor.go

@@ -11,6 +11,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/terraform-provider-databricks/common"
 	pluginfwcommon "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/common"
+	pluginfwcontext "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/context"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/converters"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/tfschema"
 	"github.com/databricks/terraform-provider-databricks/internal/service/catalog_tf"
@@ -21,6 +22,8 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework/types"
 )
 
+const resourceName = "quality_monitor"
+
 const qualityMonitorDefaultProvisionTimeout = 15 * time.Minute
 
 var _ resource.ResourceWithConfigure = &QualityMonitorResource{}
@@ -62,7 +65,7 @@ type QualityMonitorResource struct {
 }
 
 func (r *QualityMonitorResource) Metadata(ctx context.Context, req resource.MetadataRequest, resp *resource.MetadataResponse) {
-	resp.TypeName = "databricks_quality_monitor_pluginframework"
+	resp.TypeName = pluginfwcommon.GetDatabricksStagingName(resourceName)
 }
 
 func (r *QualityMonitorResource) Schema(ctx context.Context, req resource.SchemaRequest, resp *resource.SchemaResponse) {
@@ -95,6 +98,7 @@ func (d *QualityMonitorResource) ImportState(ctx context.Context, req resource.I
 }
 
 func (r *QualityMonitorResource) Create(ctx context.Context, req resource.CreateRequest, resp *resource.CreateResponse) {
+	ctx = pluginfwcontext.SetResourceNameInContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -131,6 +135,7 @@ func (r *QualityMonitorResource) Create(ctx context.Context, req resource.Create
 }
 
 func (r *QualityMonitorResource) Read(ctx context.Context, req resource.ReadRequest, resp *resource.ReadResponse) {
+	ctx = pluginfwcontext.SetResourceNameInContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -161,6 +166,7 @@ func (r *QualityMonitorResource) Read(ctx context.Context, req resource.ReadRequ
 }
 
 func (r *QualityMonitorResource) Update(ctx context.Context, req resource.UpdateRequest, resp *resource.UpdateResponse) {
+	ctx = pluginfwcontext.SetResourceNameInContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -203,6 +209,7 @@ func (r *QualityMonitorResource) Update(ctx context.Context, req resource.Update
 }
 
 func (r *QualityMonitorResource) Delete(ctx context.Context, req resource.DeleteRequest, resp *resource.DeleteResponse) {
+	ctx = pluginfwcontext.SetResourceNameInContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {

internal/providers/pluginfw/resources/volume/data_volumes.go

@@ -8,13 +8,16 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/terraform-provider-databricks/common"
 	pluginfwcommon "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/common"
+	pluginfwcontext "github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/context"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/converters"
 	"github.com/databricks/terraform-provider-databricks/internal/providers/pluginfw/tfschema"
 	"github.com/hashicorp/terraform-plugin-framework/datasource"
 	"github.com/hashicorp/terraform-plugin-framework/datasource/schema"
 	"github.com/hashicorp/terraform-plugin-framework/types"
 )
 
+const dataSourceName = "volumes"
+
 func DataSourceVolumes() datasource.DataSource {
 	return &VolumesDataSource{}
 }
@@ -32,7 +35,7 @@ type VolumesList struct {
 }
 
 func (d *VolumesDataSource) Metadata(ctx context.Context, req datasource.MetadataRequest, resp *datasource.MetadataResponse) {
-	resp.TypeName = "databricks_volumes_pluginframework"
+	resp.TypeName = pluginfwcommon.GetDatabricksStagingName(dataSourceName)
 }
 
 func (d *VolumesDataSource) Schema(ctx context.Context, req datasource.SchemaRequest, resp *datasource.SchemaResponse) {
@@ -50,6 +53,7 @@ func (d *VolumesDataSource) Configure(_ context.Context, req datasource.Configur
 }
 
 func (d *VolumesDataSource) Read(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse) {
+	ctx = pluginfwcontext.SetDataSourceNameInContext(ctx, dataSourceName)
 	w, diags := d.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {