Rabbit BigQuery Job Optimizer Python Client

This is the official Python client library for the Rabbit BigQuery Job Optimizer API.

Installation

pip install rabbit-bq-job-optimizer

Usage

from rabbit_bq_job_optimizer import RabbitBQJobOptimizer, OptimizationConfig

# Initialize the client
client = RabbitBQJobOptimizer(
    api_key="your-api-key"
    # base_url is optional - will use production endpoint by default
    # timeout is optional - defaults to 5 seconds
)

# Optimize a BigQuery job configuration
job_config = {
    "configuration": {
        "query": {
            "query": "SELECT * FROM my_table",
            "useLegacySql": False,
            "priority": "INTERACTIVE"
        }
    }
}

optimization_config = OptimizationConfig(
    type="reservation_assignment",
    config={
        "defaultPricingMode": "on_demand",
        "reservationIds": [
            "my-project:US.my-reservation-us",
            "my-project:EU.my-reservation-eu"
        ]
    }
)

# Optimize the job
result = client.optimize_job(
    configuration=job_config,
    enabledOptimizations=[optimization_config]
)

# Access the optimized configuration
optimized_config = result.optimizedJob

# Access optimization results
for optimization in result.optimizationResults:
    print(f"Type: {optimization.type}")
    print(f"Applied: {optimization.performed}")
    print(f"Estimated Savings: {optimization.estimatedSavings}")

Error Handling

The client raises exceptions for API errors:

from rabbit_bq_job_optimizer import RabbitBQJobOptimizerError

try:
    result = client.optimize_job(job_config)
except RabbitBQJobOptimizerError as e:
    print(f"Error: {e.message}")
    print(f"Status Code: {e.status_code}")

Configuration

Authentication

The API key is required to use the client. It can be obtained from the Rabbit team by contacting success@followrabbit.ai. The base URL is optional and defaults to the production endpoint.

There are three ways to configure authentication:

1. Pass API key directly to the constructor (uses default production endpoint):

from rabbit_bq_job_optimizer import RabbitBQJobOptimizer

client = RabbitBQJobOptimizer(
    api_key="your-api-key"
)

2. Pass both API key and custom base URL:

from rabbit_bq_job_optimizer import RabbitBQJobOptimizer

client = RabbitBQJobOptimizer(
    api_key="your-api-key",
    base_url="https://your-custom-endpoint.com"
)

3. Set environment variables:

export RABBIT_API_KEY="your-api-key"
export RABBIT_API_BASE_URL="your-base-url"  # Optional - will use default if not set

Then initialize the client without parameters:

from rabbit_bq_job_optimizer import RabbitBQJobOptimizer

client = RabbitBQJobOptimizer()  # Will automatically use environment variables

The environment variables are:

• RABBIT_API_KEY: Your Rabbit API key (required)
• RABBIT_API_BASE_URL: The base URL for the Rabbit API (optional, defaults to production endpoint)

Request Timeout

The client includes a configurable request timeout to prevent indefinite hangs. By default, requests will timeout after 5 seconds.

You can customize the timeout when initializing the client:

from rabbit_bq_job_optimizer import RabbitBQJobOptimizer

# Use default 5 second timeout
client = RabbitBQJobOptimizer(api_key="your-api-key")

# Use custom timeout (in seconds)
client = RabbitBQJobOptimizer(
    api_key="your-api-key",
    timeout=30  # 30 seconds
)

If the API doesn't respond within the specified timeout, a RabbitBQJobOptimizerError will be raised.

Development

To install the package in development mode:

pip install -e .

License

Apache License 2.0