Merge branch 'fix/pull-image-not-found' into 'master'

fix(retrieval): pull image when inspect returns wrapped not-found

Closes #700

See merge request postgres-ai/database-lab!1129

Author: Sarumyan

engine/internal/retrieval/engine/postgres/logical/dump.go

@@ -286,7 +286,7 @@ func (d *DumpJob) Run(ctx context.Context) (err error) {
 	}
 
 	if err := tools.PullImage(ctx, d.dockerClient, d.DockerImage); err != nil {
-		return errors.Wrap(err, "failed to scan pulling image response")
+		return fmt.Errorf("failed to prepare dump image %q: %w", d.DockerImage, err)
 	}
 
 	if err := os.MkdirAll(d.DumpOptions.DumpLocation, 0666); err != nil {

engine/internal/retrieval/engine/postgres/tools/tools.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/AlekSi/pointer"
 	"github.com/ahmetalpbalkan/dlog"
+	cerrdefs "github.com/containerd/errdefs"
 	"github.com/docker/cli/cli/streams"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
@@ -30,7 +31,6 @@ import (
 	"github.com/docker/docker/api/types/mount"
 	"github.com/docker/docker/api/types/network"
 	"github.com/docker/docker/client"
-	"github.com/docker/docker/errdefs"
 	"github.com/docker/docker/pkg/jsonmessage"
 	"github.com/docker/docker/pkg/stdcopy"
 	"github.com/pkg/errors"
@@ -518,13 +518,18 @@ func RemoveContainer(ctx context.Context, dockerClient *client.Client, container
 	log.Msg(fmt.Sprintf("Container %q has been removed", containerID))
 }
 
-// PullImage pulls a Docker image.
-func PullImage(ctx context.Context, dockerClient *client.Client, image string) error {
+// imagePuller is the minimal subset of the Docker client used by PullImage.
+// Defined as an interface so PullImage can be unit-tested with a fake client.
+type imagePuller interface {
+	ImageInspect(ctx context.Context, imageID string, opts ...client.ImageInspectOption) (imagetypes.InspectResponse, error)
+	ImagePull(ctx context.Context, refStr string, options imagetypes.PullOptions) (io.ReadCloser, error)
+}
+
+// PullImage pulls a Docker image if it is not already present locally.
+func PullImage(ctx context.Context, dockerClient imagePuller, image string) error {
 	inspectionResult, err := dockerClient.ImageInspect(ctx, image)
-	if err != nil {
-		if _, ok := err.(errdefs.ErrNotFound); !ok {
-			return errors.Wrapf(err, "failed to inspect image %q", image)
-		}
+	if err != nil && !cerrdefs.IsNotFound(err) {
+		return errors.Wrapf(err, "failed to inspect image %q", image)
 	}
 
 	if err == nil && inspectionResult.ID != "" {
@@ -540,7 +545,7 @@ func PullImage(ctx context.Context, dockerClient *client.Client, image string) e
 	defer func() { _ = pullOutput.Close() }()
 
 	if err := jsonmessage.DisplayJSONMessagesToStream(pullOutput, streams.NewOut(os.Stdout), nil); err != nil {
-		log.Err("failed to render pull image output: ", err)
+		return errors.Wrapf(err, "failed to pull image %q", image)
 	}
 
 	return nil

engine/internal/retrieval/engine/postgres/tools/tools_test.go

@@ -6,12 +6,20 @@ package tools
 
 import (
 	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
+	cerrdefs "github.com/containerd/errdefs"
 	"github.com/docker/docker/api/types/container"
+	imagetypes "github.com/docker/docker/api/types/image"
 	"github.com/docker/docker/api/types/mount"
+	"github.com/docker/docker/client"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -317,3 +325,86 @@ func TestIsEmptyDirectory_NonExistent(t *testing.T) {
 	_, err := IsEmptyDirectory("/nonexistent/path/12345")
 	assert.Error(t, err)
 }
+
+type inspectFn func(ctx context.Context, imageID string, opts ...client.ImageInspectOption) (imagetypes.InspectResponse, error)
+type pullFn func(ctx context.Context, refStr string, options imagetypes.PullOptions) (io.ReadCloser, error)
+
+// fakePuller implements imagePuller for testing PullImage in isolation.
+type fakePuller struct {
+	inspect    inspectFn
+	pull       pullFn
+	pullCalled bool
+}
+
+func (f *fakePuller) ImageInspect(ctx context.Context, imageID string, opts ...client.ImageInspectOption) (imagetypes.InspectResponse, error) {
+	if f.inspect != nil {
+		return f.inspect(ctx, imageID, opts...)
+	}
+	return imagetypes.InspectResponse{}, errors.New("inspect not implemented")
+}
+
+func (f *fakePuller) ImagePull(ctx context.Context, refStr string, options imagetypes.PullOptions) (io.ReadCloser, error) {
+	f.pullCalled = true
+	if f.pull != nil {
+		return f.pull(ctx, refStr, options)
+	}
+	return nil, errors.New("pull not implemented")
+}
+
+func TestPullImage(t *testing.T) {
+	notFound := fmt.Errorf("no such image: %w", cerrdefs.ErrNotFound)
+	wrappedNotFound := fmt.Errorf("inspect failed: %w", notFound)
+
+	localImage := func(context.Context, string, ...client.ImageInspectOption) (imagetypes.InspectResponse, error) {
+		return imagetypes.InspectResponse{ID: "sha256:abc"}, nil
+	}
+	notFoundInspect := func(context.Context, string, ...client.ImageInspectOption) (imagetypes.InspectResponse, error) {
+		return imagetypes.InspectResponse{}, notFound
+	}
+	wrappedNotFoundInspect := func(context.Context, string, ...client.ImageInspectOption) (imagetypes.InspectResponse, error) {
+		return imagetypes.InspectResponse{}, wrappedNotFound
+	}
+	daemonDownInspect := func(context.Context, string, ...client.ImageInspectOption) (imagetypes.InspectResponse, error) {
+		return imagetypes.InspectResponse{}, errors.New("daemon down")
+	}
+
+	emptyStream := func(context.Context, string, imagetypes.PullOptions) (io.ReadCloser, error) {
+		return io.NopCloser(strings.NewReader("")), nil
+	}
+	errorStream := func(context.Context, string, imagetypes.PullOptions) (io.ReadCloser, error) {
+		body := `{"errorDetail":{"message":"manifest for postgresai/extended-postgres:99-bogus not found: manifest unknown"},"error":"manifest unknown"}`
+		return io.NopCloser(strings.NewReader(body)), nil
+	}
+
+	testCases := []struct {
+		name            string
+		inspect         inspectFn
+		pull            pullFn
+		wantErr         bool
+		wantErrContains string
+		wantPullCalled  bool
+	}{
+		{name: "image already local skips pull", inspect: localImage},
+		{name: "not found triggers pull", inspect: notFoundInspect, pull: emptyStream, wantPullCalled: true},
+		{name: "wrapped not found triggers pull", inspect: wrappedNotFoundInspect, pull: emptyStream, wantPullCalled: true},
+		{name: "non-not-found inspect error propagates", inspect: daemonDownInspect, wantErr: true, wantErrContains: "failed to inspect image"},
+		{name: "pull stream error propagates", inspect: notFoundInspect, pull: errorStream, wantErr: true, wantErrContains: "failed to pull image", wantPullCalled: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			puller := &fakePuller{inspect: tc.inspect, pull: tc.pull}
+
+			err := PullImage(context.Background(), puller, "postgresai/extended-postgres:test")
+
+			if tc.wantErr {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tc.wantErrContains)
+			} else {
+				require.NoError(t, err)
+			}
+
+			assert.Equal(t, tc.wantPullCalled, puller.pullCalled)
+		})
+	}
+}

engine/internal/runci/handlers.go

@@ -163,7 +163,7 @@ func (s *Server) runMigration(w http.ResponseWriter, r *http.Request) {
 func (s *Server) runCommands(ctx context.Context, clone *models.Clone, runID string, volumes, tags map[string]string,
 	commands, migrationEnvs []string, cfg dblab_types.Config) (*observer.Session, error) {
 	if err := tools.PullImage(ctx, s.docker, s.config.Runner.Image); err != nil {
-		return nil, errors.Wrap(err, "failed to scan pulling image response")
+		return nil, fmt.Errorf("failed to prepare runner image %q: %w", s.config.Runner.Image, err)
 	}
 
 	containerCfg := s.buildContainerConfig(clone, migrationEnvs)