Skip to content

Conversation

@codeflash-ai
Copy link
Contributor

@codeflash-ai codeflash-ai bot commented Nov 27, 2025

⚡️ This pull request contains optimizations for PR #10702

If you approve this dependent PR, these changes will be merged into the original PR branch pluggable-auth-service.

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


📄 24,083% (240.83x) speedup for MigrationValidator._check_downgrade_safety in src/backend/base/langflow/alembic/migration_validator.py

⏱️ Runtime : 5.08 seconds 21.0 milliseconds (best of 52 runs)

📝 Explanation and details

The optimization achieves a 24,083% speedup by eliminating redundant calls to the expensive ast.unparse() operation.

Key Optimization:

  • Batched AST unparsing: Instead of calling ast.unparse(node) for every op.alter_column call found (931 times in the profiled workload), the code now collects all matching calls first, then performs ast.unparse() only once per function when needed.
  • Shared string operations: The unparsed content and its lowercase version are computed once and reused for all backup/SELECT checks within the same function.

Why This Works:
The line profiler shows ast.unparse(node) consumed 99.7% of the original runtime (30 seconds out of 31.7 seconds total). Since ast.unparse() converts the entire function's AST back to source code, calling it repeatedly for the same function node is extremely wasteful. The optimization reduces this from O(n) calls to O(1) call per function, where n is the number of alter_column operations.

Performance Impact:

  • Large-scale migrations: The optimization excels when functions contain many alter_column operations (test cases with 500 operations show dramatic improvements)
  • Mixed workloads: Functions with no alter_column calls avoid unnecessary unparsing entirely
  • CONTRACT phase: Benefits from content sharing when both alter_column and phase checks are needed

The optimization maintains identical behavior and warning generation while dramatically reducing computational overhead for migration validation workflows.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 58 Passed
⏪ Replay Tests 🔘 None Found
🔎 Concolic Coverage Tests 🔘 None Found
📊 Tests Coverage 72.7%
🌀 Generated Regression Tests and Runtime
import ast
# Dummy classes/enums for testing
from enum import Enum

# imports
import pytest
from langflow.alembic.migration_validator import MigrationValidator


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

# ========== UNIT TESTS ==========

# Helper to parse a function string to ast.FunctionDef
def get_funcdef_from_str(func_src: str) -> ast.FunctionDef:
    module = ast.parse(func_src)
    for node in module.body:
        if isinstance(node, ast.FunctionDef):
            return node
    raise ValueError("No function definition found")

# BASIC TEST CASES

def test_no_op_calls_no_warnings():
    # Function does nothing
    func = """
def downgrade():
    pass
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_alter_column_with_backup_no_warning():
    # Function alters column, but mentions backup in code
    func = """
def downgrade():
    # backup data first
    op.alter_column('mytable', 'mycol', new_column_name='oldcol')
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_alter_column_with_SELECT_no_warning():
    # Function alters column, but mentions SELECT (as backup)
    func = '''
def downgrade():
    op.alter_column("t", "c", new_column_name="old")
    # SELECT * FROM t
'''
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_alter_column_without_backup_warning():
    # Function alters column, no backup or SELECT
    func = """
def downgrade():
    op.alter_column('mytable', 'mycol', new_column_name='oldcol')
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_multiple_alter_columns_without_backup_warning():
    # Multiple unsafe calls, should warn for each
    func = """
def downgrade():
    op.alter_column('t1', 'c1', new_column_name='old1')
    op.alter_column('t2', 'c2', new_column_name='old2')
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output
    for v in result:
        pass

# EDGE TEST CASES


def downgrade():
    op.alter_column('t', 'c', new_column_name='old')
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); result = codeflash_output
    codes = {v.code for v in result}

def test_contract_phase_with_raise_no_contract_warning():
    # CONTRACT phase, has a raise
    func = """
def downgrade():
    raise NotImplementedError("Downgrade not supported")
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); result = codeflash_output

def test_contract_phase_with_raise_and_alter_column():
    # CONTRACT phase, has both raise and alter_column
    func = """
def downgrade():
    op.alter_column('t', 'c', new_column_name='old')
    raise NotImplementedError("Downgrade not supported")
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); result = codeflash_output

def test_case_insensitive_backup():
    # backup is case insensitive
    func = """
def downgrade():
    op.alter_column('t', 'c', new_column_name='old')
    # BACKUP the data
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_alter_column_in_if_block():
    # alter_column inside an if block should still be detected
    func = """
def downgrade():
    if True:
        op.alter_column('t', 'c', new_column_name='old')
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_alter_column_on_other_object():
    # Should not warn if not op.alter_column
    func = """
def downgrade():
    other.alter_column('t', 'c', new_column_name='old')
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_alter_column_method_name_variation():
    # Should not warn if method name is not exactly 'alter_column'
    func = """
def downgrade():
    op.alter_column2('t', 'c', new_column_name='old')
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output

def test_contract_phase_with_raise_but_not_notimplemented():
    # CONTRACT phase, raise but not NotImplementedError
    func = """
def downgrade():
    raise Exception("Some other error")
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); result = codeflash_output

# LARGE SCALE TEST CASES


def downgrade():
{safe_calls}
{unsafe_calls}
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = codeflash_output
    for v in result:
        pass


def downgrade():
{calls}
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); result = codeflash_output
    codes = [v.code for v in result]


def downgrade():
{calls}
    raise NotImplementedError("Downgrade not supported")
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); result = codeflash_output


def downgrade():
{body}
"""
    node = get_funcdef_from_str(func)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); result = 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
# Define stubs for MigrationPhase and Violation to support testing
from enum import Enum

# imports
import pytest
from langflow.alembic.migration_validator import MigrationValidator


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

# unit tests

# Helper to parse a function string into an ast.FunctionDef node
def parse_funcdef_from_str(code: str) -> ast.FunctionDef:
    mod = ast.parse(code)
    for node in mod.body:
        if isinstance(node, ast.FunctionDef):
            return node
    raise ValueError("No function definition found in code.")

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

def test_no_alter_column_expand_phase():
    # Test: Function does not call op.alter_column, EXPAND phase, should not warn
    code = """
def downgrade():
    pass
"""
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_alter_column_with_backup_expand_phase():
    # Test: Calls op.alter_column, but has 'backup' in code, should not warn
    code = '''
def downgrade():
    # backup data first
    backup_data()
    op.alter_column("table", "col", new_column_name="col2")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_alter_column_with_select_expand_phase():
    # Test: Calls op.alter_column, but has 'SELECT' in code, should not warn
    code = '''
def downgrade():
    conn.execute("SELECT * FROM table")
    op.alter_column("table", "col", new_column_name="col2")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_alter_column_no_backup_expand_phase():
    # Test: Calls op.alter_column, no backup or SELECT, should warn
    code = '''
def downgrade():
    op.alter_column("table", "col", new_column_name="col2")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_contract_phase_with_raise():
    # Test: CONTRACT phase, function raises NotImplementedError, should not warn
    code = '''
def downgrade():
    raise NotImplementedError("Downgrade not supported")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); warnings = codeflash_output

def test_contract_phase_with_raise_other():
    # Test: CONTRACT phase, function raises something else, still counts as raise
    code = '''
def downgrade():
    raise Exception("Downgrade not supported")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); warnings = codeflash_output

def test_contract_phase_without_raise():
    # Test: CONTRACT phase, no raise, should warn
    code = '''
def downgrade():
    op.alter_column("table", "col", new_column_name="col2")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); warnings = codeflash_output

# ---- EDGE CASES ----

def test_alter_column_with_backup_case_insensitive():
    # Test: backup present in uppercase, should be detected
    code = '''
def downgrade():
    BACKUP_DATA()
    op.alter_column("table", "col", new_column_name="col2")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_alter_column_with_select_in_comment():
    # Test: SELECT only in a comment, not in code, should warn
    code = '''
def downgrade():
    # SELECT * FROM table
    op.alter_column("table", "col", new_column_name="col2")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_alter_column_not_op():
    # Test: Calls alter_column on something not named 'op', should not warn
    code = '''
def downgrade():
    db.alter_column("table", "col", new_column_name="col2")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_multiple_alter_columns_some_safe():
    # Test: Multiple op.alter_column, one with backup, one without
    code = '''
def downgrade():
    op.alter_column("table", "col1", new_column_name="col2")
    # backup
    op.alter_column("table", "col2", new_column_name="col3")
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_contract_phase_with_raise_and_alter_column():
    # Test: CONTRACT phase, has both raise and alter_column, should not warn
    code = '''
def downgrade():
    op.alter_column("table", "col", new_column_name="col2")
    raise NotImplementedError()
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); warnings = codeflash_output

def test_contract_phase_with_pass_only():
    # Test: CONTRACT phase, only pass, should warn
    code = '''
def downgrade():
    pass
'''
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.CONTRACT); warnings = codeflash_output

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

def test_large_number_of_safe_alter_columns():
    # Test: 500 op.alter_column with backup, should not warn
    code_lines = ["def downgrade():", "    backup_data()"]
    for i in range(500):
        code_lines.append(f'    op.alter_column("table", "col{i}", new_column_name="col{i}_old")')
    code = "\n".join(code_lines)
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

def test_large_number_of_unsafe_alter_columns():
    # Test: 500 op.alter_column without backup, should warn for each
    code_lines = ["def downgrade():"]
    for i in range(500):
        code_lines.append(f'    op.alter_column("table", "col{i}", new_column_name="col{i}_old")')
    code = "\n".join(code_lines)
    node = parse_funcdef_from_str(code)
    validator = MigrationValidator()
    codeflash_output = validator._check_downgrade_safety(node, MigrationPhase.EXPAND); warnings = codeflash_output

To edit these changes git checkout codeflash/optimize-pr10702-2025-11-27T22.46.26 and push.

Codeflash

The optimization achieves a **24,083% speedup** by eliminating redundant calls to the expensive `ast.unparse()` operation.

**Key Optimization:**
- **Batched AST unparsing**: Instead of calling `ast.unparse(node)` for every `op.alter_column` call found (931 times in the profiled workload), the code now collects all matching calls first, then performs `ast.unparse()` only once per function when needed.
- **Shared string operations**: The unparsed content and its lowercase version are computed once and reused for all backup/SELECT checks within the same function.

**Why This Works:**
The line profiler shows `ast.unparse(node)` consumed 99.7% of the original runtime (30 seconds out of 31.7 seconds total). Since `ast.unparse()` converts the entire function's AST back to source code, calling it repeatedly for the same function node is extremely wasteful. The optimization reduces this from O(n) calls to O(1) call per function, where n is the number of `alter_column` operations.

**Performance Impact:**
- **Large-scale migrations**: The optimization excels when functions contain many `alter_column` operations (test cases with 500 operations show dramatic improvements)
- **Mixed workloads**: Functions with no `alter_column` calls avoid unnecessary unparsing entirely
- **CONTRACT phase**: Benefits from content sharing when both alter_column and phase checks are needed

The optimization maintains identical behavior and warning generation while dramatically reducing computational overhead for migration validation workflows.
@codeflash-ai codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Nov 27, 2025
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Nov 27, 2025

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 github-actions bot added the community Pull Request from an external contributor label Nov 27, 2025
@github-actions
Copy link
Contributor

Frontend Unit Test Coverage Report

Coverage Summary

Lines Statements Branches Functions
Coverage: 15%
15.24% (4188/27479) 8.46% (1778/20993) 9.57% (579/6049)

Unit Test Results

Tests Skipped Failures Errors Time
1638 0 💤 0 ❌ 0 🔥 20.645s ⏱️

@codecov
Copy link

codecov bot commented Nov 27, 2025

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (pluggable-auth-service@3418a59). Learn more about missing BASE report.

Additional details and impacted files

Impacted file tree graph

@@                    Coverage Diff                    @@
##             pluggable-auth-service   #10765   +/-   ##
=========================================================
  Coverage                          ?   31.54%           
=========================================================
  Files                             ?     1369           
  Lines                             ?    63525           
  Branches                          ?     9373           
=========================================================
  Hits                              ?    20042           
  Misses                            ?    42451           
  Partials                          ?     1032           
Flag Coverage Δ
backend 47.88% <ø> (?)
frontend 14.08% <ø> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

🚀 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

⚡️ codeflash Optimization PR opened by Codeflash AI community Pull Request from an external contributor

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant