This repository has been archived by the owner on Oct 9, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 53
/
remote_file_output_writer.go
86 lines (68 loc) · 2.19 KB
/
remote_file_output_writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
package ioutils
import (
"context"
"fmt"
"github.com/lyft/flyteidl/gen/pb-go/flyteidl/core"
"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io"
"github.com/lyft/flytestdlib/storage"
)
type RemoteFileOutputWriter struct {
io.OutputFilePaths
store storage.ProtobufStore
}
type RemoteFileOutputPaths struct {
outputPrefix storage.DataReference
store storage.ReferenceConstructor
}
var (
_ io.OutputWriter = RemoteFileOutputWriter{}
_ io.OutputFilePaths = RemoteFileOutputPaths{}
)
func (w RemoteFileOutputPaths) GetOutputPrefixPath() storage.DataReference {
return w.outputPrefix
}
func (w RemoteFileOutputPaths) GetOutputPath() storage.DataReference {
return constructPath(w.store, w.outputPrefix, OutputsSuffix)
}
func (w RemoteFileOutputPaths) GetErrorPath() storage.DataReference {
return constructPath(w.store, w.outputPrefix, ErrorsSuffix)
}
func (w RemoteFileOutputPaths) GetFuturesPath() storage.DataReference {
return constructPath(w.store, w.outputPrefix, FuturesSuffix)
}
func (w RemoteFileOutputWriter) Put(ctx context.Context, reader io.OutputReader) error {
literals, executionErr, err := reader.Read(ctx)
if err != nil {
return err
}
if executionErr != nil {
errorKind := core.ContainerError_RECOVERABLE
if !executionErr.IsRecoverable {
errorKind = core.ContainerError_NON_RECOVERABLE
}
errDoc := &core.ErrorDocument{
Error: &core.ContainerError{
Code: executionErr.Code,
Message: executionErr.Message,
Kind: errorKind,
},
}
return w.store.WriteProtobuf(ctx, w.GetErrorPath(), storage.Options{}, errDoc)
}
if literals != nil {
return w.store.WriteProtobuf(ctx, w.GetOutputPath(), storage.Options{}, literals)
}
return fmt.Errorf("no data found to write")
}
func NewRemoteFileOutputPaths(_ context.Context, store storage.ReferenceConstructor, outputPrefix storage.DataReference) RemoteFileOutputPaths {
return RemoteFileOutputPaths{
store: store,
outputPrefix: outputPrefix,
}
}
func NewRemoteFileOutputWriter(_ context.Context, store storage.ProtobufStore, outputFilePaths io.OutputFilePaths) RemoteFileOutputWriter {
return RemoteFileOutputWriter{
OutputFilePaths: outputFilePaths,
store: store,
}
}