diff --git a/.github/docs/contribution-guide/resource.go b/.github/docs/contribution-guide/resource.go
index d12f50a2..3d604d11 100644
--- a/.github/docs/contribution-guide/resource.go
+++ b/.github/docs/contribution-guide/resource.go
@@ -211,7 +211,7 @@ func (r *barResource) Create(ctx context.Context, req resource.CreateRequest, re
 	// only in case the create bar API call is asynchronous (Make sure to include *ALL* fields which are part of the
 	// internal terraform resource id! And please include the comment below in your code):
 	// Write id attributes to state before polling via the wait handler - just in case anything goes wrong during the wait handler
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
 		"project_id": projectId,
 		"region":     region,
 		"bar_id":     resp.BarId,
@@ -335,7 +335,7 @@ func (r *barResource) ImportState(ctx context.Context, req resource.ImportStateR
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": idParts[0],
 		"region":     idParts[1],
 		"bar_id":     idParts[2],
diff --git a/stackit/internal/services/dns/recordset/resource.go b/stackit/internal/services/dns/recordset/resource.go
index 2077e09f..badf20bf 100644
--- a/stackit/internal/services/dns/recordset/resource.go
+++ b/stackit/internal/services/dns/recordset/resource.go
@@ -223,7 +223,7 @@ func (r *recordSetResource) Create(ctx context.Context, req resource.CreateReque
 	ctx = core.LogResponse(ctx)
 
 	// Write id attributes to state before polling via the wait handler - just in case anything goes wrong during the wait handler
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":    projectId,
 		"zone_id":       zoneId,
 		"record_set_id": *recordSetResp.Rrset.Id,
@@ -399,7 +399,7 @@ func (r *recordSetResource) ImportState(ctx context.Context, req resource.Import
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
 		"project_id":    idParts[0],
 		"zone_id":       idParts[1],
 		"record_set_id": idParts[2],
diff --git a/stackit/internal/services/dns/zone/resource.go b/stackit/internal/services/dns/zone/resource.go
index 9fd92436..b7c76e11 100644
--- a/stackit/internal/services/dns/zone/resource.go
+++ b/stackit/internal/services/dns/zone/resource.go
@@ -306,7 +306,7 @@ func (r *zoneResource) Create(ctx context.Context, req resource.CreateRequest, r
 
 	// Write id attributes to state before polling via the wait handler - just in case anything goes wrong during the wait handler
 	zoneId := *createResp.Zone.Id
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
 		"project_id": projectId,
 		"zone_id":    zoneId,
 	})
@@ -477,7 +477,7 @@ func (r *zoneResource) ImportState(ctx context.Context, req resource.ImportState
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
 		"project_id": idParts[0],
 		"zone_id":    idParts[1],
 	})
diff --git a/stackit/internal/services/iaas/affinitygroup/resource.go b/stackit/internal/services/iaas/affinitygroup/resource.go
index 6597ff65..726ac8ee 100644
--- a/stackit/internal/services/iaas/affinitygroup/resource.go
+++ b/stackit/internal/services/iaas/affinitygroup/resource.go
@@ -325,7 +325,7 @@ func (r *affinityGroupResource) ImportState(ctx context.Context, req resource.Im
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":        idParts[0],
 		"region":            idParts[1],
 		"affinity_group_id": idParts[2],
diff --git a/stackit/internal/services/iaas/image/resource.go b/stackit/internal/services/iaas/image/resource.go
index 0699093f..7581bd38 100644
--- a/stackit/internal/services/iaas/image/resource.go
+++ b/stackit/internal/services/iaas/image/resource.go
@@ -650,7 +650,7 @@ func (r *imageResource) ImportState(ctx context.Context, req resource.ImportStat
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": idParts[0],
 		"region":     idParts[1],
 		"image_id":   idParts[2],
diff --git a/stackit/internal/services/iaas/networkarearegion/resource.go b/stackit/internal/services/iaas/networkarearegion/resource.go
index 36dd3a1a..6eaae25c 100644
--- a/stackit/internal/services/iaas/networkarearegion/resource.go
+++ b/stackit/internal/services/iaas/networkarearegion/resource.go
@@ -289,7 +289,7 @@ func (r *networkAreaRegionResource) Create(ctx context.Context, req resource.Cre
 	ctx = core.LogResponse(ctx)
 
 	// Write id attributes to state before polling via the wait handler - just in case anything goes wrong during the wait handler
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"organization_id": organizationId,
 		"network_area_id": networkAreaId,
 		"region":          region,
@@ -482,7 +482,7 @@ func (r *networkAreaRegionResource) ImportState(ctx context.Context, req resourc
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"organization_id": idParts[0],
 		"network_area_id": idParts[1],
 		"region":          idParts[2],
diff --git a/stackit/internal/services/iaas/networkarearoute/resource.go b/stackit/internal/services/iaas/networkarearoute/resource.go
index f5ba3bd9..3694312e 100644
--- a/stackit/internal/services/iaas/networkarearoute/resource.go
+++ b/stackit/internal/services/iaas/networkarearoute/resource.go
@@ -547,7 +547,7 @@ func (r *networkAreaRouteResource) ImportState(ctx context.Context, req resource
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"organization_id":       idParts[0],
 		"network_area_id":       idParts[1],
 		"region":                idParts[2],
diff --git a/stackit/internal/services/iaas/networkinterface/resource.go b/stackit/internal/services/iaas/networkinterface/resource.go
index 8ced0477..007e17f0 100644
--- a/stackit/internal/services/iaas/networkinterface/resource.go
+++ b/stackit/internal/services/iaas/networkinterface/resource.go
@@ -476,7 +476,7 @@ func (r *networkInterfaceResource) ImportState(ctx context.Context, req resource
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":           idParts[0],
 		"region":               idParts[1],
 		"network_id":           idParts[2],
diff --git a/stackit/internal/services/iaas/networkinterfaceattach/resource.go b/stackit/internal/services/iaas/networkinterfaceattach/resource.go
index 2b8d4240..7c7f8c04 100644
--- a/stackit/internal/services/iaas/networkinterfaceattach/resource.go
+++ b/stackit/internal/services/iaas/networkinterfaceattach/resource.go
@@ -317,7 +317,7 @@ func (r *networkInterfaceAttachResource) ImportState(ctx context.Context, req re
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
 		"project_id":           idParts[0],
 		"region":               idParts[1],
 		"server_id":            idParts[2],
diff --git a/stackit/internal/services/iaas/publicip/resource.go b/stackit/internal/services/iaas/publicip/resource.go
index aa8ac637..0d69eab1 100644
--- a/stackit/internal/services/iaas/publicip/resource.go
+++ b/stackit/internal/services/iaas/publicip/resource.go
@@ -374,7 +374,7 @@ func (r *publicIpResource) ImportState(ctx context.Context, req resource.ImportS
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":   idParts[0],
 		"region":       idParts[1],
 		"public_ip_id": idParts[2],
diff --git a/stackit/internal/services/iaas/publicipassociate/resource.go b/stackit/internal/services/iaas/publicipassociate/resource.go
index 66028381..6d13bc1f 100644
--- a/stackit/internal/services/iaas/publicipassociate/resource.go
+++ b/stackit/internal/services/iaas/publicipassociate/resource.go
@@ -335,7 +335,7 @@ func (r *publicIpAssociateResource) ImportState(ctx context.Context, req resourc
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":           idParts[0],
 		"region":               idParts[1],
 		"public_ip_id":         idParts[2],
diff --git a/stackit/internal/services/iaas/securitygroup/resource.go b/stackit/internal/services/iaas/securitygroup/resource.go
index 07b510ac..07554fe3 100644
--- a/stackit/internal/services/iaas/securitygroup/resource.go
+++ b/stackit/internal/services/iaas/securitygroup/resource.go
@@ -393,7 +393,7 @@ func (r *securityGroupResource) ImportState(ctx context.Context, req resource.Im
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":        idParts[0],
 		"region":            idParts[1],
 		"security_group_id": idParts[2],
diff --git a/stackit/internal/services/iaas/securitygrouprule/resource.go b/stackit/internal/services/iaas/securitygrouprule/resource.go
index ab075b47..866580e7 100644
--- a/stackit/internal/services/iaas/securitygrouprule/resource.go
+++ b/stackit/internal/services/iaas/securitygrouprule/resource.go
@@ -606,7 +606,7 @@ func (r *securityGroupRuleResource) ImportState(ctx context.Context, req resourc
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":             idParts[0],
 		"region":                 idParts[1],
 		"security_group_id":      idParts[2],
diff --git a/stackit/internal/services/iaas/server/resource.go b/stackit/internal/services/iaas/server/resource.go
index 5af72de8..cd75a552 100644
--- a/stackit/internal/services/iaas/server/resource.go
+++ b/stackit/internal/services/iaas/server/resource.go
@@ -871,7 +871,7 @@ func (r *serverResource) ImportState(ctx context.Context, req resource.ImportSta
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": idParts[0],
 		"region":     idParts[1],
 		"server_id":  idParts[2],
diff --git a/stackit/internal/services/iaas/serviceaccountattach/resource.go b/stackit/internal/services/iaas/serviceaccountattach/resource.go
index 2063f15c..bf6d31e9 100644
--- a/stackit/internal/services/iaas/serviceaccountattach/resource.go
+++ b/stackit/internal/services/iaas/serviceaccountattach/resource.go
@@ -312,7 +312,7 @@ func (r *serviceAccountAttachResource) ImportState(ctx context.Context, req reso
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":            idParts[0],
 		"region":                idParts[1],
 		"server_id":             idParts[2],
diff --git a/stackit/internal/services/iaas/volume/resource.go b/stackit/internal/services/iaas/volume/resource.go
index 0fc3a9e6..1ce0d5d7 100644
--- a/stackit/internal/services/iaas/volume/resource.go
+++ b/stackit/internal/services/iaas/volume/resource.go
@@ -558,7 +558,7 @@ func (r *volumeResource) ImportState(ctx context.Context, req resource.ImportSta
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": idParts[0],
 		"region":     idParts[1],
 		"volume_id":  idParts[2],
diff --git a/stackit/internal/services/iaas/volumeattach/resource.go b/stackit/internal/services/iaas/volumeattach/resource.go
index f297f16d..d4af203e 100644
--- a/stackit/internal/services/iaas/volumeattach/resource.go
+++ b/stackit/internal/services/iaas/volumeattach/resource.go
@@ -314,7 +314,7 @@ func (r *volumeAttachResource) ImportState(ctx context.Context, req resource.Imp
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": idParts[0],
 		"region":     idParts[1],
 		"server_id":  idParts[2],
diff --git a/stackit/internal/services/kms/key/resource.go b/stackit/internal/services/kms/key/resource.go
index 9f56669b..41707dc7 100644
--- a/stackit/internal/services/kms/key/resource.go
+++ b/stackit/internal/services/kms/key/resource.go
@@ -274,7 +274,7 @@ func (r *keyResource) Create(ctx context.Context, req resource.CreateRequest, re
 
 	keyId := *createResponse.Id
 	// Write id attributes to state before polling via the wait handler - just in case anything goes wrong during the wait handler
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": projectId,
 		"region":     region,
 		"keyring_id": keyRingId,
@@ -392,7 +392,7 @@ func (r *keyResource) ImportState(ctx context.Context, req resource.ImportStateR
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": idParts[0],
 		"region":     idParts[1],
 		"keyring_id": idParts[2],
diff --git a/stackit/internal/services/kms/keyring/resource.go b/stackit/internal/services/kms/keyring/resource.go
index 3627f3f9..6b2f8e9a 100644
--- a/stackit/internal/services/kms/keyring/resource.go
+++ b/stackit/internal/services/kms/keyring/resource.go
@@ -206,7 +206,7 @@ func (r *keyRingResource) Create(ctx context.Context, req resource.CreateRequest
 
 	keyRingId := *createResponse.Id
 	// Write id attributes to state before polling via the wait handler - just in case anything goes wrong during the wait handler
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": projectId,
 		"region":     region,
 		"keyring_id": keyRingId,
@@ -308,7 +308,7 @@ func (r *keyRingResource) ImportState(ctx context.Context, req resource.ImportSt
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id": idParts[0],
 		"region":     idParts[1],
 		"keyring_id": idParts[2],
diff --git a/stackit/internal/services/kms/wrapping-key/resource.go b/stackit/internal/services/kms/wrapping-key/resource.go
index d534b115..3d2cfc62 100644
--- a/stackit/internal/services/kms/wrapping-key/resource.go
+++ b/stackit/internal/services/kms/wrapping-key/resource.go
@@ -280,7 +280,7 @@ func (r *wrappingKeyResource) Create(ctx context.Context, req resource.CreateReq
 	wrappingKeyId := *createWrappingKeyResp.Id
 
 	// Write id attributes to state before polling via the wait handler - just in case anything goes wrong during the wait handler
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]interface{}{
 		"project_id":      projectId,
 		"region":          region,
 		"keyring_id":      keyRingId,
@@ -395,7 +395,7 @@ func (r *wrappingKeyResource) ImportState(ctx context.Context, req resource.Impo
 		return
 	}
 
-	utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
+	ctx = utils.SetAndLogStateFields(ctx, &resp.Diagnostics, &resp.State, map[string]any{
 		"project_id":      idParts[0],
 		"region":          idParts[1],
 		"keyring_id":      idParts[2],
diff --git a/stackit/internal/utils/utils.go b/stackit/internal/utils/utils.go
index a7614113..b7a479af 100644
--- a/stackit/internal/utils/utils.go
+++ b/stackit/internal/utils/utils.go
@@ -176,9 +176,10 @@ func CheckListRemoval(ctx context.Context, configModelList, planModelList types.
 }
 
 // SetAndLogStateFields writes the given map of key-value pairs to the state
-func SetAndLogStateFields(ctx context.Context, diags *diag.Diagnostics, state *tfsdk.State, values map[string]any) {
+func SetAndLogStateFields(ctx context.Context, diags *diag.Diagnostics, state *tfsdk.State, values map[string]any) context.Context {
 	for key, val := range values {
 		ctx = tflog.SetField(ctx, key, val)
 		diags.Append(state.SetAttribute(ctx, path.Root(key), val)...)
 	}
+	return ctx
 }