⚡️ Speed up function foreign_key_exists by 51% in src/backend/base/langflow/utils/migration.py

[codeflash-ai]

📄 foreign_key_exists in src/backend/base/langflow/utils/migration.py

✨ Performance Summary:

• Speed Increase: 📈 51% (0.51x faster)
• Runtime Reduction: ⏱️ From 171 milliseconds down to 113 milliseconds (best of 5 runs)

---

📝 Explanation and details

To optimize the given code for checking if a foreign key exists in a table using SQLAlchemy, we can make the following changes.

1. Avoid the creation of a list and use a set for faster lookup.
2. Use Inspector directly from sqlalchemy.inspect which is more common than sqlalchemy.engine.reflection.Inspector.

Here's the optimized version.

Changes made.

1. Switched from Inspector.from_engine to inspect(conn) which is considered more direct and readable.
2. Used any generator for checking the presence of the foreign key instead of creating a list, which could consume more memory and time. This results in short-circuit evaluation and immediate return upon finding a match.

---

✅ Correctness verification

The new optimized code was tested for correctness. The results are listed below:

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

🌀 Generated Regression Tests Details

Click to view details

import pytest  # used for our unit tests
from langflow.utils.migration import foreign_key_exists
from sqlalchemy import (Column, ForeignKey, Integer, MetaData, Table,
                        create_engine)
# function to test
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.exc import NoSuchTableError

# unit tests

@pytest.fixture(scope='module')
def setup_database():
    # Create an in-memory SQLite database and setup schema for testing
    engine = create_engine('sqlite:///:memory:')
    metadata = MetaData()

    # Define tables and foreign keys
    users = Table('users', metadata,
                  Column('id', Integer, primary_key=True),
                  Column('profile_id', Integer))

    profiles = Table('profiles', metadata,
                     Column('id', Integer, primary_key=True))

    orders = Table('orders', metadata,
                   Column('id', Integer, primary_key=True),
                   Column('customer_id', Integer, ForeignKey('users.id')),
                   Column('product_id', Integer, ForeignKey('profiles.id')))

    metadata.create_all(engine)
    return engine




def test_non_existent_table(setup_database):
    engine = setup_database
    with pytest.raises(NoSuchTableError):
        foreign_key_exists('non_existent_table', 'fk_anything', engine)












import pytest  # used for our unit tests
from langflow.utils.migration import foreign_key_exists
from sqlalchemy import (Column, ForeignKey, Integer, MetaData, Table,
                        create_engine)
# function to test
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.orm import sessionmaker

# unit tests

# Setup in-memory SQLite database for testing
@pytest.fixture(scope='module')
def connection():
    engine = create_engine('sqlite:///:memory:')
    metadata = MetaData()

    # Create tables and foreign keys for testing
    users = Table('users', metadata,
                  Column('id', Integer, primary_key=True),
                  Column('profile_id', Integer, ForeignKey('profiles.id', name='fk_user_profile'))
                  )

    profiles = Table('profiles', metadata,
                     Column('id', Integer, primary_key=True)
                     )

    orders = Table('orders', metadata,
                   Column('id', Integer, primary_key=True),
                   Column('customer_id', Integer, ForeignKey('customers.id', name='fk_order_customer'))
                   )

    customers = Table('customers', metadata,
                      Column('id', Integer, primary_key=True)
                      )

    metadata.create_all(engine)
    conn = engine.connect()
    yield conn
    conn.close()

# Basic Functionality Tests


def test_foreign_key_exists_concurrent_access(connection):
    from concurrent.futures import ThreadPoolExecutor

    def check_fk():
        return foreign_key_exists('users', 'fk_user_profile', connection)

    with ThreadPoolExecutor(max_workers=10) as executor:
        results = list(executor.map(check_fk, range(10)))

# Database-Specific Scenarios Tests
@pytest.mark.parametrize("db_url", [
    'sqlite:///:memory:',
    'postgresql://user:password@localhost/testdb',
    'mysql+pymysql://user:password@localhost/testdb'
])
def test_foreign_key_exists_different_backends(db_url):
    engine = create_engine(db_url)
    metadata = MetaData()

    users = Table('users', metadata,
                  Column('id', Integer, primary_key=True),
                  Column('profile_id', Integer, ForeignKey('profiles.id', name='fk_user_profile'))
                  )

    profiles = Table('profiles', metadata,
                     Column('id', Integer, primary_key=True)
                     )

    metadata.create_all(engine)
    conn = engine.connect()
    codeflash_output = foreign_key_exists('users', 'fk_user_profile', conn)
    conn.close()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

---

📣 **Feedback**

If you have any feedback or need assistance, feel free to join our Discord community: