Add GetByIds (#364)

Author: adiom-mark

connectors/airbyte/conn.go

@@ -15,6 +15,7 @@ import (
 )
 
 type source struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	dockerImage string
 	config      string
 

connectors/dynamodb/conn.go

@@ -24,6 +24,7 @@ import (
 )
 
 type conn struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	client        *client
 	streamsClient *dynamodbstreams.Client
 	spec          string

connectors/mongo/conn.go

@@ -64,6 +64,7 @@ type buffer struct {
 }
 
 type conn struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	client *mongo.Client
 
 	settings ConnectorSettings
@@ -326,6 +327,7 @@ func (c *conn) GetInfo(ctx context.Context, r *connect.Request[adiomv1.GetInfoRe
 				LsnStream:          true,
 				MultiNamespacePlan: true,
 				DefaultPlan:        !c.settings.PerNamespaceStreams,
+				GetByIds:           true,
 			},
 			Sink: &adiomv1.Capabilities_Sink{
 				SupportedDataTypes: []adiomv1.DataType{adiomv1.DataType_DATA_TYPE_MONGO_BSON},
@@ -1046,3 +1048,48 @@ func maybeUnavailableError(err error) error {
 	}
 	return connect.NewError(connect.CodeInternal, err)
 }
+
+// GetByIds implements adiomv1connect.ConnectorServiceHandler.
+func (c *conn) GetByIds(ctx context.Context, r *connect.Request[adiomv1.GetByIdsRequest]) (*connect.Response[adiomv1.GetByIdsResponse], error) {
+	col, _, ok := GetCol(c.client, r.Msg.GetNamespace())
+	if !ok {
+		return nil, connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("namespace should be fully qualified"))
+	}
+
+	// TODO: maybe use batch endpoint if we need to optimize
+	res := make([]*adiomv1.GetByIdsResponse_ResponseItem, len(r.Msg.GetIds()))
+	var eg errgroup.Group
+	for i, id := range r.Msg.GetIds() {
+		eg.Go(func() error {
+			if len(id.GetId()) < 1 {
+				res[i] = &adiomv1.GetByIdsResponse_ResponseItem{}
+				return nil
+			}
+			bv := id.GetId()[0]
+			rawVal := bson.RawValue{
+				Type:  bsontype.Type(bv.GetType()),
+				Value: bv.GetData(),
+			}
+			v, err := col.FindOne(ctx, bson.M{"_id": rawVal}).Raw()
+			if err != nil {
+				if errors.Is(err, mongo.ErrNoDocuments) {
+					res[i] = &adiomv1.GetByIdsResponse_ResponseItem{}
+					return nil
+				}
+				return fmt.Errorf("err in findone: %w", err)
+			}
+			res[i] = &adiomv1.GetByIdsResponse_ResponseItem{
+				Data: v,
+			}
+			return nil
+		})
+	}
+
+	if err := eg.Wait(); err != nil {
+		return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("err finding ids: %w", err))
+	}
+
+	return connect.NewResponse(&adiomv1.GetByIdsResponse{
+		Data: res,
+	}), nil
+}

connectors/null/connector.go

@@ -21,6 +21,7 @@ import (
 )
 
 type conn struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	id          string
 	logJson     bool
 	sleep       time.Duration
@@ -146,5 +147,10 @@ func (c *conn) WriteUpdates(ctx context.Context, r *connect.Request[adiomv1.Writ
 }
 
 func NewConn(id string, logJson bool, sleep time.Duration, sleepJitter time.Duration) adiomv1connect.ConnectorServiceHandler {
-	return &conn{id, logJson, sleep, sleepJitter}
+	return &conn{
+		id:          id,
+		logJson:     logJson,
+		sleep:       sleep,
+		sleepJitter: sleepJitter,
+	}
 }

connectors/postgres/conn.go

@@ -52,6 +52,7 @@ type PostgresSettings struct {
 }
 
 type conn struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	replicationUrl string
 	id             uint64
 	c              *pgxpool.Pool

connectors/random/connector.go

@@ -19,6 +19,7 @@ import (
 )
 
 type conn struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	settings ConnectorSettings
 
 	docMap      map[string]*IndexMap //map of locations to map of document IDs

connectors/random/connv2.go

@@ -20,6 +20,7 @@ import (
 )
 
 type connV2 struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	payload                         map[string]any
 	namespacePrefix                 string
 	initialSource                   []byte

connectors/s3vector/conn.go

@@ -22,6 +22,7 @@ import (
 )
 
 type conn struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	client         *s3vectors.Client
 	bucketName     *string
 	vectorKey      string

connectors/testconn/connector.go

@@ -30,6 +30,7 @@ import (
 )
 
 type conn struct {
+	adiomv1connect.UnimplementedConnectorServiceHandler
 	bootstrapPath string
 	updatesPath   string
 	loop          bool