session: add session callback before export completion

Signed-off-by: Justin Chadwell <me@jedevc.com>

client/client_test.go

@@ -46,6 +46,7 @@ import (
 	gatewaypb "github.com/moby/buildkit/frontend/gateway/pb"
 	"github.com/moby/buildkit/identity"
 	"github.com/moby/buildkit/session"
+	sessionexport "github.com/moby/buildkit/session/export"
 	"github.com/moby/buildkit/session/secrets/secretsprovider"
 	"github.com/moby/buildkit/session/sshforward/sshprovider"
 	"github.com/moby/buildkit/solver/errdefs"
@@ -208,6 +209,7 @@ func TestIntegration(t *testing.T) {
 		testMultipleRecordsWithSameLayersCacheImportExport,
 		testExportLocalNoPlatformSplit,
 		testExportLocalNoPlatformSplitOverwrite,
+		testSessionExportCallback,
 	)
 }
 
@@ -9647,3 +9649,32 @@ func testClientCustomGRPCOpts(t *testing.T, sb integration.Sandbox) {
 
 	require.Contains(t, interceptedMethods, "/moby.buildkit.v1.Control/Solve")
 }
+
+func testSessionExportCallback(t *testing.T, sb integration.Sandbox) {
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	callback := sessionexport.FinalizeExportCallback(func(resp map[string]string) (map[string]string, error) {
+		resp["image.name"] = "foobartest"
+		return resp, nil
+	})
+
+	st := llb.Image("busybox:latest")
+	def, err := st.Marshal(sb.Context())
+	require.NoError(t, err)
+	resp, err := c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: ExporterImage,
+				Attrs: map[string]string{
+					"name": "foobar",
+				},
+			},
+		},
+		Session: []session.Attachable{callback},
+	}, nil)
+	require.NoError(t, err)
+
+	require.Equal(t, "foobartest", resp.ExporterResponse["image.name"])
+}

exporter/containerimage/export.go

@@ -23,6 +23,7 @@ import (
 	"github.com/moby/buildkit/exporter"
 	"github.com/moby/buildkit/exporter/containerimage/exptypes"
 	"github.com/moby/buildkit/session"
+	sessionexport "github.com/moby/buildkit/session/export"
 	"github.com/moby/buildkit/snapshot"
 	"github.com/moby/buildkit/util/compression"
 	"github.com/moby/buildkit/util/contentutil"
@@ -331,6 +332,16 @@ func (e *imageExporterInstance) Export(ctx context.Context, src *exporter.Source
 	}
 	resp[exptypes.ExporterImageDescriptorKey] = base64.StdEncoding.EncodeToString(dtdesc)
 
+	caller, err := e.opt.SessionManager.Get(ctx, sessionID, false)
+	if err != nil {
+		return nil, nil, err
+	}
+	if resp2, err := sessionexport.Finalize(ctx, caller, resp); err != nil {
+		return nil, nil, err
+	} else if resp2 != nil {
+		resp = resp2
+	}
+
 	return resp, nil, nil
 }
 

exporter/local/export.go

@@ -12,6 +12,7 @@ import (
 	"github.com/moby/buildkit/exporter/containerimage/exptypes"
 	"github.com/moby/buildkit/exporter/util/epoch"
 	"github.com/moby/buildkit/session"
+	sessionexport "github.com/moby/buildkit/session/export"
 	"github.com/moby/buildkit/session/filesync"
 	"github.com/moby/buildkit/util/progress"
 	"github.com/pkg/errors"
@@ -171,7 +172,12 @@ func (e *localExporterInstance) Export(ctx context.Context, inp *exporter.Source
 	if err := eg.Wait(); err != nil {
 		return nil, nil, err
 	}
-	return nil, nil, nil
+
+	resp, err := sessionexport.Finalize(ctx, caller, nil)
+	if err != nil {
+		return nil, nil, err
+	}
+	return resp, nil, nil
 }
 
 func NewProgressHandler(ctx context.Context, id string) func(int, bool) {

exporter/oci/export.go

@@ -19,6 +19,7 @@ import (
 	"github.com/moby/buildkit/exporter/containerimage/exptypes"
 	"github.com/moby/buildkit/session"
 	sessioncontent "github.com/moby/buildkit/session/content"
+	sessionexport "github.com/moby/buildkit/session/export"
 	"github.com/moby/buildkit/session/filesync"
 	"github.com/moby/buildkit/util/compression"
 	"github.com/moby/buildkit/util/contentutil"
@@ -271,6 +272,12 @@ func (e *imageExporterInstance) Export(ctx context.Context, src *exporter.Source
 		}
 	}
 
+	if resp2, err := sessionexport.Finalize(ctx, caller, resp); err != nil {
+		return nil, nil, err
+	} else if resp2 != nil {
+		resp = resp2
+	}
+
 	return resp, nil, nil
 }
 

exporter/tar/export.go

@@ -12,6 +12,7 @@ import (
 	"github.com/moby/buildkit/exporter/local"
 	"github.com/moby/buildkit/exporter/util/epoch"
 	"github.com/moby/buildkit/session"
+	sessionexport "github.com/moby/buildkit/session/export"
 	"github.com/moby/buildkit/session/filesync"
 	"github.com/moby/buildkit/util/progress"
 	"github.com/pkg/errors"
@@ -160,5 +161,14 @@ func (e *localExporterInstance) Export(ctx context.Context, inp *exporter.Source
 		w.Close()
 		return nil, nil, report(err)
 	}
-	return nil, nil, report(w.Close())
+	if err := w.Close(); err != nil {
+		return nil, nil, report(err)
+	}
+	report(nil)
+
+	resp, err := sessionexport.Finalize(ctx, caller, nil)
+	if err != nil {
+		return nil, nil, err
+	}
+	return resp, nil, nil
 }

session/export/export.go

@@ -0,0 +1,33 @@
+package export
+
+import (
+	"context"
+
+	"github.com/moby/buildkit/session"
+	"google.golang.org/grpc"
+)
+
+type FinalizeExportCallback func(map[string]string) (map[string]string, error)
+
+var _ session.Attachable = (*FinalizeExportCallback)(nil)
+
+func (finalizer FinalizeExportCallback) Register(server *grpc.Server) {
+	RegisterExportCallbackServer(server, finalizer)
+}
+
+func (finalizer FinalizeExportCallback) Finalize(ctx context.Context, req *FinalizeRequest) (*FinalizeResponse, error) {
+	resp, err := finalizer(req.ExporterResponse)
+	return &FinalizeResponse{ExporterResponse: resp}, err
+}
+
+func Finalize(ctx context.Context, c session.Caller, exporterResponse map[string]string) (map[string]string, error) {
+	if !c.Supports(session.MethodURL(_ExportCallback_serviceDesc.ServiceName, "finalize")) {
+		return nil, nil
+	}
+	client := NewExportCallbackClient(c.Conn())
+	resp, err := client.Finalize(ctx, &FinalizeRequest{ExporterResponse: exporterResponse})
+	if resp.ExporterResponse == nil {
+		resp.ExporterResponse = map[string]string{}
+	}
+	return resp.ExporterResponse, err
+}