-
Notifications
You must be signed in to change notification settings - Fork 342
/
orc_utils.go
145 lines (136 loc) · 4.27 KB
/
orc_utils.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
package s3inventory
import (
"context"
"fmt"
"io"
"os"
"path"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/aws/aws-sdk-go/service/s3/s3iface"
"github.com/aws/aws-sdk-go/service/s3/s3manager"
gproto "github.com/golang/protobuf/proto" //nolint:staticcheck // orc lib uses old proto
"github.com/scritchley/orc/proto"
"github.com/treeverse/lakefs/pkg/logging"
)
const (
maxPostScriptSize = 256 // from the ORC specification: https://orc.apache.org/specification/ORCv1/
orcInitialReadSize = 16000
downloadTimeout = 3 * time.Minute
)
// getTailLength reads the ORC postscript from the given file, returning the full tail length.
// The tail length equals (footer + metadata + postscript + 1) bytes.
func getTailLength(f *os.File) (int, error) {
stat, err := f.Stat()
if err != nil {
return 0, err
}
fileSize := stat.Size()
psPlusByte := int64(maxPostScriptSize + 1)
if psPlusByte > fileSize {
psPlusByte = fileSize
}
// Read the last 256 bytes into buffer to get postscript
postScriptBytes := make([]byte, psPlusByte)
sr := io.NewSectionReader(f, fileSize-psPlusByte, psPlusByte)
_, err = io.ReadFull(sr, postScriptBytes)
if err != nil {
return 0, err
}
psLen := int(postScriptBytes[len(postScriptBytes)-1])
psOffset := len(postScriptBytes) - 1 - psLen
postScript := &proto.PostScript{}
err = gproto.Unmarshal(postScriptBytes[psOffset:psOffset+psLen], postScript)
if err != nil {
return 0, err
}
footerLength := int(postScript.GetFooterLength())
metadataLength := int(postScript.GetMetadataLength())
return footerLength + metadataLength + psLen + 1, nil
}
func downloadRange(ctx context.Context, svc s3iface.S3API, logger logging.Logger, bucket string, key string, fromByte int64) (*os.File, error) {
f, err := os.CreateTemp("", path.Base(key))
if err != nil {
return nil, err
}
defer func() {
if err := os.Remove(f.Name()); err != nil {
logger.WithError(err).WithField("file", f.Name()).Error("Failed to remove orc file after download")
}
}()
downloader := s3manager.NewDownloaderWithClient(svc)
var rng *string
if fromByte > 0 {
rng = aws.String(fmt.Sprintf("bytes=%d-", fromByte))
}
logger.Debugf("start downloading %s[%s] to local file %s", key, aws.StringValue(rng), f.Name())
timeoutCtx, cancelFn := context.WithTimeout(ctx, downloadTimeout)
defer cancelFn()
_, err = downloader.DownloadWithContext(timeoutCtx, f, &s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
Range: rng,
})
if err != nil {
logger.WithError(err).
WithFields(logging.Fields{
"download_from_bucket": bucket,
"download_from_key": key,
"download_to": f.Name(),
}).Debug("error when downloading orc file")
return nil, err
}
logger.Debugf("finished downloading %s to local file %s", key, f.Name())
return f, nil
}
// DownloadOrc downloads a file from s3 and returns a ReaderSeeker to it.
// If tailOnly is set to true, download only the tail (metadata+footer) by trying the last `orcInitialReadSize` bytes of the file.
// Then, check the last byte to see if the whole tail was downloaded. If not, download again with the actual tail length.
func DownloadOrc(ctx context.Context, svc s3iface.S3API, logger logging.Logger, bucket string, key string, tailOnly bool) (*OrcFile, error) {
var size int64
if tailOnly {
headObject, err := svc.HeadObject(&s3.HeadObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
})
if err != nil {
return nil, err
}
size = *headObject.ContentLength
}
f, err := downloadRange(ctx, svc, logger, bucket, key, size-orcInitialReadSize)
if err != nil {
return nil, err
}
if tailOnly {
tailLength, err := getTailLength(f)
if err != nil {
return nil, err
}
if tailLength > orcInitialReadSize {
// tail didn't fit in initially downloaded file
if err = f.Close(); err != nil {
logger.WithError(err).WithField("file", f.Name()).Error("failed to close orc file")
}
f, err = downloadRange(ctx, svc, logger, bucket, key, size-int64(tailLength))
if err != nil {
return nil, err
}
}
}
return &OrcFile{f}, nil
}
type OrcFile struct {
*os.File
}
func (or *OrcFile) Size() int64 {
stats, err := or.Stat()
if err != nil {
return 0
}
return stats.Size()
}
func (or *OrcFile) Close() error {
return or.File.Close()
}