From 3515bec30b8c7bd66b1db431623bc27f18bd0943 Mon Sep 17 00:00:00 2001
From: jo <ljonas@riseup.net>
Date: Tue, 12 Dec 2023 17:20:54 +0100
Subject: [PATCH] wip

---
 internal/sshkey/data_source.go | 52 +++++++++++++++++++---------------
 1 file changed, 29 insertions(+), 23 deletions(-)

diff --git a/internal/sshkey/data_source.go b/internal/sshkey/data_source.go
index 4e846db3d..8710ddc85 100644
--- a/internal/sshkey/data_source.go
+++ b/internal/sshkey/data_source.go
@@ -34,14 +34,19 @@ type resourceData struct {
 	Fingerprint types.String `tfsdk:"fingerprint"`
 	PublicKey   types.String `tfsdk:"public_key"`
 	Labels      types.Map    `tfsdk:"labels"`
+
+	Selector     types.String `tfsdk:"selector"`
+	WithSelector types.String `tfsdk:"with_selector"`
 }
 
 var resourceDataAttrTypes = map[string]attr.Type{
-	"id":          types.Int64Type,
-	"name":        types.StringType,
-	"fingerprint": types.StringType,
-	"public_key":  types.StringType,
-	"labels":      types.MapType{ElemType: types.StringType},
+	"id":            types.Int64Type,
+	"name":          types.StringType,
+	"fingerprint":   types.StringType,
+	"public_key":    types.StringType,
+	"labels":        types.MapType{ElemType: types.StringType},
+	"selector":      types.StringType,
+	"with_selector": types.StringType,
 }
 
 func newResourceData(ctx context.Context, in *hcloud.SSHKey) (resourceData, diag.Diagnostics) {
@@ -80,6 +85,7 @@ func getCommonDataSchema() map[string]schema.Attribute {
 		},
 		"labels": schema.MapAttribute{
 			ElementType: types.StringType,
+			Optional:    true,
 			Computed:    true,
 		},
 	}
@@ -140,10 +146,9 @@ func (d *dataSource) ConfigValidators(_ context.Context) []datasource.ConfigVali
 		datasourcevalidator.ExactlyOneOf(
 			path.MatchRoot("id"),
 			path.MatchRoot("name"),
-		),
-		datasourcevalidator.ExactlyOneOf(
-			path.MatchRoot("with_selector"),
+			path.MatchRoot("fingerprint"),
 			path.MatchRoot("selector"),
+			path.MatchRoot("with_selector"),
 		),
 	}
 }
@@ -185,6 +190,7 @@ func (d *dataSource) Read(ctx context.Context, req datasource.ReadRequest, resp
 		result, _, err = d.client.SSHKey.GetByName(ctx, data.Name.ValueString())
 		if err != nil {
 			resp.Diagnostics.Append(hcclient.APIErrorDiagnostics(err)...)
+			return
 		}
 		if result == nil {
 			resp.Diagnostics.AddError(
@@ -197,6 +203,7 @@ func (d *dataSource) Read(ctx context.Context, req datasource.ReadRequest, resp
 		result, _, err = d.client.SSHKey.GetByFingerprint(ctx, data.Fingerprint.ValueString())
 		if err != nil {
 			resp.Diagnostics.Append(hcclient.APIErrorDiagnostics(err)...)
+			return
 		}
 		if result == nil {
 			resp.Diagnostics.AddError(
@@ -205,15 +212,16 @@ func (d *dataSource) Read(ctx context.Context, req datasource.ReadRequest, resp
 			)
 			return
 		}
-	case !selectorData.WithSelector.IsNull() || !selectorData.Selector.IsNull():
-		var labelSelector string
-		if !selectorData.WithSelector.IsNull() {
-			labelSelector = selectorData.WithSelector.ValueString()
-		} else if !selectorData.Selector.IsNull() {
-			labelSelector = selectorData.Selector.ValueString()
+	case !data.WithSelector.IsNull() || !data.Selector.IsNull():
+		opts := hcloud.SSHKeyListOpts{}
+
+		if !data.WithSelector.IsNull() {
+			opts.LabelSelector = data.WithSelector.ValueString()
+		} else if !data.Selector.IsNull() {
+			opts.LabelSelector = data.Selector.ValueString()
 		}
 
-		allKeys, err := d.client.SSHKey.AllWithOpts(ctx, hcloud.SSHKeyListOpts{ListOpts: hcloud.ListOpts{LabelSelector: labelSelector}})
+		allKeys, err := d.client.SSHKey.AllWithOpts(ctx, opts)
 		if err != nil {
 			resp.Diagnostics.Append(hcclient.APIErrorDiagnostics(err)...)
 		}
@@ -223,7 +231,7 @@ func (d *dataSource) Read(ctx context.Context, req datasource.ReadRequest, resp
 				fmt.Sprintf(
 					"No ssh key found for label selector.\n\n"+
 						"Label selector: %s\n",
-					labelSelector,
+					opts.LabelSelector,
 				),
 			)
 			return
@@ -234,7 +242,7 @@ func (d *dataSource) Read(ctx context.Context, req datasource.ReadRequest, resp
 				fmt.Sprintf(
 					"More than one ssh key found for label selector.\n\n"+
 						"Label selector: %s\n",
-					labelSelector,
+					opts.LabelSelector,
 				),
 			)
 			return
@@ -303,6 +311,8 @@ func (d *dataSourceList) Schema(_ context.Context, _ datasource.SchemaRequest, r
 type resourceDataList struct {
 	ID      types.String `tfsdk:"id"`
 	SSHKeys types.List   `tfsdk:"ssh_keys"`
+
+	WithSelector types.String `tfsdk:"with_selector"`
 }
 
 func newResourceDataList(ctx context.Context, in []*hcloud.SSHKey) (resourceDataList, diag.Diagnostics) {
@@ -336,12 +346,8 @@ func newResourceDataList(ctx context.Context, in []*hcloud.SSHKey) (resourceData
 // ReadRequest and new state values set on the ReadResponse.
 func (d *dataSourceList) Read(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse) {
 	var data resourceDataList
-	var selectorData struct {
-		WithSelector types.String `tfsdk:"with_selector"`
-	}
 
 	resp.Diagnostics.Append(req.Config.Get(ctx, &data)...)
-	resp.Diagnostics.Append(req.Config.Get(ctx, &selectorData)...)
 	if resp.Diagnostics.HasError() {
 		return
 	}
@@ -350,8 +356,8 @@ func (d *dataSourceList) Read(ctx context.Context, req datasource.ReadRequest, r
 	var err error
 
 	opts := hcloud.SSHKeyListOpts{}
-	if !selectorData.WithSelector.IsNull() {
-		opts.LabelSelector = selectorData.WithSelector.ValueString()
+	if !data.WithSelector.IsNull() {
+		opts.LabelSelector = data.WithSelector.ValueString()
 	}
 
 	result, err = d.client.SSHKey.AllWithOpts(ctx, opts)