Harden dashboard UID handling and JSON normalization

Amp-Thread-ID: https://ampcode.com/threads/T-019c9bc7-a165-774f-b025-2e92d6e0b9d5
Co-authored-by: Amp <amp@ampcode.com>

Author: toppercodes

axiom/data_source_dashboard.go

@@ -62,7 +62,7 @@ func (d *DashboardDataSource) Schema(_ context.Context, _ datasource.SchemaReque
 			},
 			"id": schema.StringAttribute{
 				Computed:            true,
-				MarkdownDescription: "Dashboard identifier (same value as `uid`).",
+				MarkdownDescription: "Internal dashboard identifier returned by the API.",
 			},
 			"dashboard": schema.StringAttribute{
 				Computed:            true,
@@ -116,7 +116,7 @@ func (d *DashboardDataSource) Read(ctx context.Context, req datasource.ReadReque
 		return
 	}
 
-	dashboardJSON, err := normalizeDashboardRaw(dashboard.Dashboard)
+	dashboardJSON, err := normalizeDashboardRaw(dashboard.Dashboard, types.StringValue(string(dashboard.Dashboard)))
 	if err != nil {
 		resp.Diagnostics.AddError("Failed to read dashboard", fmt.Sprintf("Unable to normalize dashboard payload: %s", err))
 		return

axiom/resource_dashboard.go

@@ -19,6 +19,15 @@ import (
 	"github.com/axiomhq/axiom-go/axiom"
 )
 
+var dashboardServerManagedFields = map[string]struct{}{
+	"id":        {},
+	"version":   {},
+	"createdAt": {},
+	"updatedAt": {},
+	"createdBy": {},
+	"updatedBy": {},
+}
+
 var (
 	_ resource.Resource                = &DashboardResource{}
 	_ resource.ResourceWithImportState = &DashboardResource{}
@@ -75,7 +84,7 @@ func (r *DashboardResource) Schema(_ context.Context, _ resource.SchemaRequest,
 		Attributes: map[string]schema.Attribute{
 			"id": schema.StringAttribute{
 				Computed:            true,
-				MarkdownDescription: "Dashboard identifier (same value as `uid`).",
+				MarkdownDescription: "Internal dashboard identifier returned by the API.",
 				PlanModifiers: []planmodifier.String{
 					stringplanmodifier.UseStateForUnknown(),
 				},
@@ -179,7 +188,7 @@ func (r *DashboardResource) Create(ctx context.Context, req resource.CreateReque
 		return
 	}
 
-	state, err := flattenDashboardResource(created.Dashboard, plan.Overwrite)
+	state, err := flattenDashboardResource(created.Dashboard, plan.Overwrite, plan.Dashboard)
 	if err != nil {
 		resp.Diagnostics.AddError("Failed to create dashboard", err.Error())
 		return
@@ -222,7 +231,7 @@ func (r *DashboardResource) Read(ctx context.Context, req resource.ReadRequest,
 		return
 	}
 
-	flattened, err := flattenDashboardResource(*dashboard, state.Overwrite)
+	flattened, err := flattenDashboardResource(*dashboard, state.Overwrite, state.Dashboard)
 	if err != nil {
 		resp.Diagnostics.AddError("Failed to read dashboard", err.Error())
 		return
@@ -271,7 +280,7 @@ func (r *DashboardResource) Update(ctx context.Context, req resource.UpdateReque
 		return
 	}
 
-	flattened, err := flattenDashboardResource(updated.Dashboard, plan.Overwrite)
+	flattened, err := flattenDashboardResource(updated.Dashboard, plan.Overwrite, plan.Dashboard)
 	if err != nil {
 		resp.Diagnostics.AddError("Failed to update dashboard", err.Error())
 		return
@@ -310,7 +319,7 @@ func (r *DashboardResource) ImportState(ctx context.Context, req resource.Import
 func dashboardUpsertPayloadFromModel(plan DashboardResourceModel, fallbackUID string, currentVersion int64, isCreate bool) (dashboardUpsertRequest, string, diag.Diagnostics) {
 	var diags diag.Diagnostics
 
-	normalizedDashboard, err := normalizeDashboardString(plan.Dashboard.ValueString())
+	normalizedDashboard, dashboardMap, err := normalizeDashboardString(plan.Dashboard.ValueString())
 	if err != nil {
 		diags.AddError("Invalid dashboard JSON", fmt.Sprintf("`dashboard` must be valid JSON: %s", err))
 		return dashboardUpsertRequest{}, "", diags
@@ -321,6 +330,15 @@ func dashboardUpsertPayloadFromModel(plan DashboardResourceModel, fallbackUID st
 		uid = plan.UID.ValueString()
 	}
 
+	dashboardUID, hasDashboardUID := stringValueFromMap(dashboardMap, "uid")
+	if hasDashboardUID {
+		if uid != "" && uid != dashboardUID {
+			diags.AddError("UID mismatch", "`uid` must match `dashboard.uid` when both are set.")
+			return dashboardUpsertRequest{}, "", diags
+		}
+		uid = dashboardUID
+	}
+
 	payload := dashboardUpsertRequest{
 		Dashboard: normalizedDashboard,
 		UID:       uid,
@@ -356,8 +374,8 @@ func dashboardUIDFromState(state DashboardResourceModel) string {
 	return ""
 }
 
-func flattenDashboardResource(in dashboardResourcePayload, overwrite types.Bool) (DashboardResourceModel, error) {
-	normalizedDashboard, err := normalizeDashboardRaw(in.Dashboard)
+func flattenDashboardResource(in dashboardResourcePayload, overwrite types.Bool, configuredDashboard types.String) (DashboardResourceModel, error) {
+	normalizedDashboard, err := normalizeDashboardRaw(in.Dashboard, configuredDashboard)
 	if err != nil {
 		return DashboardResourceModel{}, fmt.Errorf("unable to normalize dashboard document: %w", err)
 	}
@@ -396,21 +414,21 @@ func decodeDashboardResource(raw []byte) (*dashboardResourcePayload, error) {
 	return out, nil
 }
 
-func normalizeDashboardString(raw string) (json.RawMessage, error) {
+func normalizeDashboardString(raw string) (json.RawMessage, map[string]any, error) {
 	parsed := make(map[string]any)
 	if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	normalized, err := json.Marshal(parsed)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return normalized, nil
+	return normalized, parsed, nil
 }
 
-func normalizeDashboardRaw(raw json.RawMessage) (string, error) {
+func normalizeDashboardRaw(raw json.RawMessage, configured types.String) (string, error) {
 	if len(raw) == 0 {
 		return "", errors.New("dashboard payload is empty")
 	}
@@ -420,6 +438,14 @@ func normalizeDashboardRaw(raw json.RawMessage) (string, error) {
 		return "", err
 	}
 
+	for field := range dashboardServerManagedFields {
+		delete(parsed, field)
+	}
+
+	if !dashboardConfigHasUID(configured) {
+		delete(parsed, "uid")
+	}
+
 	normalized, err := json.Marshal(parsed)
 	if err != nil {
 		return "", err
@@ -428,6 +454,34 @@ func normalizeDashboardRaw(raw json.RawMessage) (string, error) {
 	return string(normalized), nil
 }
 
+func dashboardConfigHasUID(configured types.String) bool {
+	if configured.IsNull() || configured.IsUnknown() || configured.ValueString() == "" {
+		return false
+	}
+
+	parsed := make(map[string]any)
+	if err := json.Unmarshal([]byte(configured.ValueString()), &parsed); err != nil {
+		return false
+	}
+
+	_, ok := stringValueFromMap(parsed, "uid")
+	return ok
+}
+
+func stringValueFromMap(in map[string]any, key string) (string, bool) {
+	v, ok := in[key]
+	if !ok {
+		return "", false
+	}
+
+	s, ok := v.(string)
+	if !ok || s == "" {
+		return "", false
+	}
+
+	return s, true
+}
+
 func addDashboardUpdateError(resp *resource.UpdateResponse, err error, uid string, localVersion int64) {
 	addDashboardWriteErrorDiagnostics(&resp.Diagnostics, err, uid, localVersion)
 }

axiom/resource_dashboard_integration_test.go

@@ -136,6 +136,48 @@ func TestAccAxiomDashboardResource_ServerGeneratedUID(t *testing.T) {
 	})
 }
 
+func TestAccAxiomDashboardResource_UIDInDashboardDocument(t *testing.T) {
+	if os.Getenv("TF_ACC") == "" {
+		t.Skip("acceptance tests skipped unless TF_ACC is set")
+	}
+	testAccPreCheck(t)
+
+	client, err := ax.NewClient()
+	assert.NoError(t, err)
+
+	uid := "tf-dashboard-doc-uid-" + strings.ReplaceAll(uuid.NewString(), "_", "-")
+	resourceName := "axiom_dashboard.test"
+
+	resource.Test(t, resource.TestCase{
+		ProtoV6ProviderFactories: map[string]func() (tfprotov6.ProviderServer, error){
+			"axiom": providerserver.NewProtocol6WithError(NewAxiomProvider()),
+		},
+		CheckDestroy: testAccCheckAxiomResourcesDestroyed(client),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccAxiomDashboardConfigDocumentUID(uid, "doc-uid-create", false, true),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAxiomResourcesExist(client, resourceName),
+					resource.TestCheckResourceAttr(resourceName, "uid", uid),
+				),
+			},
+			{
+				Config: testAccAxiomDashboardConfigDocumentUID(uid, "doc-uid-update", false, false),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAxiomResourcesExist(client, resourceName),
+					resource.TestCheckResourceAttr(resourceName, "uid", uid),
+					resource.TestCheckResourceAttrWith(resourceName, "dashboard", func(v string) error {
+						if strings.Contains(v, `"uid"`) {
+							return fmt.Errorf("expected dashboard JSON in state to omit uid when not configured, got %s", v)
+						}
+						return nil
+					}),
+				),
+			},
+		},
+	})
+}
+
 func testAccAxiomDashboardConfig(uid, name string, overwrite bool) string {
 	return fmt.Sprintf(`
 provider "axiom" {
@@ -183,6 +225,34 @@ resource "axiom_dashboard" "test" {
 `, os.Getenv("AXIOM_TOKEN"), os.Getenv("AXIOM_URL"), overwrite, name)
 }
 
+func testAccAxiomDashboardConfigDocumentUID(uid, name string, overwrite, includeUIDInDocument bool) string {
+	uidField := ""
+	if includeUIDInDocument {
+		uidField = fmt.Sprintf("uid             = %q\n", uid)
+	}
+
+	return fmt.Sprintf(`
+provider "axiom" {
+  api_token = %q
+  base_url  = %q
+}
+
+resource "axiom_dashboard" "test" {
+  overwrite = %t
+  dashboard = jsonencode({
+    %sname            = %q
+    description     = "terraform acceptance dashboard"
+    refreshTime     = 60
+    schemaVersion   = 2
+    timeWindowStart = "qr-now-1h"
+    timeWindowEnd   = "qr-now"
+    charts          = []
+    layout          = []
+  })
+}
+`, os.Getenv("AXIOM_TOKEN"), os.Getenv("AXIOM_URL"), overwrite, uidField, name)
+}
+
 func testAccCaptureDashboardUID(resourceName string, out *string) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		rs, ok := s.RootModule().Resources[resourceName]