From bdd9225041d1dbbc834f7df6adc7463bb3c09b9f Mon Sep 17 00:00:00 2001
From: Pavel Boldyrev <627562+bpg@users.noreply.github.com>
Date: Sun, 2 Jun 2024 20:51:44 -0400
Subject: [PATCH] fix(access): fix `user` and `role` import (#1346)

* fix(access): fix `user` and `role` import

Signed-off-by: Pavel Boldyrev <627562+bpg@users.noreply.github.com>
---
 proxmoxtf/resource/role.go | 19 ++++++++++++++++++-
 proxmoxtf/resource/user.go | 12 +++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/proxmoxtf/resource/role.go b/proxmoxtf/resource/role.go
index 275c2af3..3b47ead4 100644
--- a/proxmoxtf/resource/role.go
+++ b/proxmoxtf/resource/role.go
@@ -9,6 +9,7 @@ package resource
 import (
 	"context"
 	"errors"
+	"fmt"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@@ -46,13 +47,23 @@ func Role() *schema.Resource {
 		UpdateContext: roleUpdate,
 		DeleteContext: roleDelete,
 		Importer: &schema.ResourceImporter{
-			StateContext: schema.ImportStatePassthroughContext,
+			StateContext: func(_ context.Context, d *schema.ResourceData, _ interface{}) ([]*schema.ResourceData, error) {
+				roleID := d.Id()
+
+				err := d.Set(mkResourceVirtualEnvironmentRoleRoleID, roleID)
+				if err != nil {
+					return nil, fmt.Errorf("failed setting state during import: %w", err)
+				}
+
+				return []*schema.ResourceData{d}, nil
+			},
 		},
 	}
 }
 
 func roleCreate(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics {
 	config := m.(proxmoxtf.ProviderConfiguration)
+
 	client, err := config.GetClient()
 	if err != nil {
 		return diag.FromErr(err)
@@ -83,12 +94,14 @@ func roleCreate(ctx context.Context, d *schema.ResourceData, m interface{}) diag
 
 func roleRead(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics {
 	config := m.(proxmoxtf.ProviderConfiguration)
+
 	client, err := config.GetClient()
 	if err != nil {
 		return diag.FromErr(err)
 	}
 
 	roleID := d.Id()
+
 	role, err := client.Access().GetRole(ctx, roleID)
 	if err != nil {
 		if errors.Is(err, api.ErrResourceDoesNotExist) {
@@ -96,6 +109,7 @@ func roleRead(ctx context.Context, d *schema.ResourceData, m interface{}) diag.D
 
 			return nil
 		}
+
 		return diag.FromErr(err)
 	}
 
@@ -108,11 +122,13 @@ func roleRead(ctx context.Context, d *schema.ResourceData, m interface{}) diag.D
 	}
 
 	err = d.Set(mkResourceVirtualEnvironmentRolePrivileges, privileges)
+
 	return diag.FromErr(err)
 }
 
 func roleUpdate(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics {
 	config := m.(proxmoxtf.ProviderConfiguration)
+
 	client, err := config.GetClient()
 	if err != nil {
 		return diag.FromErr(err)
@@ -140,6 +156,7 @@ func roleUpdate(ctx context.Context, d *schema.ResourceData, m interface{}) diag
 
 func roleDelete(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics {
 	config := m.(proxmoxtf.ProviderConfiguration)
+
 	client, err := config.GetClient()
 	if err != nil {
 		return diag.FromErr(err)
diff --git a/proxmoxtf/resource/user.go b/proxmoxtf/resource/user.go
index f0de113a..4315d7c9 100644
--- a/proxmoxtf/resource/user.go
+++ b/proxmoxtf/resource/user.go
@@ -9,6 +9,7 @@ package resource
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
@@ -147,7 +148,16 @@ func User() *schema.Resource {
 		UpdateContext: userUpdate,
 		DeleteContext: userDelete,
 		Importer: &schema.ResourceImporter{
-			StateContext: schema.ImportStatePassthroughContext,
+			StateContext: func(_ context.Context, d *schema.ResourceData, _ interface{}) ([]*schema.ResourceData, error) {
+				roleID := d.Id()
+
+				err := d.Set(mkResourceVirtualEnvironmentUserUserID, roleID)
+				if err != nil {
+					return nil, fmt.Errorf("failed setting state during import: %w", err)
+				}
+
+				return []*schema.ResourceData{d}, nil
+			},
 		},
 	}
 }