actually use digest if one is provided

beam-cloud · Jan 10, 2025 · 6b17281 · 6b17281
1 parent a21a99c
commit 6b17281
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/pkg/abstractions/image/build.go b/pkg/abstractions/image/build.go
@@ -55,6 +55,7 @@ type BuildOpts struct {
 	BaseImageRegistry  string
 	BaseImageName      string
 	BaseImageTag       string
+	BaseImageDigest    string
 	BaseImageCreds     string
 	ExistingImageUri   string
 	ExistingImageCreds map[string]string
@@ -75,6 +76,7 @@ func (o *BuildOpts) String() string {
 	fmt.Fprintf(&b, "  \"BaseImageRegistry\": %q,", o.BaseImageRegistry)
 	fmt.Fprintf(&b, "  \"BaseImageName\": %q,", o.BaseImageName)
 	fmt.Fprintf(&b, "  \"BaseImageTag\": %q,", o.BaseImageTag)
+	fmt.Fprintf(&b, "  \"BaseImageDigest\": %q,", o.BaseImageDigest)
 	fmt.Fprintf(&b, "  \"BaseImageCreds\": %q,", o.BaseImageCreds)
 	fmt.Fprintf(&b, "  \"ExistingImageUri\": %q,", o.ExistingImageUri)
 	fmt.Fprintf(&b, "  \"ExistingImageCreds\": %#v,", o.ExistingImageCreds)
@@ -106,6 +108,7 @@ func (o *BuildOpts) setCustomImageBuildOptions() error {
 	o.BaseImageRegistry = baseImage.Registry
 	o.BaseImageName = baseImage.Repo
 	o.BaseImageTag = baseImage.Tag
+	o.BaseImageDigest = baseImage.Digest
 
 	return nil
 }
@@ -180,7 +183,7 @@ func (b *Builder) GetImageId(opts *BuildOpts) (string, error) {
 
 	bodyToHash := &ImageIdHash{
 		BaseImageName:   opts.BaseImageName,
-		BaseImageTag:    opts.BaseImageTag,
+		BaseImageTag:    tagOrDigest(opts.BaseImageDigest, opts.BaseImageTag),
 		PythonVersion:   opts.PythonVersion,
 		PythonPackages:  opts.PythonPackages,
 		ExitingImageUri: opts.ExistingImageUri,
@@ -232,6 +235,7 @@ func (b *Builder) Build(ctx context.Context, opts *BuildOpts, outputChan chan co
 		BaseImageRegistry: opts.BaseImageRegistry,
 		BaseImageName:     opts.BaseImageName,
 		BaseImageTag:      opts.BaseImageTag,
+		BaseImageDigest:   opts.BaseImageDigest,
 		ExistingImageUri:  opts.ExistingImageUri,
 		EnvVars:           opts.EnvVars,
 		Dockerfile:        opts.Dockerfile,
@@ -242,7 +246,14 @@ func (b *Builder) Build(ctx context.Context, opts *BuildOpts, outputChan chan co
 		return err
 	}
 
-	sourceImage := fmt.Sprintf("%s/%s:%s", opts.BaseImageRegistry, opts.BaseImageName, opts.BaseImageTag)
+	var sourceImage string
+	switch {
+	case opts.BaseImageDigest != "":
+		sourceImage = fmt.Sprintf("%s/%s@%s", opts.BaseImageRegistry, opts.BaseImageName, opts.BaseImageDigest)
+	default:
+		sourceImage = fmt.Sprintf("%s/%s:%s", opts.BaseImageRegistry, opts.BaseImageName, opts.BaseImageTag)
+	}
+
 	containerId := b.genContainerId()
 
 	// Allow config to override default build container settings
@@ -748,3 +759,10 @@ func extractPackageName(pkg string) string {
 	// Handle regular packages
 	return strings.FieldsFunc(pkg, func(c rune) bool { return c == '=' || c == '>' || c == '<' || c == '[' || c == ';' })[0]
 }
+
+func tagOrDigest(digest string, tag string) string {
+	if tag != "" {
+		return tag
+	}
+	return digest
+}
diff --git a/pkg/abstractions/image/build_test.go b/pkg/abstractions/image/build_test.go
@@ -89,6 +89,12 @@ func TestExtractImageNameAndTag(t *testing.T) {
 			wantRepo:     "gis-ops/docker-valhalla/valhalla",
 			wantRegistry: "ghcr.io",
 		},
+		{
+			ref:          "us-east1-docker.pkg.dev/test/ds-us-east1/test-test-test@sha256:c31c45b6fdc3d01c131a6dcae1daed008e3df4001bb43e555e49d82ac8d779e4",
+			wantDigest:   "sha256:c31c45b6fdc3d01c131a6dcae1daed008e3df4001bb43e555e49d82ac8d779e4",
+			wantRepo:     "test/ds-us-east1/test-test-test",
+			wantRegistry: "us-east1-docker.pkg.dev",
+		},
 	}
 
 	for _, test := range tests {