⚡️ Speed up method AstraDBVectorStoreComponent.get_database_list_static by 12% in PR #6048 (bugfix-dev-astradb)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #6048

If you approve this dependent PR, these changes will be merged into the original PR branch bugfix-dev-astradb.

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

---

📄 12% (0.12x) speedup for AstraDBVectorStoreComponent.get_database_list_static in src/backend/base/langflow/components/vectorstores/astradb.py

⏱️ Runtime : 44.6 milliseconds → 39.9 milliseconds (best of 5 runs)

📝 Explanation and details

To optimize the given Python program for better runtime performance, I would recommend minimizing the number of API calls and avoiding unnecessary list and dictionary operations. We will refactor the code to fetch necessary data in bulk where possible.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock, patch

# imports
import pytest  # used for our unit tests
# function to test
from astrapy import DataAPIClient
from langflow.base.vectorstores.model import LCVectorStoreComponent
from langflow.components.vectorstores.astradb import \
    AstraDBVectorStoreComponent


# unit tests
@pytest.fixture
def mock_data_api_client():
    with patch('astrapy.DataAPIClient') as mock:
        yield mock























from unittest.mock import Mock, patch

# imports
import pytest  # used for our unit tests
# function to test
from astrapy import DataAPIClient
from langflow.base.vectorstores.model import LCVectorStoreComponent
from langflow.components.vectorstores.astradb import \
    AstraDBVectorStoreComponent

# unit tests

# Basic Functionality
def test_basic_functionality_valid_token():
    token = "valid_token"
    environment = None
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = "example_db"
        mock_db.info.id = "example_id"
        mock_db.info.region = "example_region"
        mock_db.info.keyspace = "example_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        with patch('astrapy.DataAPIClient.get_database') as mock_get_database:
            mock_db_client = Mock()
            mock_db_client.list_collection_names.return_value = ["collection1"]
            mock_get_database.return_value = mock_db_client

            codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Environment Handling
def test_environment_handling_dev():
    token = "valid_token"
    environment = "dev"
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = "example_db"
        mock_db.info.id = "example_id"
        mock_db.info.region = "example_region"
        mock_db.info.keyspace = "example_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        with patch('astrapy.DataAPIClient.get_database') as mock_get_database:
            mock_db_client = Mock()
            mock_db_client.list_collection_names.return_value = ["collection1"]
            mock_get_database.return_value = mock_db_client

            codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Empty Database List
def test_empty_database_list():
    token = "valid_token"
    environment = None
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_admin_client.list_databases.return_value = []
        mock_get_admin.return_value = mock_admin_client

        codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Database with No Collections
def test_database_with_no_collections():
    token = "valid_token"
    environment = None
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = "test_db"
        mock_db.info.id = "test_id"
        mock_db.info.region = "test_region"
        mock_db.info.keyspace = "test_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        with patch('astrapy.DataAPIClient.get_database') as mock_get_database:
            mock_db_client = Mock()
            mock_db_client.list_collection_names.return_value = []
            mock_get_database.return_value = mock_db_client

            codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Invalid Token

def test_performance_large_number_of_databases():
    token = "valid_token"
    environment = None
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = "test_db"
        mock_db.info.id = "test_id"
        mock_db.info.region = "test_region"
        mock_db.info.keyspace = "test_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db] * 1000
        mock_get_admin.return_value = mock_admin_client

        with patch('astrapy.DataAPIClient.get_database') as mock_get_database:
            mock_db_client = Mock()
            mock_db_client.list_collection_names.return_value = ["collection1", "collection2"]
            mock_get_database.return_value = mock_db_client

            codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Database with Missing Information
def test_database_missing_info_id():
    token = "valid_token"
    environment = None
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = "test_db"
        mock_db.info.region = "test_region"
        mock_db.info.keyspace = "test_keyspace"
        del mock_db.info.id  # Simulate missing ID
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Special Characters in Database Information
def test_special_characters_in_database_name():
    token = "valid_token"
    environment = None
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = "db name!@#"
        mock_db.info.id = "test_id"
        mock_db.info.region = "test_region"
        mock_db.info.keyspace = "test_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Database with Extremely Long Names or IDs
def test_extremely_long_database_name():
    token = "valid_token"
    environment = None
    long_name = "a" * 256
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = long_name
        mock_db.info.id = "test_id"
        mock_db.info.region = "test_region"
        mock_db.info.keyspace = "test_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Database with Unicode Characters
def test_unicode_characters_in_database_name():
    token = "valid_token"
    environment = None
    unicode_name = "数据库"
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = unicode_name
        mock_db.info.id = "test_id"
        mock_db.info.region = "test_region"
        mock_db.info.keyspace = "test_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)

# Intermittent Network Failures
def test_intermittent_network_failures():
    token = "valid_token"
    environment = None
    with patch('astrapy.DataAPIClient.get_admin') as mock_get_admin:
        mock_admin_client = Mock()
        mock_db = Mock()
        mock_db.info.name = "test_db"
        mock_db.info.id = "test_id"
        mock_db.info.region = "test_region"
        mock_db.info.keyspace = "test_keyspace"
        mock_admin_client.list_databases.return_value = [mock_db]
        mock_get_admin.return_value = mock_admin_client

        with patch('astrapy.DataAPIClient.get_database') as mock_get_database:
            mock_db_client = Mock()
            mock_db_client.list_collection_names.side_effect = [Exception("Network error"), ["collection1"]]
            mock_get_database.return_value = mock_db_client

            codeflash_output = AstraDBVectorStoreComponent.get_database_list_static(token, environment)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

[codspeed-hq]

CodSpeed Performance Report

Merging #6084 will degrade performances by 21.44%

_{Comparing codeflash/optimize-pr6048-2025-02-03T13.49.31 (0bdc12a) with bugfix-dev-astradb (61a8b61)}

Summary

❌ 1 regressions
✅ 8 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Benchmarks breakdown

	Benchmark	BASE	HEAD	Change
❌	test_invalid_run_with_input_type_chat	16.7 ms	21.2 ms	-21.44%