Merge pull request #24 from chenmingyong0423/feature/build

implements the lookup method for build the $lookup stage
chenmingyong0423 · Apr 6, 2024 · 82a7d52 · 82a7d52
2 parents 4abfe12 + 606b88d
commit 82a7d52
Show file tree

Hide file tree

Showing 3 changed files with 122 additions and 0 deletions.
diff --git a/builder/aggregation/aggregation_stage_builder.go b/builder/aggregation/aggregation_stage_builder.go
@@ -140,6 +140,23 @@ func (b *StageBuilder) Count(countName string) *StageBuilder {
 	return b
 }
 
+func (b *StageBuilder) Lookup(from, as string, opt *LookUpOptions) *StageBuilder {
+	d := bson.D{bson.E{Key: "from", Value: from}}
+	if opt.LocalField != "" && opt.ForeignField != "" {
+		d = append(d, bson.E{Key: "localField", Value: opt.LocalField})
+		d = append(d, bson.E{Key: "foreignField", Value: opt.ForeignField})
+	}
+	if len(opt.Let) > 0 {
+		d = append(d, bson.E{Key: "let", Value: opt.Let})
+	}
+	if len(opt.Pipeline) > 0 {
+		d = append(d, bson.E{Key: "pipeline", Value: opt.Pipeline})
+	}
+	d = append(d, bson.E{Key: "as", Value: as})
+	b.pipeline = append(b.pipeline, bson.D{bson.E{Key: StageLookUp, Value: d}})
+	return b
+}
+
 func (b *StageBuilder) Build() mongo.Pipeline {
 	return b.pipeline
 }
diff --git a/builder/aggregation/aggregation_stage_builder_test.go b/builder/aggregation/aggregation_stage_builder_test.go
@@ -654,3 +654,77 @@ func TestStageBuilder_SortByCount(t *testing.T) {
 func TestStageBuilder_Count(t *testing.T) {
 	assert.Equal(t, mongo.Pipeline{bson.D{bson.E{Key: "$count", Value: "passing_scores"}}}, StageBsonBuilder().Count("passing_scores").Build())
 }
+
+func TestStageBuilder_Lookup(t *testing.T) {
+	testCases := []struct {
+		name string
+		from string
+		as   string
+		opt  *LookUpOptions
+
+		want mongo.Pipeline
+	}{
+		{
+			name: "basic",
+			from: "orders",
+			opt: &LookUpOptions{
+				LocalField:   "_id",
+				ForeignField: "userId",
+				Let:          nil,
+				Pipeline:     nil,
+			},
+			as: "userOrders",
+			want: mongo.Pipeline{
+				{
+					bson.E{Key: "$lookup", Value: bson.D{
+						bson.E{Key: "from", Value: "orders"},
+						bson.E{Key: "localField", Value: "_id"},
+						bson.E{Key: "foreignField", Value: "userId"},
+						bson.E{Key: "as", Value: "userOrders"},
+					}},
+				},
+			},
+		},
+		{
+			name: "advanced case",
+			from: "orders",
+			opt: &LookUpOptions{
+				LocalField:   "",
+				ForeignField: "",
+				Let:          bson.D{bson.E{Key: "userId", Value: "$_id"}},
+				Pipeline: mongo.Pipeline{
+					{
+						bson.E{Key: "$match", Value: bson.D{bson.E{Key: "$expr", Value: bson.D{bson.E{Key: "$and", Value: []any{
+							bson.D{bson.E{Key: "$eq", Value: []any{"$userId", "$$userId"}}},
+							bson.D{bson.E{Key: "$gt", Value: []any{"$totalAmount", 100}}},
+						}}}}}},
+					},
+				},
+			},
+			as: "largeOrders",
+			want: mongo.Pipeline{
+				{
+					bson.E{Key: "$lookup", Value: bson.D{
+						bson.E{Key: "from", Value: "orders"},
+						bson.E{Key: "let", Value: bson.D{bson.E{Key: "userId", Value: "$_id"}}},
+						bson.E{Key: "pipeline", Value: mongo.Pipeline{
+							{
+								bson.E{Key: "$match", Value: bson.D{bson.E{Key: "$expr", Value: bson.D{bson.E{Key: "$and", Value: []any{
+									bson.D{bson.E{Key: "$eq", Value: []any{"$userId", "$$userId"}}},
+									bson.D{bson.E{Key: "$gt", Value: []any{"$totalAmount", 100}}},
+								}}}}}},
+							},
+						}},
+						bson.E{Key: "as", Value: "largeOrders"},
+					},
+					},
+				},
+			},
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.want, StageBsonBuilder().Lookup(tc.from, tc.as, tc.opt).Build())
+		})
+	}
+}
diff --git a/builder/aggregation/types.go b/builder/aggregation/types.go
@@ -0,0 +1,31 @@
+// Copyright 2024 chenmingyong0423
+
+// Licensed 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 aggregation
+
+import (
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+const (
+	StageLookUp = "$lookup"
+)
+
+type LookUpOptions struct {
+	LocalField   string
+	ForeignField string
+	Let          bson.D
+	Pipeline     mongo.Pipeline
+}