Skip to content

migrate: handle new context methods #205

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ReneWerner87 merged 1 commit into from
Sep 3, 2025
Merged

migrate: handle new context methods #205

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 56 additions & 57 deletions cmd/internal/migrations/v3/context_methods.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,51 @@ func MigrateContextMethods(cmd *cobra.Command, cwd string, _, _ *semver.Version)
changed, err := internal.ChangeFileContent(cwd, func(content string) string {
orig := content

// UserContext() removed - Ctx implements context.Context
// old Context() returned fasthttp.RequestCtx
if !strings.Contains(orig, ".SetContext(") {
Comment on lines +23 to +24
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector Bot Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Skip Context rewrite only when SetContext detected

The AST rewrite that renames Context() calls to RequestCtx() only skips when the original source contains .SetContext(. Files that initially call c.UserContext() but never SetUserContext are migrated in two steps: the first pass turns UserContext() into Context() but introduces no SetContext, so a second run enters this block and rewrites that newly produced Context() to RequestCtx(). The return type then changes from context.Context to *fasthttp.RequestCtx, causing type errors for code that re-runs the migrator on already migrated files that only used UserContext().

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

@ReneWerner87 ReneWerner87 Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is it not working, we have a unittest for this case and it already test this

fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "", content, parser.ParseComments)
if err == nil {
modified := false
baseIdent := func(expr ast.Expr) *ast.Ident {
for {
switch e := expr.(type) {
case *ast.Ident:
return e
case *ast.SelectorExpr:
expr = e.X
case *ast.CallExpr:
expr = e.Fun
default:
return nil
}
}
}
ast.Inspect(file, func(n ast.Node) bool {
call, ok := n.(*ast.CallExpr)
if !ok {
return true
}
sel, ok := call.Fun.(*ast.SelectorExpr)
if !ok || sel.Sel.Name != "Context" || len(call.Args) != 0 {
return true
}
if ident := baseIdent(sel.X); ident != nil && isFiberCtx(orig, ident.Name) {
sel.Sel.Name = "RequestCtx"
modified = true
}
return true
})
if modified {
var buf bytes.Buffer
if err := format.Node(&buf, fset, file); err == nil {
content = buf.String()
}
}
}
}

Comment on lines +23 to +66
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Guard can silently miss/produce invalid migrations in mixed files; handle SetUserContext args and chained Context() regardless of the guard.

If a file already contains SetContext(, the guarded pass skips all Context()→RequestCtx() rewrites. That can leave:

  • SetUserContext(c.Context()) → SetContext(c.Context()) (wrong type)
  • Chains like c.Status(...).Context().SetBodyStreamWriter(...) unchanged (compile error)

Add a targeted, always-on AST pass that:

  • Renames SetUserContext → SetContext for fiber Ctx receivers
  • Rewrites its argument from .Context() → .RequestCtx()
  • Optionally rewrites chained .Context(). → .RequestCtx().

Apply within-file call to a helper (see diff in lines 81-93) and add the helper below:

// place near the bottom of this file
func fixSetUserContextCalls(orig, content string) string {
	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, "", content, parser.ParseComments)
	if err != nil {
		return content
	}
	modified := false

	baseIdent := func(expr ast.Expr) *ast.Ident {
		for {
			switch e := expr.(type) {
			case *ast.Ident:
				return e
			case *ast.SelectorExpr:
				expr = e.X
			case *ast.CallExpr:
				expr = e.Fun
			case *ast.IndexExpr, *ast.IndexListExpr:
				expr = e.(interface{ GetX() ast.Expr }).GetX()
			case *ast.ParenExpr:
				expr = e.X
			case *ast.TypeAssertExpr:
				expr = e.X
			case *ast.StarExpr:
				expr = e.X
			default:
				return nil
			}
		}
	}

	ast.Inspect(file, func(n ast.Node) bool {
		ce, ok := n.(*ast.CallExpr)
		if !ok {
			return true
		}
		sel, ok := ce.Fun.(*ast.SelectorExpr)
		if !ok {
			return true
		}
		recv := baseIdent(sel.X)
		if recv == nil || !isFiberCtx(orig, recv.Name) {
			return true
		}

		// Normalize SetUserContext → SetContext
		if sel.Sel.Name == "SetUserContext" {
			sel.Sel.Name = "SetContext"
			modified = true
		}

		// If first arg is <recv>.Context(), make it <recv>.RequestCtx()
		if sel.Sel.Name == "SetContext" && len(ce.Args) >= 1 {
			if argCall, ok := ce.Args[0].(*ast.CallExpr); ok {
				if argSel, ok := argCall.Fun.(*ast.SelectorExpr); ok && argSel.Sel.Name == "Context" {
					if aRecv := baseIdent(argSel.X); aRecv != nil && aRecv.Name == recv.Name {
						argSel.Sel.Name = "RequestCtx"
						modified = true
					}
				}
			}
		}
		return true
	})

	if !modified {
		return content
	}
	var buf bytes.Buffer
	if err := format.Node(&buf, fset, file); err != nil {
		return content
	}
	return buf.String()
}

// UserContext() replaced by Context()
reUserCtx := regexp.MustCompile(`(\w+)\.UserContext\(\)`)
content = reUserCtx.ReplaceAllStringFunc(content, func(match string) string {
parts := reUserCtx.FindStringSubmatch(match)
Expand All @@ -29,68 +73,23 @@ func MigrateContextMethods(cmd *cobra.Command, cwd string, _, _ *semver.Version)
}
ident := parts[1]
if isFiberCtx(orig, ident) {
return ident
return ident + ".Context()"
}
return match
})

// old Context() returned fasthttp.RequestCtx
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "", content, parser.ParseComments)
if err == nil {
modified := false
baseIdent := func(expr ast.Expr) *ast.Ident {
for {
switch e := expr.(type) {
case *ast.Ident:
return e
case *ast.SelectorExpr:
expr = e.X
case *ast.CallExpr:
expr = e.Fun
default:
return nil
}
}
}
ast.Inspect(file, func(n ast.Node) bool {
call, ok := n.(*ast.CallExpr)
if !ok {
return true
}
sel, ok := call.Fun.(*ast.SelectorExpr)
if !ok || sel.Sel.Name != "Context" || len(call.Args) != 0 {
return true
}
if ident := baseIdent(sel.X); ident != nil && isFiberCtx(orig, ident.Name) {
sel.Sel.Name = "RequestCtx"
modified = true
}
return true
})
if modified {
var buf bytes.Buffer
if err := format.Node(&buf, fset, file); err == nil {
content = buf.String()
}
}
}

// SetUserContext removed - comment out the call
reSetUserCtx := regexp.MustCompile(`(?m)^(\s*)(.*?\b(\w+)\.SetUserContext\([^\n]*\).*)$`)
content = reSetUserCtx.ReplaceAllStringFunc(content, func(line string) string {
if strings.Contains(line, "TODO: SetUserContext was removed") {
return line
}
parts := reSetUserCtx.FindStringSubmatch(line)
if len(parts) != 4 {
return line
// SetUserContext renamed to SetContext
reSetUserCtx := regexp.MustCompile(`(\w+)\.SetUserContext\(`)
content = reSetUserCtx.ReplaceAllStringFunc(content, func(match string) string {
parts := reSetUserCtx.FindStringSubmatch(match)
if len(parts) != 2 {
return match
}
ident := parts[3]
if !isFiberCtx(orig, ident) {
return line
ident := parts[1]
if isFiberCtx(orig, ident) {
return ident + ".SetContext("
}
return fmt.Sprintf("%s// TODO: SetUserContext was removed, please migrate manually: %s", parts[1], parts[2])
return match
})
Comment on lines +81 to 93
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Regex-only SetUserContext rename can break types in mixed files; replace with AST helper.

Tie the rename to argument fixups even when SetContext already exists.

Apply:

-		// SetUserContext renamed to SetContext
-		reSetUserCtx := regexp.MustCompile(`(\w+)\.SetUserContext\(`)
-		content = reSetUserCtx.ReplaceAllStringFunc(content, func(match string) string {
-			parts := reSetUserCtx.FindStringSubmatch(match)
-			if len(parts) != 2 {
-				return match
-			}
-			ident := parts[1]
-			if isFiberCtx(orig, ident) {
-				return ident + ".SetContext("
-			}
-			return match
-		})
+		// SetUserContext → SetContext + fix arg: <recv>.Context() → <recv>.RequestCtx()
+		content = fixSetUserContextCalls(orig, content)

Add the helper per earlier comment.

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents 
In cmd/internal/migrations/v3/context_methods.go around lines 81 to 93, the
current regex-only replacement for renaming SetUserContext -> SetContext can
break types in mixed files and misses tying the rename to argument fixups;
replace this block with an AST-based helper call (the helper referenced in the
earlier comment) that walks the Go AST, finds selector expressions calling
.SetUserContext on idents, verifies the receiver is a fiber.Context via
isFiberCtx, replaces the selector to .SetContext, and performs the necessary
argument adjustments (including adding/removing context arg) even when a
.SetContext already exists; invoke that helper here instead of using
regexp.ReplaceAllStringFunc to ensure type-safe, idempotent transformations.


return content
Expand Down
52 changes: 46 additions & 6 deletions cmd/internal/migrations/v3/context_methods_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,11 @@ func handler(c fiber.Ctx) error {
require.NoError(t, v3.MigrateContextMethods(cmd, dir, nil, nil))

content := readFile(t, file)
assert.Contains(t, content, ".RequestCtx()")
assert.NotContains(t, content, ".Context()")
assert.Contains(t, content, `// TODO: SetUserContext was removed, please migrate manually: c.SetUserContext(ctx)`)
assert.Contains(t, content, "uc := c")
assert.Contains(t, content, "ctx := c.RequestCtx()")
assert.Contains(t, content, "uc := c.Context()")
assert.Contains(t, content, "c.SetContext(ctx)")
assert.NotContains(t, content, ".UserContext()")
assert.NotContains(t, content, "SetUserContext")
assert.Contains(t, buf.String(), "Migrating context methods")
}

Expand All @@ -62,8 +63,9 @@ func handler(ctx fiber.Ctx) error {
require.NoError(t, v3.MigrateContextMethods(cmd, dir, nil, nil))

content := readFile(t, file)
assert.Contains(t, content, `// TODO: SetUserContext was removed, please migrate manually: res := ctx.SetUserContext(ctx.RequestCtx())`)
assert.Contains(t, content, `res := ctx.SetContext(ctx.RequestCtx())`)
assert.NotContains(t, content, `.UserContext()`)
assert.NotContains(t, content, "SetUserContext")
assert.Contains(t, buf.String(), "Migrating context methods")
}

Expand All @@ -88,7 +90,45 @@ func handler(c fiber.Ctx) error {
require.NoError(t, v3.MigrateContextMethods(cmd, dir, nil, nil))
second := readFile(t, file)
assert.Equal(t, first, second)
assert.Equal(t, 1, strings.Count(second, "TODO: SetUserContext was removed"))
assert.Equal(t, 1, strings.Count(second, "SetContext("))
Comment on lines 91 to +93
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The assertion assert.Equal(t, 1, strings.Count(second, "TODO: SetUserContext was removed")) is no longer valid since SetUserContext is now replaced with SetContext. This test should be updated to reflect the new code.

I suggest changing the assertion to check for the correct number of SetContext( calls after the migration. Since the goal is to ensure idempotency, the count should remain consistent after multiple runs.

Consider updating the assertion to check for the correct number of SetContext( calls after the migration. This will ensure that the test accurately reflects the intended behavior of the migration.

I recommend updating the assertion to check for the correct number of SetContext( calls after the migration. This will ensure that the test accurately reflects the intended behavior of the migration.

Suggested change
second := readFile(t, file)
assert.Equal(t, first, second)
assert.Equal(t, 1, strings.Count(second, "TODO: SetUserContext was removed"))
assert.Equal(t, 1, strings.Count(second, "SetContext("))
assert.Equal(t, 1, strings.Count(second, "SetContext("))

}

func Test_MigrateContextMethods_MultipleRuns(t *testing.T) {
t.Parallel()

dir, err := os.MkdirTemp("", "mcmtestmulti2")
require.NoError(t, err)
defer func() { require.NoError(t, os.RemoveAll(dir)) }()

file := writeTempFile(t, dir, `package main
import "github.com/gofiber/fiber/v2"
func handler(c fiber.Ctx) error {
ctx := c.Context()
uc := c.UserContext()
c.SetUserContext(ctx)
c.SetUserContext(c.Context())
_ = uc
return nil
}`)

var buf bytes.Buffer
cmd := newCmd(&buf)
require.NoError(t, v3.MigrateContextMethods(cmd, dir, nil, nil))
first := readFile(t, file)

require.Contains(t, first, "ctx := c.RequestCtx()")
require.Contains(t, first, "uc := c.Context()")
require.Contains(t, first, "c.SetContext(ctx)")
require.Contains(t, first, "c.SetContext(c.RequestCtx())")
require.NotContains(t, first, ".UserContext()")
require.NotContains(t, first, "SetUserContext")

require.NoError(t, v3.MigrateContextMethods(cmd, dir, nil, nil))
second := readFile(t, file)
assert.Equal(t, first, second)
assert.Equal(t, 1, strings.Count(second, "uc := c.Context()"))
assert.Equal(t, 2, strings.Count(second, ".RequestCtx()"))
assert.Equal(t, 2, strings.Count(second, "SetContext("))
}
Comment on lines +96 to 132
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add a mixed-file test to catch the guard edge case.

This reproduces SetUserContext→SetContext with an existing SetContext elsewhere; ensures the argument becomes RequestCtx.

Apply:

@@
 func Test_MigrateContextMethods_MultipleRuns(t *testing.T) {
@@
 }
 
+func Test_MigrateContextMethods_MixedFileWithExistingSetContext(t *testing.T) {
+	t.Parallel()
+	dir, err := os.MkdirTemp("", "mcmtest-mixed")
+	require.NoError(t, err)
+	defer func() { require.NoError(t, os.RemoveAll(dir)) }()
+
+	file := writeTempFile(t, dir, `package main
+import "github.com/gofiber/fiber/v2"
+func handler(c fiber.Ctx) error {
+    // already migrated in this file:
+    c.SetContext(c.RequestCtx())
+    // still old call:
+    _ = c.SetUserContext(c.Context())
+    // chain that must stay valid:
+    c.Status(200).Context().SetBodyStreamWriter(nil)
+    return nil
+}`)
+
+	var buf bytes.Buffer
+	cmd := newCmd(&buf)
+	require.NoError(t, v3.MigrateContextMethods(cmd, dir, nil, nil))
+	content := readFile(t, file)
+	assert.Contains(t, content, "c.SetContext(c.RequestCtx())")                 // unchanged existing
+	assert.NotContains(t, content, "SetUserContext(")                           // renamed
+	assert.Contains(t, content, "SetContext(c.RequestCtx())")                   // fixed arg
+	assert.NotContains(t, content, ".Context().SetBodyStreamWriter")            // chained fix
+	assert.Contains(t, content, ".RequestCtx().SetBodyStreamWriter")            // chained fix
+}
🤖 Prompt for AI Agents 
In cmd/internal/migrations/v3/context_methods_test.go around lines 96–132, add a
mixed-file test that reproduces the guard edge case where an existing SetContext
call in the file prevents proper transformation of SetUserContext → SetContext;
the test should create a temp file containing both SetUserContext calls and
existing SetContext usages, run v3.MigrateContextMethods twice, then assert that
SetUserContext calls were converted to SetContext with RequestCtx() arguments,
that existing SetContext calls remain unchanged, and that counts of "uc :=
c.Context()", ".RequestCtx()", and "SetContext(" match the expected numbers to
ensure idempotence.


func Test_MigrateContextMethods_SkipNonFiber(t *testing.T) {
Expand Down
Loading