lostluck commented on code in PR #29410: URL: https://github.com/apache/beam/pull/29410#discussion_r1416102636
########## sdks/go/pkg/beam/io/natsio/read.go: ########## @@ -0,0 +1,289 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package natsio + +import ( + "context" + "fmt" + "math" + "reflect" + "time" + + "github.com/apache/beam/sdks/v2/go/pkg/beam" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf" + "github.com/apache/beam/sdks/v2/go/pkg/beam/io/rtrackers/offsetrange" + "github.com/apache/beam/sdks/v2/go/pkg/beam/register" + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" +) + +func init() { + register.DoFn5x2[ + context.Context, *watermarkEstimator, *sdf.LockRTracker, []byte, + func(beam.EventTime, ConsumerMessage), sdf.ProcessContinuation, error, + ]( + &readFn{}, + ) + register.Emitter2[beam.EventTime, ConsumerMessage]() + beam.RegisterType(reflect.TypeOf((*ConsumerMessage)(nil)).Elem()) +} + +const ( + defaultFetchSize = 100 + defaultStartSeqNo = 1 + defaultEndSeqNo = math.MaxInt64 + fetchTimeout = 3 * time.Second + assumedLag = 1 * time.Second + resumeDelay = 5 * time.Second +) + +type ConsumerMessage struct { + Subject string + PublishingTime time.Time + ID string + Headers map[string][]string + Data []byte +} + +// Read reads messages from NATS JetStream and returns a PCollection<ConsumerMessage>. +// Read takes a variable number of ReadOptionFn to configure the read operation: +// - UserCredentials: path to the user credentials file. Defaults to empty. +// - ProcessingTimePolicy: whether to use the pipeline processing time of the messages as the event +// time. Defaults to true. +// - PublishingTimePolicy: whether to use the publishing time of the messages as the event time. +// Defaults to false. +// - FetchSize: the maximum number of messages to retrieve at a time. Defaults to 100. +// - StartSeqNo: the start sequence number of messages to read. Defaults to 1. +// - EndSeqNo: the end sequence number of messages to read (exclusive). Defaults to math.MaxInt64. +func Read( + s beam.Scope, + uri string, + stream string, + subject string, + opts ...ReadOptionFn, +) beam.PCollection { + s = s.Scope("natsio.Read") + + option := &readOption{ + TimePolicy: processingTimePolicy, + FetchSize: defaultFetchSize, + StartSeqNo: defaultStartSeqNo, + EndSeqNo: defaultEndSeqNo, + } + + for _, opt := range opts { + if err := opt(option); err != nil { + panic(fmt.Sprintf("natsio.Read: invalid option: %v", err)) + } + } + + imp := beam.Impulse(s) + return beam.ParDo(s, newReadFn(uri, stream, subject, option), imp) +} + +type readFn struct { + natsFn + Stream string + Subject string + TimePolicy timePolicy + FetchSize int + StartSeqNo int64 + EndSeqNo int64 + timestampFn timestampFn +} + +func newReadFn(uri string, stream string, subject string, option *readOption) *readFn { + return &readFn{ + natsFn: natsFn{ + URI: uri, + CredsFile: option.CredsFile, + }, + Stream: stream, + Subject: subject, + TimePolicy: option.TimePolicy, + FetchSize: option.FetchSize, + StartSeqNo: option.StartSeqNo, + EndSeqNo: option.EndSeqNo, + } +} + +func (fn *readFn) Setup() error { + if err := fn.natsFn.Setup(); err != nil { + return err + } + + fn.timestampFn = fn.TimePolicy.TimestampFn() + return nil +} + +func (fn *readFn) CreateInitialRestriction(_ []byte) offsetrange.Restriction { + return offsetrange.Restriction{ + Start: fn.StartSeqNo, + End: fn.EndSeqNo, + } +} + +func (fn *readFn) SplitRestriction( + _ []byte, + rest offsetrange.Restriction, +) []offsetrange.Restriction { + return []offsetrange.Restriction{rest} Review Comment: OK, so this means no initial splits, so a single stream to start, and only dynamic splits will occur. I think this is correct WRT the rest of the code. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
