/
c2k.go
139 lines (126 loc) · 4.89 KB
/
c2k.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
package main
import (
"bufio"
"flag"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/credentials"
"github.com/aws/aws-sdk-go/service/firehose"
"github.com/aws/aws-sdk-go/service/kinesis"
"io"
"log"
"os"
)
const (
defaultDelimiter = "\n"
delimiterUsage = "Delimiter to split on (defaults to newline)"
ItrUsage = "Type of Shard Iterator to use. Valid choices: AT_SEQUENCE_NUMBER, AFTER_SEQUENCE_NUMBER, TRIM_HORIZON"
listenUsage = "Listen to stream instead of sending data"
defaultProfile = "default"
profileUsage = "AWS Profile name to use for authentication"
defaultPartitionKey = "1"
partitionKeyUsage = "Partition key"
defaultRegion = "us-east-1"
regionUsage = "AWS region, defaults to us-east-1"
startingSeqNumUsage = "Sequence number to use for iterators that use a sequence number"
shardIdUsage = "Shard ID for listen purposes"
defaultShardId string = "ALL"
streamNameUsage = "Stream name to put data"
incompleteRead = "c2k: incomplete read of stream"
noSuchFile = "c2k: %s: no such file"
TrimHorizon string = "TRIM_HORIZON"
AtSequenceNum string = "AT_SEQUENCE_NUMBER"
AfterSequenceNum string = "AFTER_SEQUENCE_NUMBER"
Latest string = "LATEST"
)
type Options struct {
Delimiter, Profile, Region, ShardId, StartingSeqNum, StreamName, PartitionKey, ItrType string
Firehose bool
}
func main() {
var listen bool
opts := parseArgs(&listen)
svc := createService(opts.Profile, opts.Region)
fsvc := createFirehoseService(opts.Profile, opts.Region)
if listen {
listener := NewListener(opts, svc)
listener.Listen(os.Stdout)
} else {
if len(flag.Args()) == 0 {
uploadFile("-", opts, svc, fsvc)
} else {
for _, fileName := range flag.Args() {
uploadFile(fileName, opts, svc, fsvc)
}
}
}
}
func createService(profile, region string) *kinesis.Kinesis {
creds := credentials.NewSharedCredentials("", profile)
return kinesis.New(&aws.Config{Region: aws.String(region), Credentials: creds})
}
func createFirehoseService(profile, region string) *firehose.Firehose {
creds := credentials.NewSharedCredentials("", profile)
return firehose.New(&aws.Config{Region: aws.String(region), Credentials: creds})
}
func parseArgs(listen *bool) Options {
opts := Options{}
flag.StringVar(&opts.Delimiter, "delimiter", defaultDelimiter, delimiterUsage)
flag.StringVar(&opts.Delimiter, "d", defaultDelimiter, delimiterUsage+" (short)")
flag.BoolVar(&opts.Firehose, "f", false, "Firehose mode")
flag.StringVar(&opts.ItrType, "iter", TrimHorizon, ItrUsage)
flag.StringVar(&opts.ItrType, "i", TrimHorizon, ItrUsage+" (short)")
flag.BoolVar(listen, "listen", false, listenUsage)
flag.BoolVar(listen, "l", false, listenUsage+" (short)")
flag.StringVar(&opts.PartitionKey, "partitionKey", defaultPartitionKey, partitionKeyUsage)
flag.StringVar(&opts.PartitionKey, "pk", defaultPartitionKey, partitionKeyUsage+" (short)")
flag.StringVar(&opts.Profile, "profile", defaultProfile, profileUsage)
flag.StringVar(&opts.Profile, "p", defaultProfile, profileUsage+" (short)")
flag.StringVar(&opts.Region, "region", defaultRegion, regionUsage)
flag.StringVar(&opts.Region, "r", defaultRegion, regionUsage+" (short)")
flag.StringVar(&opts.StartingSeqNum, "startingSeqNum", "", startingSeqNumUsage)
flag.StringVar(&opts.StartingSeqNum, "sn", "", startingSeqNumUsage+" (short)")
flag.StringVar(&opts.ShardId, "shardId", defaultShardId, shardIdUsage)
flag.StringVar(&opts.ShardId, "sId", defaultShardId, shardIdUsage+" (short)")
flag.StringVar(&opts.StreamName, "streamName", "", streamNameUsage)
flag.StringVar(&opts.StreamName, "s", "", streamNameUsage+" (short)")
flag.Parse()
if opts.StreamName == "" {
log.Fatal("streamName is a required parameter")
}
if *listen && opts.ItrType != TrimHorizon && opts.ItrType != Latest && opts.ItrType != AfterSequenceNum && opts.ItrType != AtSequenceNum {
log.Fatal("Invalid iter type given ", opts.ItrType)
}
return opts
}
func uploadFile(fileName string, opts Options, svc *kinesis.Kinesis, fsvc *firehose.Firehose) {
var handle *os.File
var err error
if fileName == "-" {
log.Println("Reading from stdin")
handle = os.Stdin
} else {
handle, err = os.Open(fileName)
if err != nil {
log.Printf(noSuchFile, fileName)
return
}
}
defer handle.Close()
rdr := bufio.NewReader(handle)
putFromReader(rdr, opts, svc, fsvc)
}
func putFromReader(rdr *bufio.Reader, opts Options, svc *kinesis.Kinesis, fsvc *firehose.Firehose) {
uploader := NewUploader(svc, fsvc, opts)
defer uploader.Flush()
for {
data, err := rdr.ReadBytes([]byte(opts.Delimiter)[0])
if err == io.EOF {
break
}
if err != nil {
log.Printf(incompleteRead)
break
}
uploader.Upload(data)
}
}