diff --git a/bundle/deploy/terraform/init.go b/bundle/deploy/terraform/init.go index d1847cf2..d480242c 100644 --- a/bundle/deploy/terraform/init.go +++ b/bundle/deploy/terraform/init.go @@ -218,6 +218,23 @@ func setProxyEnvVars(ctx context.Context, environ map[string]string, b *bundle.B return nil } +func setUserAgentExtraEnvVar(environ map[string]string, b *bundle.Bundle) error { + var products []string + + if experimental := b.Config.Experimental; experimental != nil { + if experimental.PyDABs.Enabled { + products = append(products, "databricks-pydabs/0.0.0") + } + } + + userAgentExtra := strings.Join(products, " ") + if userAgentExtra != "" { + environ["DATABRICKS_USER_AGENT_EXTRA"] = userAgentExtra + } + + return nil +} + func (m *initialize) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { tfConfig := b.Config.Bundle.Terraform if tfConfig == nil { @@ -262,6 +279,11 @@ func (m *initialize) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnosti return diag.FromErr(err) } + err = setUserAgentExtraEnvVar(environ, b) + if err != nil { + return diag.FromErr(err) + } + // Configure environment variables for auth for Terraform to use. log.Debugf(ctx, "Environment variables for Terraform: %s", strings.Join(maps.Keys(environ), ", ")) err = tf.SetEnv(environ) diff --git a/bundle/deploy/terraform/init_test.go b/bundle/deploy/terraform/init_test.go index 421e9be3..aa9b2f77 100644 --- a/bundle/deploy/terraform/init_test.go +++ b/bundle/deploy/terraform/init_test.go @@ -248,6 +248,27 @@ func TestSetProxyEnvVars(t *testing.T) { assert.ElementsMatch(t, []string{"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY"}, maps.Keys(env)) } +func TestSetUserAgentExtraEnvVar(t *testing.T) { + b := &bundle.Bundle{ + RootPath: t.TempDir(), + Config: config.Root{ + Experimental: &config.Experimental{ + PyDABs: config.PyDABs{ + Enabled: true, + }, + }, + }, + } + + env := make(map[string]string, 0) + err := setUserAgentExtraEnvVar(env, b) + + require.NoError(t, err) + assert.Equal(t, map[string]string{ + "DATABRICKS_USER_AGENT_EXTRA": "databricks-pydabs/0.0.0", + }, env) +} + func TestInheritEnvVars(t *testing.T) { env := map[string]string{}