r/aws_sagemaker_image_version: read correct version after create

Previously the read operation would always retrieve the latest image version matching the `image_name` argument, rather than tracking the version created by a given resource instance. This change updates the `id` attribute to a comma-delimited string concatenating `image_name` and `version`, enabling subsequent read operations to correctly match the version created by this resource, and supporting import of existing versions which are not the most recently created. It also fixes issues observed when managing multiple image versions of the same name in a single workspace.

```console
% SAGEMAKER_IMAGE_VERSION_BASE_IMAGE="<redacted>" make testacc PKG=sagemaker TESTS=TestAccSageMakerImageVersion_
make: Verifying source code with gofmt...
==> Checking that code complies with gofmt requirements...
TF_ACC=1 go1.23.8 test ./internal/service/sagemaker/... -v -count 1 -parallel 20 -run='TestAccSageMakerImageVersion_'  -timeout 360m -vet=off
2025/05/08 10:39:43 Initializing Terraform AWS Provider...

--- PASS: TestAccSageMakerImageVersion_Disappears_image (80.09s)
--- PASS: TestAccSageMakerImageVersion_basic (82.82s)
--- PASS: TestAccSageMakerImageVersion_disappears (84.49s)
--- PASS: TestAccSageMakerImageVersion_multiple (88.17s)
--- PASS: TestAccSageMakerImageVersion_update (91.47s)
--- PASS: TestAccSageMakerImageVersion_upgrade_V5_98_0 (102.47s)
PASS
ok      github.com/hashicorp/terraform-provider-aws/internal/service/sagemaker  108.081s
```

Author: jar-b

.changelog/42536.txt

@@ -0,0 +1,6 @@
+```release-note:bug
+resource/aws_sagemaker_image_version: Read the correct image version after creation rather than always fetching the latest
+```
+```release-note:breaking-change
+resource/aws_sagemaker_image_version: `id` is now a comma-delimited string concatenating `image_name` and `version`
+```

internal/service/sagemaker/exports_test.go

@@ -49,7 +49,7 @@ var (
 	FindHubByName                             = findHubByName
 	FindHumanTaskUIByName                     = findHumanTaskUIByName
 	FindImageByName                           = findImageByName
-	FindImageVersionByName                    = findImageVersionByName
+	FindImageVersionByTwoPartKey              = findImageVersionByTwoPartKey
 	FindMlflowTrackingServerByName            = findMlflowTrackingServerByName
 	FindModelByName                           = findModelByName
 	FindModelPackageGroupByName               = findModelPackageGroupByName

internal/service/sagemaker/image_version.go

@@ -6,6 +6,7 @@ package sagemaker
 import (
 	"context"
 	"log"
+	"strconv"
 
 	"github.com/YakDriver/regexache"
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -20,21 +21,34 @@ import (
 	"github.com/hashicorp/terraform-provider-aws/internal/enum"
 	"github.com/hashicorp/terraform-provider-aws/internal/errs"
 	"github.com/hashicorp/terraform-provider-aws/internal/errs/sdkdiag"
+	"github.com/hashicorp/terraform-provider-aws/internal/flex"
 	"github.com/hashicorp/terraform-provider-aws/internal/tfresource"
 	"github.com/hashicorp/terraform-provider-aws/names"
 )
 
+const imageVersionResourcePartCount = 2
+
 // @SDKResource("aws_sagemaker_image_version", name="Image Version")
 func resourceImageVersion() *schema.Resource {
 	return &schema.Resource{
 		CreateWithoutTimeout: resourceImageVersionCreate,
 		ReadWithoutTimeout:   resourceImageVersionRead,
 		UpdateWithoutTimeout: resourceImageVersionUpdate,
 		DeleteWithoutTimeout: resourceImageVersionDelete,
+
 		Importer: &schema.ResourceImporter{
 			StateContext: schema.ImportStatePassthroughContext,
 		},
 
+		StateUpgraders: []schema.StateUpgrader{
+			{
+				Type:    resourceImageVersionV0().CoreConfigSchema().ImpliedType(),
+				Upgrade: imageVersionStateUpgradeV0,
+				Version: 0,
+			},
+		},
+
+		SchemaVersion: 1,
 		Schema: map[string]*schema.Schema{
 			names.AttrARN: {
 				Type:     schema.TypeString,
@@ -105,7 +119,7 @@ func resourceImageVersionCreate(ctx context.Context, d *schema.ResourceData, met
 	conn := meta.(*conns.AWSClient).SageMakerClient(ctx)
 
 	name := d.Get("image_name").(string)
-	input := &sagemaker.CreateImageVersionInput{
+	input := sagemaker.CreateImageVersionInput{
 		ImageName:   aws.String(name),
 		BaseImage:   aws.String(d.Get("base_image").(string)),
 		ClientToken: aws.String(id.UniqueId()),
@@ -139,25 +153,37 @@ func resourceImageVersionCreate(ctx context.Context, d *schema.ResourceData, met
 		input.ProgrammingLang = aws.String(v.(string))
 	}
 
-	_, err := conn.CreateImageVersion(ctx, input)
-	if err != nil {
+	if _, err := conn.CreateImageVersion(ctx, &input); err != nil {
 		return sdkdiag.AppendErrorf(diags, "creating SageMaker AI Image Version %s: %s", name, err)
 	}
 
-	d.SetId(name)
-
-	if _, err := waitImageVersionCreated(ctx, conn, d.Id()); err != nil {
+	out, err := waitImageVersionCreated(ctx, conn, name)
+	if err != nil {
 		return sdkdiag.AppendErrorf(diags, "waiting for SageMaker AI Image Version (%s) to be created: %s", d.Id(), err)
 	}
 
+	parts := []string{name, strconv.Itoa(int(aws.ToInt32(out.Version)))}
+	id, err := flex.FlattenResourceId(parts, imageVersionResourcePartCount, false)
+	if err != nil {
+		return sdkdiag.AppendErrorf(diags, "creating SageMaker AI Image Version %s: %s", name, err)
+	}
+
+	d.SetId(id)
+
 	return append(diags, resourceImageVersionRead(ctx, d, meta)...)
 }
 
 func resourceImageVersionRead(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	var diags diag.Diagnostics
 	conn := meta.(*conns.AWSClient).SageMakerClient(ctx)
 
-	image, err := findImageVersionByName(ctx, conn, d.Id())
+	name, version, err := expandImageVersionResourceID(d.Id())
+	if err != nil {
+		return sdkdiag.AppendErrorf(diags, "reading SageMaker AI Image Version (%s): %s", d.Id(), err)
+	}
+	d.Set("image_name", name) // to support import
+
+	image, err := findImageVersionByTwoPartKey(ctx, conn, name, version)
 
 	if !d.IsNewResource() && tfresource.NotFound(err) {
 		d.SetId("")
@@ -174,7 +200,6 @@ func resourceImageVersionRead(ctx context.Context, d *schema.ResourceData, meta
 	d.Set("image_arn", image.ImageArn)
 	d.Set("container_image", image.ContainerImage)
 	d.Set(names.AttrVersion, image.Version)
-	d.Set("image_name", d.Id())
 	d.Set("horovod", image.Horovod)
 	d.Set("job_type", image.JobType)
 	d.Set("processor", image.Processor)
@@ -190,9 +215,12 @@ func resourceImageVersionUpdate(ctx context.Context, d *schema.ResourceData, met
 	var diags diag.Diagnostics
 	conn := meta.(*conns.AWSClient).SageMakerClient(ctx)
 
-	input := &sagemaker.UpdateImageVersionInput{
-		ImageName: aws.String(d.Id()),
-		Version:   aws.Int32(int32(d.Get(names.AttrVersion).(int))),
+	name := d.Get("image_name").(string)
+	version := d.Get(names.AttrVersion).(int)
+
+	input := sagemaker.UpdateImageVersionInput{
+		ImageName: aws.String(name),
+		Version:   aws.Int32(int32(version)),
 	}
 
 	if d.HasChange("horovod") {
@@ -223,7 +251,7 @@ func resourceImageVersionUpdate(ctx context.Context, d *schema.ResourceData, met
 		input.ProgrammingLang = aws.String(d.Get("programming_lang").(string))
 	}
 
-	if _, err := conn.UpdateImageVersion(ctx, input); err != nil {
+	if _, err := conn.UpdateImageVersion(ctx, &input); err != nil {
 		return sdkdiag.AppendErrorf(diags, "updating SageMaker AI Image Version (%s): %s", d.Id(), err)
 	}
 
@@ -234,9 +262,12 @@ func resourceImageVersionDelete(ctx context.Context, d *schema.ResourceData, met
 	var diags diag.Diagnostics
 	conn := meta.(*conns.AWSClient).SageMakerClient(ctx)
 
+	name := d.Get("image_name").(string)
+	version := d.Get(names.AttrVersion).(int)
+
 	input := &sagemaker.DeleteImageVersionInput{
-		ImageName: aws.String(d.Id()),
-		Version:   aws.Int32(int32(d.Get(names.AttrVersion).(int))),
+		ImageName: aws.String(name),
+		Version:   aws.Int32(int32(version)),
 	}
 
 	if _, err := conn.DeleteImageVersion(ctx, input); err != nil {
@@ -246,19 +277,24 @@ func resourceImageVersionDelete(ctx context.Context, d *schema.ResourceData, met
 		return sdkdiag.AppendErrorf(diags, "deleting SageMaker AI Image Version (%s): %s", d.Id(), err)
 	}
 
-	if _, err := waitImageVersionDeleted(ctx, conn, d.Id()); err != nil {
-		return sdkdiag.AppendErrorf(diags, "waiting for SageMaker AI Image Version (%s) to delete: %s", d.Id(), err)
+	if _, err := waitImageVersionDeleted(ctx, conn, name, version); err != nil {
+		return sdkdiag.AppendErrorf(diags, "waiting for SageMaker AI Image Version (%s) deletion: %s", d.Id(), err)
 	}
 
 	return diags
 }
 
+// findImageVersionByName is used immediately after creation to poll for status of
+// the most recently created image version
+//
+// findImageVersionByTwoPartKey should be used for all subsequent operations once the
+// version number is assigned.
 func findImageVersionByName(ctx context.Context, conn *sagemaker.Client, name string) (*sagemaker.DescribeImageVersionOutput, error) {
-	input := &sagemaker.DescribeImageVersionInput{
+	input := sagemaker.DescribeImageVersionInput{
 		ImageName: aws.String(name),
 	}
 
-	output, err := conn.DescribeImageVersion(ctx, input)
+	output, err := conn.DescribeImageVersion(ctx, &input)
 
 	if errs.IsAErrorMessageContains[*awstypes.ResourceNotFound](err, "does not exist") {
 		return nil, &retry.NotFoundError{
@@ -277,3 +313,48 @@ func findImageVersionByName(ctx context.Context, conn *sagemaker.Client, name st
 
 	return output, nil
 }
+
+// findImageVersionByTwoPartKey is used to fetch a specific version once a version number
+// is assigned
+func findImageVersionByTwoPartKey(ctx context.Context, conn *sagemaker.Client, name string, version int) (*sagemaker.DescribeImageVersionOutput, error) {
+	input := sagemaker.DescribeImageVersionInput{
+		ImageName: aws.String(name),
+		Version:   aws.Int32(int32(version)),
+	}
+
+	output, err := conn.DescribeImageVersion(ctx, &input)
+
+	if errs.IsAErrorMessageContains[*awstypes.ResourceNotFound](err, "does not exist") {
+		return nil, &retry.NotFoundError{
+			LastError:   err,
+			LastRequest: input,
+		}
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	if output == nil {
+		return nil, tfresource.NewEmptyResultError(input)
+	}
+
+	return output, nil
+}
+
+// expandImageVersionResourceID wraps flex.ExpandResourceId and handles conversion
+// of the version portion to an integer
+func expandImageVersionResourceID(id string) (string, int, error) {
+	parts, err := flex.ExpandResourceId(id, imageVersionResourcePartCount, false)
+	if err != nil {
+		return "", 0, err
+	}
+
+	name := parts[0]
+	version, err := strconv.Atoi(parts[1])
+	if err != nil {
+		return name, version, err
+	}
+
+	return name, version, nil
+}

internal/service/sagemaker/image_version_migrate.go

@@ -0,0 +1,100 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package sagemaker
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/YakDriver/regexache"
+	awstypes "github.com/aws/aws-sdk-go-v2/service/sagemaker/types"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
+	"github.com/hashicorp/terraform-provider-aws/internal/enum"
+	"github.com/hashicorp/terraform-provider-aws/names"
+)
+
+func resourceImageVersionV0() *schema.Resource {
+	return &schema.Resource{
+		Schema: map[string]*schema.Schema{
+			names.AttrARN: {
+				Type:     schema.TypeString,
+				Computed: true,
+			},
+			"base_image": {
+				Type:     schema.TypeString,
+				Required: true,
+				ForceNew: true,
+			},
+			"container_image": {
+				Type:     schema.TypeString,
+				Computed: true,
+			},
+			"horovod": {
+				Type:     schema.TypeBool,
+				Optional: true,
+			},
+			"image_arn": {
+				Type:     schema.TypeString,
+				Computed: true,
+			},
+			"image_name": {
+				Type:     schema.TypeString,
+				Required: true,
+				ForceNew: true,
+			},
+			"job_type": {
+				Type:             schema.TypeString,
+				Optional:         true,
+				ValidateDiagFunc: enum.Validate[awstypes.JobType](),
+			},
+			"ml_framework": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				ValidateFunc: validation.StringMatch(regexache.MustCompile(`^[a-zA-Z]+ ?\d+\.\d+(\.\d+)?$`), ""),
+			},
+			"processor": {
+				Type:             schema.TypeString,
+				Optional:         true,
+				ValidateDiagFunc: enum.Validate[awstypes.Processor](),
+			},
+			"programming_lang": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				ValidateFunc: validation.StringMatch(regexache.MustCompile(`^[a-zA-Z]+ ?\d+\.\d+(\.\d+)?$`), ""),
+			},
+			"release_notes": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				ValidateFunc: validation.StringLenBetween(0, 255),
+			},
+			"vendor_guidance": {
+				Type:             schema.TypeString,
+				Optional:         true,
+				ValidateDiagFunc: enum.Validate[awstypes.VendorGuidance](),
+			},
+			names.AttrVersion: {
+				Type:     schema.TypeInt,
+				Computed: true,
+			},
+		},
+	}
+}
+
+func imageVersionStateUpgradeV0(_ context.Context, rawState map[string]any, meta any) (map[string]any, error) {
+	if rawState == nil {
+		rawState = map[string]any{}
+	}
+
+	// The version attribute must be asserted into an int, otherwise fmt will assume
+	// a float64 when a numerical any value is detected.
+	version, ok := rawState[names.AttrVersion].(int)
+	if !ok {
+		version = 1
+	}
+
+	rawState[names.AttrID] = fmt.Sprintf("%s,%d", rawState["image_name"], version)
+
+	return rawState, nil
+}