Merge pull request #90 from axiomhq/axm-11643-support-edge-based-dataset-creation-in-terraform

Support edge deployment for dataset resource

Author: ronoc

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

@@ -44,6 +44,7 @@ type DatasetResourceModel struct {
 	Name               types.String `tfsdk:"name"`
 	Kind               types.String `tfsdk:"kind"`
 	Description        types.String `tfsdk:"description"`
+	EdgeDeployment     types.String `tfsdk:"edge_deployment"`
 	ID                 types.String `tfsdk:"id"`
 	UseRetentionPeriod types.Bool   `tfsdk:"use_retention_period"`
 	RetentionDays      types.Int64  `tfsdk:"retention_days"`
@@ -90,6 +91,14 @@ func (r *DatasetResource) Schema(_ context.Context, _ resource.SchemaRequest, re
 					stringplanmodifier.UseStateForUnknown(),
 				},
 			},
+			"edge_deployment": schema.StringAttribute{
+				Optional:            true,
+				Computed:            true,
+				MarkdownDescription: "Edge deployment for the dataset (for example, 'cloud.eu-central-1.aws')",
+				PlanModifiers: []planmodifier.String{
+					stringplanmodifier.UseStateForUnknown(),
+				},
+			},
 			"id": schema.StringAttribute{
 				Computed:            true,
 				MarkdownDescription: "Dataset identifier",
@@ -170,20 +179,20 @@ func (r *DatasetResource) Create(ctx context.Context, req resource.CreateRequest
 		return
 	}
 
-	ds, err := r.client.Datasets.Create(ctx, axiom.DatasetCreateRequest{
-		Name:               plan.Name.ValueString(),
-		Kind:               plan.Kind.ValueString(),
-		Description:        plan.Description.ValueString(),
-		UseRetentionPeriod: plan.UseRetentionPeriod.ValueBool(),
-		RetentionDays:      int(plan.RetentionDays.ValueInt64()),
-	})
+	ds, err := r.client.Datasets.Create(ctx, datasetCreateRequestFromPlan(plan))
 	if err != nil {
 		resp.Diagnostics.AddError("Client Error", fmt.Sprintf("Unable to create dataset, got error: %s", err))
 		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
 	}
@@ -203,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)...)
 	}
 }
 
@@ -229,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) {
@@ -260,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
 	}
@@ -281,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)...)
 	}
 }
 
@@ -303,13 +338,20 @@ 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
 
 	if dataset.Description != "" {
 		description = types.StringValue(dataset.Description)
 	}
 
+	if dataset.EdgeDeployment != "" {
+		edgeDeployment = types.StringValue(dataset.EdgeDeployment)
+	} else if defaultEdgeDeployment != "" {
+		edgeDeployment = types.StringValue(defaultEdgeDeployment)
+	}
+
 	mapFields := make([]attr.Value, 0, len(dataset.MapFields))
 	for _, fieldName := range dataset.MapFields {
 		mapFields = append(mapFields, types.StringValue(fieldName))
@@ -320,8 +362,66 @@ func flattenDataset(dataset *axiom.Dataset) DatasetResourceModel {
 		Name:               types.StringValue(dataset.Name),
 		Kind:               types.StringValue(dataset.Kind),
 		Description:        description,
+		EdgeDeployment:     edgeDeployment,
 		UseRetentionPeriod: types.BoolValue(dataset.UseRetentionPeriod),
 		RetentionDays:      types.Int64Value(int64(dataset.RetentionDays)),
 		MapFields:          types.ListValueMust(types.StringType, mapFields),
 	}
 }
+
+func datasetCreateRequestFromPlan(plan DatasetResourceModel) axiom.DatasetCreateRequest {
+	return axiom.DatasetCreateRequest{
+		Name:               plan.Name.ValueString(),
+		Kind:               plan.Kind.ValueString(),
+		Description:        plan.Description.ValueString(),
+		EdgeDeployment:     edgeDeploymentValue(plan.EdgeDeployment),
+		UseRetentionPeriod: plan.UseRetentionPeriod.ValueBool(),
+		RetentionDays:      int(plan.RetentionDays.ValueInt64()),
+	}
+}
+
+func edgeDeploymentValue(edgeDeployment types.String) string {
+	if edgeDeployment.IsNull() || edgeDeployment.IsUnknown() {
+		return ""
+	}
+
+	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

@@ -0,0 +1,144 @@
+package axiom
+
+import (
+	"testing"
+
+	"github.com/hashicorp/terraform-plugin-framework/types"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/axiomhq/axiom-go/axiom"
+)
+
+func TestDatasetCreateRequestFromPlan(t *testing.T) {
+	t.Parallel()
+
+	request := datasetCreateRequestFromPlan(DatasetResourceModel{
+		Name:               types.StringValue("dataset-a"),
+		Kind:               types.StringValue("axiom:events:v1"),
+		Description:        types.StringValue("dataset description"),
+		EdgeDeployment:     types.StringValue("cloud.eu-central-1.aws"),
+		UseRetentionPeriod: types.BoolValue(true),
+		RetentionDays:      types.Int64Value(30),
+	})
+
+	assert.Equal(t, axiom.DatasetCreateRequest{
+		Name:               "dataset-a",
+		Kind:               "axiom:events:v1",
+		Description:        "dataset description",
+		EdgeDeployment:     "cloud.eu-central-1.aws",
+		UseRetentionPeriod: true,
+		RetentionDays:      30,
+	}, request)
+}
+
+func TestDatasetCreateRequestFromPlan_NullEdgeDeployment(t *testing.T) {
+	t.Parallel()
+
+	request := datasetCreateRequestFromPlan(DatasetResourceModel{
+		Name:           types.StringValue("dataset-a"),
+		Kind:           types.StringValue("axiom:events:v1"),
+		EdgeDeployment: types.StringNull(),
+	})
+
+	assert.Empty(t, request.EdgeDeployment)
+}
+
+func TestFlattenDataset_EdgeDeployment(t *testing.T) {
+	t.Parallel()
+
+	t.Run("sets edge deployment in state", func(t *testing.T) {
+		t.Parallel()
+
+		state := flattenDataset(&axiom.Dataset{
+			ID:             "dataset-a",
+			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()
+
+	t.Run("returns empty string for null", func(t *testing.T) {
+		t.Parallel()
+
+		assert.Empty(t, edgeDeploymentValue(types.StringNull()))
+	})
+
+	t.Run("returns empty string for unknown", func(t *testing.T) {
+		t.Parallel()
+
+		assert.Empty(t, edgeDeploymentValue(types.StringUnknown()))
+	})
+
+	t.Run("returns value when known", func(t *testing.T) {
+		t.Parallel()
+
+		assert.Equal(t, "cloud.eu-central-1.aws", edgeDeploymentValue(types.StringValue("cloud.eu-central-1.aws")))
+	})
+}

docs/data-sources/dataset.md

@@ -22,6 +22,7 @@ description: |-
 ### Read-Only
 
 - `description` (String) Dataset description
+- `edge_deployment` (String) Edge deployment for the dataset (for example, 'cloud.eu-central-1.aws')
 - `kind` (String) Dataset kind. Must be one of: 'axiom:events:v1', 'otel:metrics:v1', 'otel:traces:v1', 'otel:logs:v1'. Defaults to 'axiom:events:v1'
 - `map_fields` (List of String) Map fields for the dataset
 - `name` (String) Dataset name

docs/resources/dataset.md

@@ -22,6 +22,7 @@ description: |-
 ### Optional
 
 - `description` (String) Dataset description
+- `edge_deployment` (String) Edge deployment for the dataset (for example, 'cloud.eu-central-1.aws')
 - `kind` (String) Dataset kind. Must be one of: 'axiom:events:v1', 'otel:metrics:v1', 'otel:traces:v1', 'otel:logs:v1'. Defaults to 'axiom:events:v1'
 - `map_fields` (List of String) Map fields for the dataset
 - `retention_days` (Number) Retention days for the dataset