chbatchclose: add support for closure inside defer

Author: cyrilou242

passes/chbatchclose/chbatchclose.go

@@ -146,28 +146,76 @@ func (a *analyzer) handleAssign(pass *analysis.Pass, assign *ast.AssignStmt, usa
 	}
 }
 
-// handleDefer checks if a defer statement calls Close() or Abort() on a tracked Batch variable.
+// handleDefer checks if a defer statement calls Close() on a tracked Batch variable.
+// Two shapes are supported:
+//   - direct selector call:           defer batch.Close()
+//   - immediately-invoked closure:    defer func() { ... batch.Close() ... }()
+//     (no params; nested FuncLits inside the closure body are not descended into)
 func handleDefer(deferStmt *ast.DeferStmt, usages map[string]*batchUsage) {
 	call := deferStmt.Call
-	sel, ok := call.Fun.(*ast.SelectorExpr)
-	if !ok {
-		return
-	}
-	varName := util.IdentName(sel.X)
-	if varName == "" {
-		return
-	}
-	u, exists := usages[varName]
-	if !exists {
-		return
-	}
 
-	switch sel.Sel.Name {
-	case "Close":
-		u.deferredClose = true
+	switch fun := call.Fun.(type) {
+	case *ast.SelectorExpr:
+		varName := util.IdentName(fun.X)
+		if varName == "" {
+			return
+		}
+		u, exists := usages[varName]
+		if !exists {
+			return
+		}
+		if fun.Sel.Name == "Close" {
+			u.deferredClose = true
+		}
+	case *ast.FuncLit:
+		// only no-arg closures are supported for the moment
+		if fun.Type.Params != nil && len(fun.Type.Params.List) > 0 {
+			return
+		}
+		if fun.Body == nil {
+			return
+		}
+		handleDeferredClosure(fun.Body, usages)
 	}
 }
 
+// handleDeferredClosure walks the body of a deferred closure looking for
+// `<tracked>.Close()` calls. It does not descend into nested FuncLits, so
+// `defer func() { go func() { batch.Close() }() }()` and similar are not
+// treated as a valid defensive close.
+func handleDeferredClosure(body *ast.BlockStmt, usages map[string]*batchUsage) {
+	ast.Inspect(body, func(n ast.Node) bool {
+		if n == nil {
+			return false
+		}
+		// don't descend into nested closures (goroutines, callbacks, etc.)
+		if n != body {
+			if _, ok := n.(*ast.FuncLit); ok {
+				return false
+			}
+		}
+		call, ok := n.(*ast.CallExpr)
+		if !ok {
+			return true
+		}
+		sel, ok := call.Fun.(*ast.SelectorExpr)
+		if !ok {
+			return true
+		}
+		if sel.Sel.Name != "Close" {
+			return true
+		}
+		varName := util.IdentName(sel.X)
+		if varName == "" {
+			return true
+		}
+		if u, exists := usages[varName]; exists {
+			u.deferredClose = true
+		}
+		return true
+	})
+}
+
 // handleReturn checks if any return value is a tracked Batch variable.
 func handleReturn(ret *ast.ReturnStmt, usages map[string]*batchUsage) {
 	for _, result := range ret.Results {

passes/chbatchclose/testdata/src/testcases/testcases.go

@@ -176,16 +176,75 @@ func reassignValidValid() {
 	_ = batch.Send()
 }
 
-// closures are not supported
-// while the code below is in theory correct, it is very likely to be a bad pattern and should be flagged by the linter
-func deferCloseIsInClosure() {
-	batch, err := conn.PrepareBatch(ctx, "INSERT INTO t") //want `clickhouse Batch batch must be closed defensively with defer batch\.Close\(\) after successful instantiation`
+// valid: defer with a closure that calls batch.Close() directly.
+// This is the IDE/errcheck-friendly pattern (allows wrapping the Close error, see test below).
+func validDeferCloseInClosure() {
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO t")
 	if err != nil {
 		return
 	}
 	defer func() { batch.Close() }()
 }
 
+// valid: defer with closure handling the Close() error (real-world pattern of above).
+func validDeferCloseInClosureWithErrCheck() {
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO t")
+	if err != nil {
+		return
+	}
+	defer func() {
+		if err = batch.Close(); err != nil {
+			_ = err // log error
+		}
+	}()
+}
+
+// invalid: defer with a closure that takes the batch as an argument.
+// this is correct in theory, but not supported for the moment (false positive)
+// Supporting this would require mapping closure params back to outer args. Assuming this pattern is not used for the moment.
+func invalidDeferCloseInClosureWithArg() {
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO t") //want `clickhouse Batch batch must be closed defensively with defer batch\.Close\(\) after successful instantiation`
+	if err != nil {
+		return
+	}
+	defer func(b driver.Batch) { _ = b.Close() }(batch)
+}
+
+// invalid: Close is called from a goroutine nested inside the deferred closure.
+// The defer itself does not synchronously close the batch.
+// in theory correct-ish but for the moment I think it's not a good pattern
+func invalidDeferCloseInNestedGoroutine() {
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO t") //want `clickhouse Batch batch must be closed defensively with defer batch\.Close\(\) after successful instantiation`
+	if err != nil {
+		return
+	}
+	defer func() {
+		go func() { _ = batch.Close() }()
+	}()
+}
+
+// invalid: 2 level of closure
+// correct in theory but not tracked for the moment - false positive
+func invalidDeferCloseInNestedCallback() {
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO t") //want `clickhouse Batch batch must be closed defensively with defer batch\.Close\(\) after successful instantiation`
+	if err != nil {
+		return
+	}
+	cb := func(fn func()) { fn() }
+	defer func() {
+		cb(func() { _ = batch.Close() })
+	}()
+}
+
+// invalid: deferred closure calls only Abort() (or other methods), not Close().
+func invalidDeferAbortInClosure() {
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO t") //want `clickhouse Batch batch must be closed defensively with defer batch\.Close\(\) after successful instantiation`
+	if err != nil {
+		return
+	}
+	defer func() { _ = batch.Abort() }()
+}
+
 // the code below is in theory correct as all error cases are handled and result in a batch.Abort().
 // we still mark this as an error as it's not defensive. A defer batch.Close() would not change the code correctness and is easy to add.
 func invalidNotDefensive() error {