dataset: resolve edge deployment from org default

Amp-Thread-ID: https://ampcode.com/threads/T-019d0c27-9d81-71e1-9db5-23c7af7a04a8
Co-authored-by: Amp <amp@ampcode.com>

Author: toppercodes

axiom/data_source_dataset.go

@@ -71,5 +71,11 @@ func (d *DatasetDataSource) Read(ctx context.Context, req datasource.ReadRequest
 		return
 	}
 
-	resp.Diagnostics.Append(resp.State.Set(ctx, flattenDataset(ds))...)
+	state, err := flattenDatasetWithOrgDefault(ctx, d.client, ds)
+	if err != nil {
+		resp.Diagnostics.AddWarning("Unable to resolve default edge deployment", err.Error())
+		state = flattenDataset(ds, "")
+	}
+
+	resp.Diagnostics.Append(resp.State.Set(ctx, state)...)
 }

axiom/resource_dataset.go

@@ -185,8 +185,14 @@ func (r *DatasetResource) Create(ctx context.Context, req resource.CreateRequest
 		return
 	}
 
+	state, err := flattenDatasetWithOrgDefault(ctx, r.client, ds)
+	if err != nil {
+		resp.Diagnostics.AddWarning("Unable to resolve default edge deployment", err.Error())
+		state = flattenDataset(ds, "")
+	}
+
 	// Set state immediately after creation to avoid orphaned resources
-	resp.Diagnostics.Append(resp.State.Set(ctx, flattenDataset(ds))...)
+	resp.Diagnostics.Append(resp.State.Set(ctx, state)...)
 	if resp.Diagnostics.HasError() {
 		return
 	}
@@ -206,7 +212,14 @@ func (r *DatasetResource) Create(ctx context.Context, req resource.CreateRequest
 		}
 
 		ds.MapFields = resMapFields
-		resp.Diagnostics.Append(resp.State.Set(ctx, flattenDataset(ds))...)
+
+		state, err := flattenDatasetWithOrgDefault(ctx, r.client, ds)
+		if err != nil {
+			resp.Diagnostics.AddWarning("Unable to resolve default edge deployment", err.Error())
+			state = flattenDataset(ds, "")
+		}
+
+		resp.Diagnostics.Append(resp.State.Set(ctx, state)...)
 	}
 }
 
@@ -232,7 +245,13 @@ func (r *DatasetResource) Read(ctx context.Context, req resource.ReadRequest, re
 		return
 	}
 
-	resp.Diagnostics.Append(resp.State.Set(ctx, flattenDataset(ds))...)
+	state, err := flattenDatasetWithOrgDefault(ctx, r.client, ds)
+	if err != nil {
+		resp.Diagnostics.AddWarning("Unable to resolve default edge deployment", err.Error())
+		state = flattenDataset(ds, "")
+	}
+
+	resp.Diagnostics.Append(resp.State.Set(ctx, state)...)
 }
 
 func (r *DatasetResource) Update(ctx context.Context, req resource.UpdateRequest, resp *resource.UpdateResponse) {
@@ -263,8 +282,14 @@ func (r *DatasetResource) Update(ctx context.Context, req resource.UpdateRequest
 		return
 	}
 
+	state, err := flattenDatasetWithOrgDefault(ctx, r.client, ds)
+	if err != nil {
+		resp.Diagnostics.AddWarning("Unable to resolve default edge deployment", err.Error())
+		state = flattenDataset(ds, "")
+	}
+
 	// Set state immediately after update to preserve changes
-	resp.Diagnostics.Append(resp.State.Set(ctx, flattenDataset(ds))...)
+	resp.Diagnostics.Append(resp.State.Set(ctx, state)...)
 	if resp.Diagnostics.HasError() {
 		return
 	}
@@ -284,7 +309,14 @@ func (r *DatasetResource) Update(ctx context.Context, req resource.UpdateRequest
 		}
 
 		ds.MapFields = resMapFields
-		resp.Diagnostics.Append(resp.State.Set(ctx, flattenDataset(ds))...)
+
+		state, err := flattenDatasetWithOrgDefault(ctx, r.client, ds)
+		if err != nil {
+			resp.Diagnostics.AddWarning("Unable to resolve default edge deployment", err.Error())
+			state = flattenDataset(ds, "")
+		}
+
+		resp.Diagnostics.Append(resp.State.Set(ctx, state)...)
 	}
 }
 
@@ -306,7 +338,7 @@ func (r *DatasetResource) ImportState(ctx context.Context, req resource.ImportSt
 	resource.ImportStatePassthroughID(ctx, path.Root("id"), req, resp)
 }
 
-func flattenDataset(dataset *axiom.Dataset) DatasetResourceModel {
+func flattenDataset(dataset *axiom.Dataset, defaultEdgeDeployment string) DatasetResourceModel {
 	var description types.String
 	var edgeDeployment types.String
 
@@ -316,6 +348,8 @@ func flattenDataset(dataset *axiom.Dataset) DatasetResourceModel {
 
 	if dataset.EdgeDeployment != "" {
 		edgeDeployment = types.StringValue(dataset.EdgeDeployment)
+	} else if defaultEdgeDeployment != "" {
+		edgeDeployment = types.StringValue(defaultEdgeDeployment)
 	}
 
 	mapFields := make([]attr.Value, 0, len(dataset.MapFields))
@@ -353,3 +387,41 @@ func edgeDeploymentValue(edgeDeployment types.String) string {
 
 	return edgeDeployment.ValueString()
 }
+
+func flattenDatasetWithOrgDefault(ctx context.Context, client *axiom.Client, dataset *axiom.Dataset) (DatasetResourceModel, error) {
+	if dataset.EdgeDeployment != "" {
+		return flattenDataset(dataset, ""), nil
+	}
+
+	defaultEdgeDeployment, err := fetchDefaultEdgeDeployment(ctx, client)
+	if err != nil {
+		return DatasetResourceModel{}, err
+	}
+
+	return flattenDataset(dataset, defaultEdgeDeployment), nil
+}
+
+func fetchDefaultEdgeDeployment(ctx context.Context, client *axiom.Client) (string, error) {
+	if client == nil {
+		return "", fmt.Errorf("client is not set")
+	}
+
+	organizations, err := client.Organizations.List(ctx)
+	if err != nil {
+		return "", fmt.Errorf("could not fetch organizations while resolving edge deployment: %w", err)
+	}
+
+	return selectDefaultEdgeDeployment(organizations), nil
+}
+
+func selectDefaultEdgeDeployment(organizations []*axiom.Organization) string {
+	for _, organization := range organizations {
+		if organization == nil || organization.DefaultEdgeDeployment == "" {
+			continue
+		}
+
+		return organization.DefaultEdgeDeployment
+	}
+
+	return ""
+}

axiom/resource_dataset_unit_test.go

@@ -54,24 +54,73 @@ func TestFlattenDataset_EdgeDeployment(t *testing.T) {
 			Name:           "dataset-a",
 			Kind:           "axiom:events:v1",
 			EdgeDeployment: "cloud.eu-central-1.aws",
-		})
+		}, "")
 
 		assert.Equal(t, "cloud.eu-central-1.aws", state.EdgeDeployment.ValueString())
 	})
 
+	t.Run("uses org default edge deployment when dataset omits it", func(t *testing.T) {
+		t.Parallel()
+
+		state := flattenDataset(&axiom.Dataset{
+			ID:   "dataset-a",
+			Name: "dataset-a",
+			Kind: "axiom:events:v1",
+		}, "cloud.us-east-1.aws")
+
+		assert.Equal(t, "cloud.us-east-1.aws", state.EdgeDeployment.ValueString())
+	})
+
 	t.Run("keeps edge deployment null when absent", func(t *testing.T) {
 		t.Parallel()
 
 		state := flattenDataset(&axiom.Dataset{
 			ID:   "dataset-a",
 			Name: "dataset-a",
 			Kind: "axiom:events:v1",
-		})
+		}, "")
 
 		assert.True(t, state.EdgeDeployment.IsNull())
 	})
 }
 
+func TestSelectDefaultEdgeDeployment(t *testing.T) {
+	t.Parallel()
+
+	t.Run("returns first default edge deployment", func(t *testing.T) {
+		t.Parallel()
+
+		selected := selectDefaultEdgeDeployment([]*axiom.Organization{
+			{DefaultEdgeDeployment: "cloud.eu-central-1.aws"},
+			{DefaultEdgeDeployment: "cloud.us-east-1.aws"},
+		})
+
+		assert.Equal(t, "cloud.eu-central-1.aws", selected)
+	})
+
+	t.Run("skips nil and empty organizations", func(t *testing.T) {
+		t.Parallel()
+
+		selected := selectDefaultEdgeDeployment([]*axiom.Organization{
+			nil,
+			{},
+			{DefaultEdgeDeployment: "cloud.us-east-1.aws"},
+		})
+
+		assert.Equal(t, "cloud.us-east-1.aws", selected)
+	})
+
+	t.Run("returns empty when none configured", func(t *testing.T) {
+		t.Parallel()
+
+		selected := selectDefaultEdgeDeployment([]*axiom.Organization{
+			{},
+		})
+
+		assert.Empty(t, selected)
+	})
+}
+
 func TestEdgeDeploymentValue(t *testing.T) {
 	t.Parallel()