⚡️ Speed up method MigrationValidator._check_upgrade_operations by 1,769% in PR #10519 (support-n-1-database-migration-guideline)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #10519

If you approve this dependent PR, these changes will be merged into the original PR branch support-n-1-database-migration-guideline.

This PR will be automatically closed if the original PR is merged.

---

📄 1,769% (17.69x) speedup for MigrationValidator._check_upgrade_operations in src/backend/base/langflow/alembic/migration_validator.py

⏱️ Runtime : 366 milliseconds → 19.6 milliseconds (best of 50 runs)

📝 Explanation and details

The optimization achieves a remarkable 17.7x speedup by addressing two critical performance bottlenecks in the AST validation logic:

Key Optimization 1: Eliminated Redundant AST Operations Matching
The original code performed expensive _is_op_call() checks for every AST call node (up to 5 checks per node). The optimized version makes a single pass through ast.walk(), directly identifying and grouping op calls by type using inline attribute checks. This eliminates thousands of redundant function calls and string comparisons.

Key Optimization 2: Cached Existence Check Analysis
The most expensive bottleneck was _has_existence_check_nearby(), which performed a full ast.walk(func_node) for every add_column call (97.4% of original runtime). The optimization builds an existence check index once per function and reuses it via func_node._existence_if_index, converting O(n*m) complexity to O(n+m) where n=AST nodes, m=add_column calls.

Performance Impact by Test Case:

• Small migrations (single operations): 10-20x faster due to reduced overhead
• Large-scale migrations (100+ operations): Even greater speedup as the existence check caching pays massive dividends
• Mixed operation types: Benefits from both optimizations working together

Behavioral Preservation:
The optimizations maintain identical violation detection logic and error messages. The existence check uses the same heuristics ("column", "inspector", "not in" keywords) with equivalent semantics.

This optimization is particularly valuable for migration validation in CI/CD pipelines where many migration files may be processed, as the performance gains compound significantly with larger AST trees and more database operations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 82 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import ast
from enum import Enum

imports

import pytest
from langflow.alembic.migration_validator import MigrationValidator

--- Supporting classes and enums for the tests ---

class MigrationPhase(Enum):
EXPAND = "EXPAND"
CONTRACT = "CONTRACT"

--- Unit tests for MigrationValidator._check_upgrade_operations ---

@pytest.fixture
def validator():
"""Fixture for MigrationValidator instance."""
return MigrationValidator()

def parse_func(src: str) -> ast.FunctionDef:
"""Helper to parse a function source string to AST FunctionDef."""
mod = ast.parse(src)
for node in mod.body:
if isinstance(node, ast.FunctionDef):
return node
raise ValueError("No function definition found")

------------------- BASIC TEST CASES -------------------

def test_add_nullable_column_with_existence_check_expand(validator):
# Should NOT raise any violations (nullable=True and existence check)
src = """
def upgrade():
if 'new_column' not in inspector.get_columns('table'):
op.add_column('table', sa.Column('new_column', sa.String(), nullable=True))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def upgrade():
op.add_column('table', sa.Column('new_column', sa.String(), nullable=False))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output
codes = {v.code for v in violations}

def test_add_non_nullable_column_with_default_and_existence_check_expand(validator):
# Should NOT raise any violations (server_default and existence check)
src = """
def upgrade():
if 'new_column' not in inspector.get_columns('table'):
op.add_column('table', sa.Column('new_column', sa.String(), nullable=False, server_default='abc'))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_direct_rename_table_expand(validator):
# Should raise DIRECT_RENAME
src = """
def upgrade():
op.rename_table('old_table', 'new_table')
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_direct_rename_column_expand(validator):
# Should raise DIRECT_RENAME
src = """
def upgrade():
op.rename_column('table', 'old_col', 'new_col')
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_alter_column_type_change_expand(validator):
# Should raise DIRECT_TYPE_CHANGE
src = """
def upgrade():
op.alter_column('table', 'col', type_=sa.Integer())
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_alter_column_nullable_false_expand(validator):
# Should raise BREAKING_ADD_COLUMN
src = """
def upgrade():
op.alter_column('table', 'col', nullable=False)
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_drop_column_expand(validator):
# Should raise IMMEDIATE_DROP
src = """
def upgrade():
op.drop_column('table', 'col')
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_drop_column_contract(validator):
# Should NOT raise any violations in CONTRACT phase
src = """
def upgrade():
op.drop_column('table', 'col')
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def upgrade():
op.add_column('table', sa.Column('new_column', sa.String(), nullable=True))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output
codes = {v.code for v in violations}

------------------- EDGE TEST CASES -------------------

def test_add_column_with_nested_existence_check(validator):
# Existence check is nested inside another if
src = """
def upgrade():
if True:
if 'new_column' not in inspector.get_columns('table'):
op.add_column('table', sa.Column('new_column', sa.String(), nullable=True))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_add_column_with_existence_check_using_column_in_condition(validator):
# Existence check uses 'column' keyword in condition
src = """
def upgrade():
if 'new_column' not in columns:
op.add_column('table', sa.Column('new_column', sa.String(), nullable=True))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def upgrade():
op.add_column('table', sa.Column('new_column', sa.String()))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output
codes = {v.code for v in violations}

def upgrade():
op.alter_column('table', 'col', type_=sa.Integer(), nullable=False)
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output
codes = {v.code for v in violations}

def test_alter_column_type_contract(validator):
# Should NOT raise DIRECT_TYPE_CHANGE in CONTRACT phase
src = """
def upgrade():
op.alter_column('table', 'col', type_=sa.Integer())
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def test_alter_column_nullable_false_contract(validator):
# Should NOT raise violation in CONTRACT phase
src = """
def upgrade():
op.alter_column('table', 'col', nullable=False)
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def upgrade():
op.add_column('table', sa.Column('new_column', sa.String(), nullable=True))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output
codes = {v.code for v in violations}

def upgrade():
op.add_column('table', sa.Column('new_column', sa.String(), server_default='abc'))
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output
codes = {v.code for v in violations}

def upgrade():
op.add_column('table', sa.Column('new_column', sa.String(), nullable=False))
op.rename_column('table', 'old_col', 'new_col')
op.drop_column('table', 'col')
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output
codes = {v.code for v in violations}

def test_no_operations(validator):
# No operations in upgrade function
src = """
def upgrade():
pass
"""
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

------------------- LARGE SCALE TEST CASES -------------------

def test_many_add_columns_all_valid(validator):
# 100 columns, all nullable=True and with existence check
src_lines = [
"def upgrade():",
" inspector = None",
]
for i in range(100):
src_lines.append(
f" if 'col_{i}' not in inspector.get_columns('table'):"
)
src_lines.append(
f" op.add_column('table', sa.Column('col_{i}', sa.String(), nullable=True))"
)
src = "\n".join(src_lines)
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_many_alter_column_type_change_expand(validator):
# 100 alter_column with type_ change in EXPAND phase
src_lines = [
"def upgrade():",
]
for i in range(100):
src_lines.append(
f" op.alter_column('table', 'col_{i}', type_=sa.Integer())"
)
src = "\n".join(src_lines)
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_many_drop_column_contract(validator):
# 100 drop_column in CONTRACT phase (all valid)
src_lines = [
"def upgrade():",
]
for i in range(100):
src_lines.append(
f" op.drop_column('table', 'col_{i}')"
)
src = "\n".join(src_lines)
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def test_many_drop_column_expand(validator):
# 100 drop_column in EXPAND phase (all invalid)
src_lines = [
"def upgrade():",
]
for i in range(100):
src_lines.append(
f" op.drop_column('table', 'col_{i}')"
)
src = "\n".join(src_lines)
node = parse_func(src)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
import ast

imports

import pytest
from langflow.alembic.migration_validator import MigrationValidator

--- Dummy classes/enums for testing ---

class MigrationPhase:
"""Enum for migration phases."""
EXPAND = "EXPAND"
CONTRACT = "CONTRACT"
# For test convenience, allow .value property
def init(self, value):
self.value = value
def eq(self, other):
if isinstance(other, MigrationPhase):
return self.value == other.value
return self.value == other

MigrationPhase.EXPAND = MigrationPhase("EXPAND")
MigrationPhase.CONTRACT = MigrationPhase("CONTRACT")

--- Helper for parsing code snippets into AST FunctionDef nodes ---

def get_functiondef_from_code(code: str) -> ast.FunctionDef:
"""Parse code and return the first FunctionDef node."""
mod = ast.parse(code)
for node in mod.body:
if isinstance(node, ast.FunctionDef):
return node
raise ValueError("No FunctionDef found")

--- Unit tests ---

@pytest.fixture
def validator():
"""Fixture for MigrationValidator instance."""
return MigrationValidator()

1. Basic Test Cases

def test_add_column_nullable_true_with_existence_check_expand(validator):
# Should NOT raise any violations (nullable=True, existence check, EXPAND phase)
code = """
def upgrade():
if "new_column" not in inspector.get_columns("my_table"):
op.add_column("my_table", sa.Column("new_column", sa.String(), nullable=True))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_add_column_with_server_default_and_existence_check_expand(validator):
# Should NOT raise any violations (server_default, existence check, EXPAND phase)
code = """
def upgrade():
if "new_column" not in inspector.get_columns("my_table"):
op.add_column("my_table", sa.Column("new_column", sa.String(), server_default="foo"))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_add_column_missing_nullable_and_default_no_existence_check_expand(validator):
# Should raise both BREAKING_ADD_COLUMN and NO_EXISTENCE_CHECK
code = """
def upgrade():
op.add_column("my_table", sa.Column("new_column", sa.String()))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_add_column_in_contract_phase(validator):
# Should raise INVALID_PHASE_OPERATION (even if nullable=True and existence check)
code = """
def upgrade():
if "new_column" not in inspector.get_columns("my_table"):
op.add_column("my_table", sa.Column("new_column", sa.String(), nullable=True))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def test_alter_column_type_change_expand(validator):
# Should raise DIRECT_TYPE_CHANGE in EXPAND phase
code = """
def upgrade():
op.alter_column("my_table", "col", type_=sa.Integer())
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_alter_column_type_change_contract(validator):
# Should NOT raise DIRECT_TYPE_CHANGE in CONTRACT phase
code = """
def upgrade():
op.alter_column("my_table", "col", type_=sa.Integer())
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def test_alter_column_nullable_to_false_expand(validator):
# Should raise BREAKING_ADD_COLUMN in EXPAND phase
code = """
def upgrade():
op.alter_column("my_table", "col", nullable=False)
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_alter_column_nullable_to_false_contract(validator):
# Should NOT raise BREAKING_ADD_COLUMN in CONTRACT phase
code = """
def upgrade():
op.alter_column("my_table", "col", nullable=False)
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def test_drop_column_in_expand_phase(validator):
# Should raise IMMEDIATE_DROP in EXPAND phase
code = """
def upgrade():
op.drop_column("my_table", "col")
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_drop_column_in_contract_phase(validator):
# Should NOT raise IMMEDIATE_DROP in CONTRACT phase
code = """
def upgrade():
op.drop_column("my_table", "col")
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def test_direct_rename_table(validator):
# Should raise DIRECT_RENAME
code = """
def upgrade():
op.rename_table("old_table", "new_table")
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_direct_rename_column(validator):
# Should raise DIRECT_RENAME
code = """
def upgrade():
op.rename_column("my_table", "old_col", "new_col")
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

2. Edge Test Cases

def test_add_column_nullable_true_no_existence_check(validator):
# Should raise NO_EXISTENCE_CHECK only
code = """
def upgrade():
op.add_column("my_table", sa.Column("new_column", sa.String(), nullable=True))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_add_column_server_default_no_existence_check(validator):
# Should raise NO_EXISTENCE_CHECK only
code = """
def upgrade():
op.add_column("my_table", sa.Column("new_column", sa.String(), server_default="foo"))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_add_column_nested_call_with_nullable_true_and_existence_check(validator):
# Should NOT raise any violations (nullable=True in nested call, existence check)
code = """
def upgrade():
if "new_column" not in inspector.get_columns("my_table"):
op.add_column("my_table", sa.Column("new_column", sa.String(), nullable=True))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_add_column_nested_call_missing_nullable_and_default(validator):
# Should raise BREAKING_ADD_COLUMN and NO_EXISTENCE_CHECK
code = """
def upgrade():
op.add_column("my_table", sa.Column("new_column", sa.String()))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_existence_check_with_different_if_condition(validator):
# Should NOT detect existence check if condition is not about columns
code = """
def upgrade():
if some_flag:
op.add_column("my_table", sa.Column("new_column", sa.String(), nullable=True))
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_alter_column_multiple_keywords(validator):
# Should raise DIRECT_TYPE_CHANGE and BREAKING_ADD_COLUMN if both present
code = """
def upgrade():
op.alter_column("my_table", "col", type_=sa.Integer(), nullable=False)
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_drop_column_within_if_statement(validator):
# Should still raise IMMEDIATE_DROP in EXPAND phase, even if inside if
code = """
def upgrade():
if "col" in inspector.get_columns("my_table"):
op.drop_column("my_table", "col")
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_no_operations(validator):
# Should NOT raise any violations if no relevant operations
code = """
def upgrade():
print("No DB ops")
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def upgrade():
op.add_column("my_table", sa.Column("new_column", sa.String()))
op.alter_column("my_table", "col", type_=sa.Integer())
op.drop_column("my_table", "col")
op.rename_column("my_table", "old_col", "new_col")
"""
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output
codes = [v.code for v in violations]

3. Large Scale Test Cases

def test_many_add_column_operations(validator):
# Test scalability with many add_column operations (all missing nullable/default/existence check)
code_lines = ["def upgrade():"]
for i in range(100):
code_lines.append(f" op.add_column('my_table', sa.Column('col{i}', sa.String()))")
code = "\n".join(code_lines)
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_many_operations_contract_phase(validator):
# All drop_column operations should NOT raise IMMEDIATE_DROP in CONTRACT phase
code_lines = ["def upgrade():"]
for i in range(50):
code_lines.append(f" op.drop_column('my_table', 'col{i}')")
code = "\n".join(code_lines)
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.CONTRACT); violations = codeflash_output

def test_many_direct_rename_operations(validator):
# Should raise DIRECT_RENAME for each rename_column
code_lines = ["def upgrade():"]
for i in range(100):
code_lines.append(f" op.rename_column('my_table', 'old_col{i}', 'new_col{i}')")
code = "\n".join(code_lines)
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

def test_large_function_with_mixed_if_existence_checks(validator):
# Half add_column ops with existence check, half without
code_lines = ["def upgrade():"]
for i in range(50):
code_lines.append(f" if 'col{i}' not in inspector.get_columns('my_table'):")
code_lines.append(f" op.add_column('my_table', sa.Column('col{i}', sa.String(), nullable=True))")
for i in range(50, 100):
code_lines.append(f" op.add_column('my_table', sa.Column('col{i}', sa.String(), nullable=True))")
code = "\n".join(code_lines)
node = get_functiondef_from_code(code)
codeflash_output = validator._check_upgrade_operations(node, MigrationPhase.EXPAND); violations = codeflash_output

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr10519-2025-11-17T06.00.22 and push.

[coderabbitai]

Important

Review skipped

Bot user detected.

To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

Frontend Unit Test Coverage Report

Coverage Summary

Lines	Statements	Branches	Functions
	14.67% (3955/26947)	7.45% (1533/20552)	9% (532/5906)

Unit Test Results

Tests	Skipped	Failures	Errors	Time
1588	0 💤	0 ❌	0 🔥	19.806s ⏱️

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 31.41%. Comparing base (b641c42) to head (cfd512d).

Additional details and impacted files

@@                             Coverage Diff                              @@
##           support-n-1-database-migration-guideline   #10618      +/-   ##
============================================================================
- Coverage                                     31.65%   31.41%   -0.25%     
============================================================================
  Files                                          1330     1325       -5     
  Lines                                         60946    59987     -959     
  Branches                                       9109     8980     -129     
============================================================================
- Hits                                          19295    18844     -451     
+ Misses                                        40739    40237     -502     
+ Partials                                        912      906       -6     

Flag	Coverage Δ
backend	50.93% <ø> (-0.91%)	⬇️
frontend	13.57% <ø> (ø)
lfx	39.37% <ø> (+0.45%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 36 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.