Merge pull request #1158 from anatoly32322/fix-function-id

Fix gstack

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed `gstack` logic for parsing `ast.BlockStmt`
+
 ## v3.59.2
 * Added internal `gstack` codegen tool for filling `stack.FunctionID` with value from call stack
 

internal/cmd/gstack/main.go

@@ -93,6 +93,14 @@ func getCallExpressionsFromStmt(statement ast.Stmt) (listOfCallExpressions []*as
 		}
 	case *ast.ExprStmt:
 		listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromExpr(stmt.X)...)
+	case *ast.AssignStmt:
+		for _, rh := range stmt.Rhs {
+			listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromExpr(rh)...)
+		}
+	case *ast.ReturnStmt:
+		for _, result := range stmt.Results {
+			listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromExpr(result)...)
+		}
 	}
 	if body != nil {
 		listOfCallExpressions = append(
@@ -106,20 +114,7 @@ func getCallExpressionsFromStmt(statement ast.Stmt) (listOfCallExpressions []*as
 
 func getListOfCallExpressionsFromBlockStmt(block *ast.BlockStmt) (listOfCallExpressions []*ast.CallExpr) {
 	for _, statement := range block.List {
-		switch expr := statement.(type) {
-		case *ast.ExprStmt:
-			listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromExpr(expr.X)...)
-		case *ast.ReturnStmt:
-			for _, result := range expr.Results {
-				listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromExpr(result)...)
-			}
-		case *ast.AssignStmt:
-			for _, rh := range expr.Rhs {
-				listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromExpr(rh)...)
-			}
-		default:
-			listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromStmt(statement)...)
-		}
+		listOfCallExpressions = append(listOfCallExpressions, getCallExpressionsFromStmt(statement)...)
 	}
 
 	return listOfCallExpressions

internal/query/client.go

@@ -158,21 +158,15 @@ func doTx(
 }
 
 func (c *Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options.DoTxOption) (err error) {
-	var (
-		onDone = trace.QueryOnDoTx(c.config.Trace(), &ctx,
-			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/query.(*Client).DoTx"),
-		)
-		attempts int
-	)
-	defer func() {
-		onDone(attempts, err)
-	}()
-
 	select {
 	case <-c.done:
 		return xerrors.WithStackTrace(errClosedClient)
 	default:
-		attempts, err = doTx(ctx, c.pool, op, c.config.Trace(), opts...)
+		onDone := trace.QueryOnDoTx(c.config.Trace(), &ctx, stack.FunctionID(
+			"github.com/ydb-platform/ydb-go-sdk/3/internal/query.(*Client).DoTx"),
+		)
+		attempts, err := doTx(ctx, c.pool, op, c.config.Trace(), opts...)
+		onDone(attempts, err)
 
 		return err
 	}