⚡️ Speed up method BigQueryUtils.dataset_exists by 14% in src/Connectors/gcp_bq_queries.py

[codeflash-ai]

📄 BigQueryUtils.dataset_exists() in src/Connectors/gcp_bq_queries.py

📈 Performance improved by 14% (0.14x faster)

⏱️ Runtime went down from 335 microseconds to 293 microseconds

Explanation and details

To optimize the given Python program for better runtime performance, we can consider the following.

1. Minimize API calls: Since the primary operation in the dataset_exists method is checking the existence of a dataset using the get_dataset method, a single API call is already quite efficient. However, we can optimize the object initialization to make sure it doesn't need to reinitialize for every check if the method is called multiple times.

2. Thread safety: Ensure that the BigQuery client initialization does not interfere with other threads if used in a multithreaded environment (although this is not directly shown in your code).

3. Remove unnecessary print statements: While they're useful for debugging, removing them can speed up execution slightly and is generally a good practice for production code.

Here's the optimized code based on the aforementioned suggestions.

Key changes.

• Initialized _client and _bqstorage_client as class-level attributes to avoid re-initialization.
• Removed print statements to slightly speed up execution and generally clean up production code.

This revised code keeps the function signatures and return values intact while focusing on minimalistic optimizations.

Correctness verification

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

🔘 (none found) − ⚙️ Existing Unit Tests

✅ 16 Passed − 🌀 Generated Regression Tests

(click to show generated tests)

# imports
from unittest import mock

import pytest  # used for our unit tests
# function to test
from google.cloud import bigquery, bigquery_storage
from google.cloud.exceptions import NotFound
from src.Connectors.gcp_bq_queries import BigQueryUtils

# unit tests

# Helper function to mock get_dataset behavior
def mock_get_dataset(dataset_id):
    if dataset_id in ["existing_dataset", "existing_dataset_special_chars!@#", "existing_dataset_with_a_very_long_name_exceeding_typical_length"]:
        return True
    else:
        raise NotFound("Dataset not found")
    # Outputs were verified to be equal to the original implementation

@pytest.fixture
def bigquery_utils():
    # Fixture to create an instance of BigQueryUtils
    return BigQueryUtils(project_id="test_project")
    # Outputs were verified to be equal to the original implementation

def test_dataset_exists_basic(bigquery_utils):
    # Mock the get_dataset method
    with mock.patch.object(bigquery.Client, 'get_dataset', side_effect=mock_get_dataset):
        # Test for existing datasets
        codeflash_output = bigquery_utils.dataset_exists("existing_dataset")
        codeflash_output = bigquery_utils.dataset_exists("existing_dataset_special_chars!@#")
        codeflash_output = bigquery_utils.dataset_exists("existing_dataset_with_a_very_long_name_exceeding_typical_length")
    # Outputs were verified to be equal to the original implementation

def test_dataset_does_not_exist_basic(bigquery_utils):
    # Mock the get_dataset method
    with mock.patch.object(bigquery.Client, 'get_dataset', side_effect=mock_get_dataset):
        # Test for non-existing datasets
        codeflash_output = bigquery_utils.dataset_exists("non_existing_dataset")
        codeflash_output = bigquery_utils.dataset_exists("non_existing_dataset_special_chars!@#")
        codeflash_output = bigquery_utils.dataset_exists("non_existing_dataset_with_a_very_long_name_exceeding_typical_length")
    # Outputs were verified to be equal to the original implementation

def test_invalid_dataset_id(bigquery_utils):
    # Mock the get_dataset method
    with mock.patch.object(bigquery.Client, 'get_dataset', side_effect=mock_get_dataset):
        # Test for invalid dataset IDs
        codeflash_output = bigquery_utils.dataset_exists("")
        codeflash_output = bigquery_utils.dataset_exists(None)
        codeflash_output = bigquery_utils.dataset_exists("    ")
    # Outputs were verified to be equal to the original implementation

def test_edge_cases(bigquery_utils):
    # Mock the get_dataset method
    with mock.patch.object(bigquery.Client, 'get_dataset', side_effect=mock_get_dataset):
        # Test for edge cases
        codeflash_output = bigquery_utils.dataset_exists("a" * 1024)
        codeflash_output = bigquery_utils.dataset_exists("データセット")
        codeflash_output = bigquery_utils.dataset_exists("dataset; DROP TABLE users; --")
    # Outputs were verified to be equal to the original implementation



def test_boundary_conditions(bigquery_utils):
    # Mock the get_dataset method
    with mock.patch.object(bigquery.Client, 'get_dataset', side_effect=mock_get_dataset):
        # Test for boundary conditions
        codeflash_output = bigquery_utils.dataset_exists("a" * 1024)
        codeflash_output = bigquery_utils.dataset_exists("a" * 1023)
    # Outputs were verified to be equal to the original implementation

def test_mixed_case_sensitivity(bigquery_utils):
    # Mock the get_dataset method
    with mock.patch.object(bigquery.Client, 'get_dataset', side_effect=mock_get_dataset):
        # Test for mixed case sensitivity
        codeflash_output = bigquery_utils.dataset_exists("DatasetMixedCASE")
        codeflash_output = bigquery_utils.dataset_exists("DATASETUPPERCASE")
    # Outputs were verified to be equal to the original implementation

🔘 (none found) − ⏪ Replay Tests