@@ -123,28 +123,28 @@ func (v *visitor) visitMethodChain(ctx mongodb.IMethodChainContext) {
123123 if ! ok {
124124 return
125125 }
126- for _ , methodCall := range mc .AllMethodCall () {
127- v .visitMethodCall (methodCall )
126+
127+ if mc .CollectionMethodCall () != nil {
128+ v .visitCollectionMethodCall (mc .CollectionMethodCall ())
129+ if v .err != nil {
130+ return
131+ }
132+ }
133+
134+ for _ , cursorCall := range mc .AllCursorMethodCall () {
135+ v .visitCursorMethodCall (cursorCall )
128136 if v .err != nil {
129137 return
130138 }
131139 }
132140}
133141
134- func (v * visitor ) visitMethodCall (ctx mongodb.IMethodCallContext ) {
135- mc , ok := ctx .(* mongodb.MethodCallContext )
142+ func (v * visitor ) visitCollectionMethodCall (ctx mongodb.ICollectionMethodCallContext ) {
143+ mc , ok := ctx .(* mongodb.CollectionMethodCallContext )
136144 if ! ok {
137145 return
138146 }
139147
140- // Determine method context for registry lookup
141- getMethodContext := func () string {
142- if v .operation .OpType == types .OpFind || v .operation .OpType == types .OpFindOne {
143- return "cursor"
144- }
145- return "collection"
146- }
147-
148148 switch {
149149 // Supported read operations
150150 case mc .FindMethod () != nil :
@@ -168,32 +168,13 @@ func (v *visitor) visitMethodCall(ctx mongodb.IMethodCallContext) {
168168 case mc .GetIndexesMethod () != nil :
169169 v .operation .OpType = types .OpGetIndexes
170170
171- // Supported cursor modifiers
172- case mc .SortMethod () != nil :
173- v .extractSort (mc .SortMethod ())
174- case mc .LimitMethod () != nil :
175- v .extractLimit (mc .LimitMethod ())
176- case mc .SkipMethod () != nil :
177- v .extractSkip (mc .SkipMethod ())
178- case mc .ProjectionMethod () != nil :
179- v .extractProjection (mc .ProjectionMethod ())
180- case mc .HintMethod () != nil :
181- v .extractHint (mc .HintMethod ())
182- case mc .MaxMethod () != nil :
183- v .extractMax (mc .MaxMethod ())
184- case mc .MinMethod () != nil :
185- v .extractMin (mc .MinMethod ())
186-
187- // Supported M2 write operations
171+ // Supported write operations
188172 case mc .InsertOneMethod () != nil :
189173 v .operation .OpType = types .OpInsertOne
190174 v .extractInsertOneArgs (mc .InsertOneMethod ())
191-
192175 case mc .InsertManyMethod () != nil :
193176 v .operation .OpType = types .OpInsertMany
194177 v .extractInsertManyArgs (mc .InsertManyMethod ())
195-
196- // Supported M2 write operations - updateOne
197178 case mc .UpdateOneMethod () != nil :
198179 v .operation .OpType = types .OpUpdateOne
199180 v .extractUpdateOneArgs (mc .UpdateOneMethod ())
@@ -219,27 +200,27 @@ func (v *visitor) visitMethodCall(ctx mongodb.IMethodCallContext) {
219200 v .operation .OpType = types .OpFindOneAndDelete
220201 v .extractFindOneAndDeleteArgs (mc .FindOneAndDeleteMethod ())
221202
222- // Supported M3 index operations
203+ // Supported index operations
223204 case mc .CreateIndexMethod () != nil :
224205 v .operation .OpType = types .OpCreateIndex
225206 v .extractCreateIndexArgs (mc .CreateIndexMethod ())
226207 case mc .CreateIndexesMethod () != nil :
227- v .handleUnsupportedMethod ("collection" , "createIndexes" ) // Lower ROI, keep as planned
208+ v .handleUnsupportedMethod ("collection" , "createIndexes" )
228209 case mc .DropIndexMethod () != nil :
229210 v .operation .OpType = types .OpDropIndex
230211 v .extractDropIndexArgs (mc .DropIndexMethod ())
231212 case mc .DropIndexesMethod () != nil :
232213 v .operation .OpType = types .OpDropIndexes
233214 v .extractDropIndexesArgs (mc .DropIndexesMethod ())
234215
235- // Supported M3 collection management
216+ // Supported collection management
236217 case mc .DropMethod () != nil :
237218 v .operation .OpType = types .OpDrop
238219 case mc .RenameCollectionMethod () != nil :
239220 v .operation .OpType = types .OpRenameCollection
240221 v .extractRenameCollectionArgs (mc .RenameCollectionMethod ())
241222
242- // Planned M3 stats operations - return PlannedOperationError for fallback
223+ // Planned stats operations
243224 case mc .StatsMethod () != nil :
244225 v .handleUnsupportedMethod ("collection" , "stats" )
245226 case mc .StorageSizeMethod () != nil :
@@ -257,34 +238,45 @@ func (v *visitor) visitMethodCall(ctx mongodb.IMethodCallContext) {
257238 case mc .LatencyStatsMethod () != nil :
258239 v .handleUnsupportedMethod ("collection" , "latencyStats" )
259240
260- // Generic method fallback - all methods going through genericMethod are unsupported
261- case mc .GenericMethod () != nil :
262- gmCtx , ok := mc .GenericMethod ().(* mongodb.GenericMethodContext )
263- if ! ok {
264- return
241+ default :
242+ methodName := extractMethodNameFromText (mc .GetText ())
243+ if methodName != "" {
244+ v .handleUnsupportedMethod ("collection" , methodName )
265245 }
266- methodName := gmCtx . Identifier (). GetText ()
267- v . handleUnsupportedMethod ( getMethodContext (), methodName )
246+ }
247+ }
268248
269- // Default: all other methods not explicitly handled
270- // These go to handleUnsupportedMethod which returns UnsupportedOperationError
271- // since they're not in the planned registry
249+ func (v * visitor ) visitCursorMethodCall (ctx mongodb.ICursorMethodCallContext ) {
250+ mc , ok := ctx .(* mongodb.CursorMethodCallContext )
251+ if ! ok {
252+ return
253+ }
254+
255+ switch {
256+ case mc .SortMethod () != nil :
257+ v .extractSort (mc .SortMethod ())
258+ case mc .LimitMethod () != nil :
259+ v .extractLimit (mc .LimitMethod ())
260+ case mc .SkipMethod () != nil :
261+ v .extractSkip (mc .SkipMethod ())
262+ case mc .ProjectionMethod () != nil :
263+ v .extractProjection (mc .ProjectionMethod ())
264+ case mc .HintMethod () != nil :
265+ v .extractHint (mc .HintMethod ())
266+ case mc .MaxMethod () != nil :
267+ v .extractMax (mc .MaxMethod ())
268+ case mc .MinMethod () != nil :
269+ v .extractMin (mc .MinMethod ())
272270 default :
273- // Extract method name from the parse tree for error message
274- methodName := v .extractMethodName (mc )
271+ methodName := extractMethodNameFromText (mc .GetText ())
275272 if methodName != "" {
276- v .handleUnsupportedMethod (getMethodContext () , methodName )
273+ v .handleUnsupportedMethod ("cursor" , methodName )
277274 }
278275 }
279276}
280277
281- // extractMethodName extracts the method name from a MethodCallContext for error messages.
282- func (v * visitor ) extractMethodName (mc * mongodb.MethodCallContext ) string {
283- // Try to get method name from various method contexts
284- // The parser creates specific method contexts for known methods
285- // For unknown methods, they go through GenericMethod which is handled separately
286- text := mc .GetText ()
287- // Extract method name before the opening parenthesis
278+ // extractMethodNameFromText extracts the method name from a parse tree text before the opening parenthesis.
279+ func extractMethodNameFromText (text string ) string {
288280 if idx := strings .Index (text , "(" ); idx > 0 {
289281 return text [:idx ]
290282 }
0 commit comments