import json import time from abc import ABC from datetime import datetime, timezone from typing import Any, Dict, List, Optional, TypedDict import asyncio import aiohttp import pandas as pd from dagster import OpExecutionContext from src.utils.database import write_to_snowflake from src.utils.snowflake.constants import QueryType, Role, Warehouse from .constants import JobName, Vendor from ..snowflake.base import SnowflakeJob QUERY_FILE = "../bots/queries/vendor_sync.sql" class RequestParams(TypedDict): method: str url: str headers: dict json: dict entity_id: str class VendorSync(SnowflakeJob): """Base class for syncing bot data with vendors (Stripe, Clerk, Braze, etc).""" BOT_VENDOR_SYNC_TABLE = "BOT_VENDOR_SYNC" def __init__( self, name: JobName, vendor: Vendor, description: str, group_name: str, batch_size: int = 1000, rate_limit_delay: int = 3, query_limit: int = 1000, dry_run: bool = False, monitored: bool = False, owners: Optional[List[str]] = None, metadata: Optional[Dict[str, Any]] = None, tags: Optional[Dict[str, Any]] = None, schedule: Optional[str] = None, query_file: Optional[str] = None, ): super().__init__( name=name, description=description, schedule=schedule, query_file=query_file or QUERY_FILE, # Use default if not provided query_type=QueryType.SELECT, warehouse=Warehouse.MEDIUM, role=Role.ACCOUNTADMIN, group_name=group_name, monitored=monitored, owners=owners, metadata=metadata, tags=tags, deps=["bot_status_changes"], ) self.vendor = vendor.value self.batch_size = batch_size self.rate_limit_delay = rate_limit_delay self.query_limit = query_limit # Store query_limit self.dry_run = dry_run def get_query_params(self, context: OpExecutionContext) -> Dict[str, Any]: """Override base class method to provide query parameters""" return {"vendor": str(self.vendor), "limit": int(self.query_limit)} def get_request_details(self, record) -> RequestParams: """Override this method to provide request details for the record""" raise NotImplementedError def get_vendor_user_id(self, record) -> str: """Override this method to provide the vendor user ID for the record""" raise NotImplementedError async def process_single_entity(self, session, record) -> Dict: """Makes the HTTP request and handles error tracking""" vendor_user_id = self.get_vendor_user_id(record) try: request = self.get_request_details(record) # In dry run mode, just return early if self.dry_run: return { "USER_ID": record["USER_ID"], "VENDOR": self.vendor, "VENDOR_USER_ID": vendor_user_id, "IS_BOT": record["IS_BOT"], "STATUS_CODE": 200, # Simulate success "ERROR": None, "CREATED_AT": datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S"), "DAGSTER_RUN_ID": self.context.run_id, "RESPONSE": "DRY RUN - No request made", } # Normal execution async with session.request( method=request["method"], url=request["url"], headers=request["headers"], json=request["json"], ) as response: response_text = await response.text() return { "USER_ID": record["USER_UID"], "VENDOR": self.vendor, "VENDOR_USER_ID": vendor_user_id, "IS_BOT": record["IS_BOT"], "STATUS_CODE": response.status, "ERROR": response_text if response.status not in (200, 201) else None, "CREATED_AT": datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S"), "DAGSTER_RUN_ID": self.context.run_id, "RESPONSE": response_text, } except Exception as e: return { "USER_ID": record["USER_ID"], "VENDOR": self.vendor, "VENDOR_USER_ID": vendor_user_id, "IS_BOT": record["IS_BOT"], "STATUS_CODE": None, "ERROR": str(e), "CREATED_AT": datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S"), "DAGSTER_RUN_ID": self.context.run_id, "RESPONSE": None, } async def process_batch(self, batch: List) -> List[Dict]: """Process a batch of entities concurrently""" async with aiohttp.ClientSession() as session: tasks = [self.process_single_entity(session, entity) for entity in batch] return await asyncio.gather(*tasks) def _write_results_to_snowflake(self, results: List[Dict], context: OpExecutionContext) -> None: """Write results to standard BOT_VENDOR_SYNC table.""" snowflake_df = pd.DataFrame(results) if self.dry_run: context.log.info( f"DRY RUN - Would write batch of {len(snowflake_df)} records to {self.BOT_VENDOR_SYNC_TABLE}" ) context.log.info(f"Snowflake DataFrame:\n{snowflake_df.head(10).to_markdown()}") else: context.log.info( f"Writing batch of {len(snowflake_df)} records to {self.BOT_VENDOR_SYNC_TABLE}" ) context.log.info(f"Snowflake DataFrame:\n{snowflake_df.head(10).to_markdown()}") write_to_snowflake(snowflake_df, self.BOT_VENDOR_SYNC_TABLE) def _log_summary( self, total_processed: int, success_count: int, failure_count: int, context: OpExecutionContext, ) -> None: """Common logging for vendor syncs.""" context.log.info( f"{self.vendor} sync complete. " f"Total processed: {total_processed}, " f"Success: {success_count}, " f"Failures: {failure_count}" ) def post_execute(self, result: Optional[pd.DataFrame], context: OpExecutionContext) -> None: """Common execution pattern for vendor syncs""" # Store context at the start self.context = context all_results = [] total_processed = 0 success_count = 0 failure_count = 0 # Process in batches for i in range(0, len(result), self.batch_size): batch = result.iloc[i : i + self.batch_size].copy() # Process batch batch_results = asyncio.run(self.process_batch(batch.to_dict("records"))) # Track results for result_item in batch_results: if result_item.get("STATUS_CODE") == 200: success_count += 1 else: failure_count += 1 all_results.extend(batch_results) total_processed += len(batch) # Write to Snowflake periodically if len(all_results) >= 1000: self._write_results_to_snowflake(all_results, context) all_results = [] # Rate limiting if i + self.batch_size < len(result): time.sleep(self.rate_limit_delay) # Write remaining results if all_results: self._write_results_to_snowflake(all_results, context) self._log_summary(total_processed, success_count, failure_count, context)