[Internal] Set SDK used in useragent

databricks · Oct 10, 2024 · 6a7f01f · 6a7f01f
1 parent 2ad3f1b
commit 6a7f01f
Show file tree

Hide file tree

Showing 8 changed files with 37 additions and 18 deletions.
diff --git a/common/context.go b/common/context.go
@@ -5,10 +5,13 @@ import (
 	"strings"
 
 	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/databricks/terraform-provider-databricks/internal/providers/common"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
+const sdkName = "sdkv2"
+
 // AddContextToAllResources ...
 func AddContextToAllResources(p *schema.Provider, prefix string) {
 	for k, r := range p.DataSourcesMap {
@@ -34,6 +37,7 @@ func (f op) addContext(k contextKey, v string) op {
 		case IsData:
 			ctx = useragent.InContext(ctx, "data", v)
 		}
+		ctx = common.SetSDKInContext(ctx, sdkName)
 		ctx = context.WithValue(ctx, k, v)
 		return f(ctx, d, m)
 	}

diff --git a/internal/providers/common/common.go b/internal/providers/common/common.go
@@ -3,6 +3,16 @@
 // Note: This is different from internal/providers which contains the changes that *depends* on both:
 // internal/providers/sdkv2 and internal/providers/pluginfw packages. Whereas, internal/providers/common package contains
 // the changes *used* by both internal/providers/sdkv2 and internal/providers/pluginfw packages.
-package internal
+package common
+
+import (
+	"context"
+
+	"github.com/databricks/databricks-sdk-go/useragent"
+)
 
 const ProviderName = "databricks-tf-provider"
+
+func SetSDKInContext(ctx context.Context, sdkUsed string) context.Context {
+	return useragent.InContext(ctx, "sdk", sdkUsed)
+}
diff --git a/internal/providers/pluginfw/context/context.go b/internal/providers/pluginfw/context/context.go
@@ -1,15 +1,20 @@
-package common
+package context
 
 import (
 	"context"
 
 	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/databricks/terraform-provider-databricks/internal/providers/common"
 )
 
-func SetResourceNameInContext(ctx context.Context, resourceName string) context.Context {
+const SdkName = "pluginframework"
+
+func SetUserAgentInResourceContext(ctx context.Context, resourceName string) context.Context {
+	ctx = common.SetSDKInContext(ctx, SdkName)
 	return useragent.InContext(ctx, "resource", resourceName)
 }
 
-func SetDataSourceNameInContext(ctx context.Context, dataSourceName string) context.Context {
+func SetUserAgentInDataSourceContext(ctx context.Context, dataSourceName string) context.Context {
+	ctx = common.SetSDKInContext(ctx, SdkName)
 	return useragent.InContext(ctx, "data", dataSourceName)
 }
diff --git a/internal/providers/pluginfw/context/context_test.go b/internal/providers/pluginfw/context/context_test.go
@@ -1,4 +1,4 @@
-package common
+package context
 
 import (
 	"context"
@@ -8,20 +8,20 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestSetResourceNameInContext(t *testing.T) {
+func TestSetUserAgentInResourceContext(t *testing.T) {
 	ctx := context.Background()
 	resourceKey := "resource"
 	resourceName := "test-resource"
-	actualContext := SetResourceNameInContext(ctx, resourceName)
+	actualContext := SetUserAgentInResourceContext(ctx, resourceName)
 	expectedContext := useragent.InContext(ctx, resourceKey, resourceName)
 	assert.Equal(t, expectedContext, actualContext)
 }
 
-func TestSetDataSourceNameInContext(t *testing.T) {
+func TestSetUserAgentInDataSourceContext(t *testing.T) {
 	ctx := context.Background()
 	dataSourceKey := "data"
 	dataSourceName := "test-datasource"
-	actualContext := SetDataSourceNameInContext(ctx, dataSourceName)
+	actualContext := SetUserAgentInDataSourceContext(ctx, dataSourceName)
 	expectedContext := useragent.InContext(ctx, dataSourceKey, dataSourceName)
 	assert.Equal(t, expectedContext, actualContext)
 }
diff --git a/internal/providers/pluginfw/resources/cluster/data_cluster.go b/internal/providers/pluginfw/resources/cluster/data_cluster.go
@@ -70,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInDataSourceContext(ctx, dataSourceName)
 	w, diags := d.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {

diff --git a/internal/providers/pluginfw/resources/library/resource_library.go b/internal/providers/pluginfw/resources/library/resource_library.go
@@ -104,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -147,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -185,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {

diff --git a/internal/providers/pluginfw/resources/qualitymonitor/resource_quality_monitor.go b/internal/providers/pluginfw/resources/qualitymonitor/resource_quality_monitor.go
@@ -98,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -135,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -166,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {
@@ -209,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInResourceContext(ctx, resourceName)
 	w, diags := r.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {

diff --git a/internal/providers/pluginfw/resources/volume/data_volumes.go b/internal/providers/pluginfw/resources/volume/data_volumes.go
@@ -53,7 +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)
+	ctx = pluginfwcontext.SetUserAgentInDataSourceContext(ctx, dataSourceName)
 	w, diags := d.Client.GetWorkspaceClient()
 	resp.Diagnostics.Append(diags...)
 	if resp.Diagnostics.HasError() {