from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
from config import Config
from typing import Dict, Any, List, Optional
import pendulum
import time
from collections import defaultdict
[docs]
class RunCalixTestOperator(BaseOperator):
"""
An operator to run Calix speed and/or latency tests on multiple devices.
Processes test requests in chunks, makes concurrent API requests (respecting rate limits),
and writes results to Redshift tables for efficient incremental inserts.
"""
ui_color = "#A6E6A6"
template_fields = ["test_configs"]
# URL constants
RUN_SPEED_URL = "https://api.calix.ai/v1/csc/operator/system/run-speed-test"
RUN_LATENCY_URL = "https://api.calix.ai/v1/csc/operator/system/run-latency-test"
TEST_TYPE_URLS = {"speed": RUN_SPEED_URL, "latency": RUN_LATENCY_URL}
# Redshift connection & table settings
REDSHIFT_CONN_ID = Config.REDSHIFT
REDSHIFT_DB = "dev"
REDSHIFT_SCHEMA = "data"
DEFAULT_TABLES = {
"speed": "t_calix_api_speed_results",
"latency": "t_calix_api_latency_results",
}
DELAY_START_SPEED_TEST = 10 # seconds
@apply_defaults
def __init__(
self,
test_configs: List[Dict],
chunk_size: int = 100,
max_concurrent_requests: int = 25,
max_retries: int = 2,
retry_delay_seconds: int = 5,
redshift_conn_id: str = REDSHIFT_CONN_ID,
redshift_schema: str = REDSHIFT_SCHEMA,
redshift_db: str = REDSHIFT_DB,
*args,
**kwargs,
):
"""
Initialize the RunCalixTestOperator.
:param test_configs: List of test configuration dictionaries with:
- serial_number: Device serial number
- token: Calix API access token
- expire_ts: Token expiration timestamp
- client_id: Calix client ID
- test_type: Type of test ('speed' or 'latency')
- output_table: (optional) Custom output table name
:type test_configs: List[Dict]
:param chunk_size: Number of devices to test in each chunk
:type chunk_size: int, optional
:param max_concurrent_requests: Maximum concurrent API requests
:type max_concurrent_requests: int, optional
:param max_retries: Maximum number of retries for failed API calls
:type max_retries: int, optional
:param retry_delay_seconds: Delay between retries in seconds
:type retry_delay_seconds: int, optional
:param redshift_conn_id: Redshift connection ID
:type redshift_conn_id: str, optional
:param redshift_schema: Redshift schema name
:type redshift_schema: str, optional
:param redshift_db: Redshift database name
:type redshift_db: str, optional
"""
super().__init__(*args, **kwargs)
self.test_configs = test_configs
self.chunk_size = chunk_size
self.max_concurrent_requests = max_concurrent_requests
self.max_retries = max_retries
self.retry_delay_seconds = retry_delay_seconds
self.redshift_conn_id = redshift_conn_id
self.redshift_schema = redshift_schema
self.redshift_db = redshift_db
[docs]
def get_output_table(self, config: Dict) -> str:
"""Get output table name from config or default"""
test_type = config.get("test_type")
return config.get("output_table", self.DEFAULT_TABLES.get(test_type))
[docs]
async def run_test(
self, session, semaphore, config, retries=None, delay_start=0
) -> Optional[Dict[str, Any]]:
"""
Submit a test request for a single device and process the response.
"""
import asyncio
if delay_start > 0:
await asyncio.sleep(delay_start)
if retries is None:
retries = self.max_retries
test_type = config.get("test_type")
if test_type not in self.TEST_TYPE_URLS:
self.log.error(
f"Invalid test_type {test_type} for {config.get('serial_number')}"
)
return None
serial_number = config.get("serial_number")
token = config.get("token")
expire_ts = config.get("expire_ts")
client_id = config.get("client_id")
endpoint = self.TEST_TYPE_URLS[test_type]
current_time = pendulum.now("UTC")
if current_time.int_timestamp > expire_ts:
self.log.warning(
f"Skipping device {serial_number} due to expired token. Will be retried in next run."
)
return None
headers = {
"X-Calix-ClientID": client_id,
"X-Calix-AccessToken": token,
"Content-Type": "application/json",
"Accept": "application/json",
}
payload = {"serialNumber": serial_number}
result = {
"serial_number": serial_number,
"client_id": client_id,
"error_message": None,
"test_timestamp": current_time.to_iso8601_string(),
"test_type": test_type,
"output_table": self.get_output_table(config),
}
async with semaphore:
try:
self.log.info(
f"Running {test_type}-test for device {serial_number} ({retries} retry[s] left)"
)
async with session.post(
endpoint, headers=headers, json=payload
) as response:
if response.status == 200:
data = await response.json()
if test_type == "speed":
result.update(
{
"calix_test_id": data.get("id"),
"ave_download_rate": data.get("aveDownloadRate"),
"ave_upload_rate": data.get("aveUploadRate"),
"phy_link_rate": data.get("phyLinkRate"),
"hit_cpu_limit": data.get("hitCpuLimit"),
"server_host": data.get("serverHost"),
"ds": data.get("ds"),
"us": data.get("us"),
"trigger": data.get("trigger"),
"ctime": data.get("ctime"),
"ds_tier": data.get("dsTier"),
"ds_level": data.get("dsLevel"),
"ds_tier_next": data.get("dsTierNext"),
"us_tier": data.get("usTier"),
"us_level": data.get("usLevel"),
"us_tier_next": data.get("usTierNext"),
"timestamp": data.get("timestamp"),
}
)
if data.get("timestamp"):
timestamp_seconds = data.get("timestamp") / 1000
test_time = pendulum.from_timestamp(
timestamp_seconds, tz="UTC"
)
result["test_timestamp"] = test_time.to_iso8601_string()
elif test_type == "latency":
result.update(
{
"calix_test_id": data.get("id"),
"latency": data.get("latency"),
"packets_sent": data.get("packetsSent"),
"packets_received": data.get("packetsReceived"),
"server_host": data.get("serverHost"),
"trigger": data.get("trigger"),
"create_time": data.get("createTime"),
}
)
if data.get("createTime"):
create_time_seconds = data.get("createTime") / 1000
create_time = pendulum.from_timestamp(
create_time_seconds, tz="UTC"
)
result["test_timestamp"] = (
create_time.to_iso8601_string()
)
self.log.info(f"COMPLETED {test_type}-test for {serial_number}")
elif response.status == 429: # Rate limit exceeded
if retries > 0:
self.log.warning(
f"Rate limit exceeded for {serial_number}. Retrying in {self.retry_delay_seconds}s..."
)
await asyncio.sleep(self.retry_delay_seconds)
return await self.run_test(
session, semaphore, config, retries - 1
)
else:
result["error_message"] = (
"Rate limit exceeded and max retries reached"
)
else:
error_text = await response.text()
result["error_message"] = error_text[:500]
self.log.error(
f"API error for {test_type}-{serial_number}: {response.status} - {error_text}"
)
except Exception as e:
self.log.error(f"Exception for {serial_number}: {str(e)}")
result["error_message"] = f"{type(e).__name__}: {str(e)[:500]}"
return result
[docs]
async def process_chunk(self, chunk_configs):
"""
Process a chunk of test configurations with staggered execution for same-device tests.
"""
import asyncio
import aiohttp
serial_configs = defaultdict(list)
for config in chunk_configs:
serial_configs[config["serial_number"]].append(config)
semaphore = asyncio.Semaphore(self.max_concurrent_requests)
tasks = []
async with aiohttp.ClientSession() as session:
for serial_number, configs in serial_configs.items():
sorted_configs = sorted(
configs, key=lambda x: 0 if x.get("test_type") == "latency" else 1
)
for i, config in enumerate(sorted_configs):
delay = (
self.DELAY_START_SPEED_TEST
if (i > 0 and config["test_type"] == "speed")
else 0
)
tasks.append(
self.run_test(session, semaphore, config, delay_start=delay)
)
test_results = await asyncio.gather(*tasks)
return [result for result in test_results if result is not None]
[docs]
def insert_to_redshift(
self, results: List[Dict], test_type: str, output_table: str
) -> int:
"""
Insert test results directly into Redshift tables using parameterized SQL.
Returns: Number of records successfully inserted
"""
if not results:
self.log.warning(f"No {test_type} results to insert")
return 0
# Get a Redshift hook
hook = RedshiftSQLHook(
redshift_conn_id=self.redshift_conn_id,
schema=self.redshift_schema,
database=self.redshift_db,
)
# Keep track of rows inserted
total_inserted = 0
for result in results:
if test_type == "speed":
# Build parameterized SQL for speed test results
sql = f"""
INSERT INTO {self.redshift_schema}.{output_table} (
serial_number, error_message, test_timestamp,
calix_test_id, ave_download_rate, ave_upload_rate,
phy_link_rate, hit_cpu_limit, server_host,
ds, us, trigger, ctime, ds_tier, ds_level, ds_tier_next,
us_tier, us_level, us_tier_next, timestamp, client_id
) VALUES (
%s, %s, %s::TIMESTAMPTZ,
%s, %s::DOUBLE PRECISION, %s::DOUBLE PRECISION,
%s::BIGINT, %s::BOOLEAN, %s,
%s::BIGINT, %s::BIGINT, %s, %s, %s::BIGINT, %s::BIGINT, %s::BIGINT,
%s::BIGINT, %s::BIGINT, %s::BIGINT, %s::BIGINT, %s
);
"""
params = (
result.get("serial_number"),
result.get("error_message"),
result.get("test_timestamp"),
result.get("calix_test_id"),
result.get("ave_download_rate"),
result.get("ave_upload_rate"),
result.get("phy_link_rate"),
result.get("hit_cpu_limit"),
result.get("server_host"),
result.get("ds"),
result.get("us"),
result.get("trigger"),
result.get("ctime"),
result.get("ds_tier"),
result.get("ds_level"),
result.get("ds_tier_next"),
result.get("us_tier"),
result.get("us_level"),
result.get("us_tier_next"),
result.get("timestamp"),
result.get("client_id"),
)
else: # latency
# Build parameterized SQL for latency test results
sql = f"""
INSERT INTO {self.redshift_schema}.{output_table} (
serial_number, error_message, test_timestamp,
calix_test_id, latency, packets_sent, packets_received,
server_host, trigger, create_time, client_id
) VALUES (
%s, %s, %s::TIMESTAMPTZ,
%s, %s::DOUBLE PRECISION, %s::INTEGER, %s::INTEGER,
%s, %s, %s::BIGINT, %s
);
"""
params = (
result.get("serial_number"),
result.get("error_message"),
result.get("test_timestamp"),
result.get("calix_test_id"),
result.get("latency"),
result.get("packets_sent"),
result.get("packets_received"),
result.get("server_host"),
result.get("trigger"),
result.get("create_time"),
result.get("client_id"),
)
# Execute the insert with proper error handling
try:
hook.run(sql=sql, parameters=params)
total_inserted += 1
except Exception as e:
self.log.error(f"Error inserting record to Redshift: {str(e)}")
# Continue with other records even if one fails
self.log.info(
f"Successfully inserted {total_inserted} records into {self.redshift_schema}.{output_table}"
)
return total_inserted
[docs]
def execute(self, context):
"""Execute the operator: run tests, process results, and store data directly to Redshift."""
import asyncio
start_time = time.time()
# Validate test_configs is a list now that templating is complete
if not isinstance(self.test_configs, list):
self.log.error(
f"test_configs must be a list, got {type(self.test_configs)}"
)
# Convert XCom reference or other iterables to list if possible
self.test_configs = list(self.test_configs)
self.log.info(
f"Converted test_configs to list with {len(self.test_configs)} items"
)
if not self.test_configs:
self.log.warning("No test configurations provided. Nothing to execute.")
return {"total_tests": 0, "successful_tests": 0, "success_rate": 0}
# Validate test types in configs
invalid_configs = []
for i, config in enumerate(self.test_configs):
test_type = config.get("test_type")
if test_type not in self.TEST_TYPE_URLS:
self.log.warning(
f"Invalid test_type '{test_type}' in config #{i} for serial_number '{config.get('serial_number')}'. "
f"Must be one of {list(self.TEST_TYPE_URLS.keys())}"
)
invalid_configs.append(i)
# Remove invalid configs
if invalid_configs:
self.log.warning(f"Removing {len(invalid_configs)} invalid configurations")
self.test_configs = [
c for i, c in enumerate(self.test_configs) if i not in invalid_configs
]
if not self.test_configs:
self.log.warning(
"No valid test configurations remained. Nothing to execute."
)
return {"total_tests": 0, "successful_tests": 0, "success_rate": 0}
all_test_results = []
total_configs = len(self.test_configs)
# Process in chunks
for i in range(0, total_configs, self.chunk_size):
chunk_end = min(i + self.chunk_size, total_configs)
chunk_configs = self.test_configs[i:chunk_end]
self.log.info(
f"Processing chunk {i//self.chunk_size + 1}/{(total_configs-1)//self.chunk_size + 1} "
f"({len(chunk_configs)} devices)"
)
chunk_results = asyncio.run(self.process_chunk(chunk_configs))
all_test_results.extend(chunk_results)
# Track inserted records count
inserted_records = 0
# Group results by test type and output table and insert directly to Redshift
if all_test_results:
results_by_table = defaultdict(list)
for result in all_test_results:
if result:
key = (result.pop("test_type"), result.pop("output_table"))
results_by_table[key].append(result)
# Insert each group directly to Redshift, bypassing the DataFrame conversion
for (test_type, output_table), results in results_by_table.items():
if results:
inserted = self.insert_to_redshift(results, test_type, output_table)
inserted_records += inserted
elapsed_time = time.time() - start_time
successful_tests = len(
[r for r in all_test_results if not r.get("error_message")]
)
result = {
"total_tests": len(all_test_results),
"successful_tests": successful_tests,
"success_rate": (
successful_tests / len(all_test_results) if all_test_results else 0
),
"execution_time": elapsed_time,
"inserted_records": inserted_records,
}
self.log.info(
f"Completed {result['total_tests']} tests "
f"({result['successful_tests']} successful, {inserted_records} records inserted) "
f"in {elapsed_time:.2f} seconds"
)
return result