1  import json
   2  import io
   3  import boto3
   4  import pandas as pd
   5  from datetime import datetime, timedelta
   6  from logger_setup import setup_logger
   7  
   8  logger = setup_logger("recall_sampler/data_writer")
   9  
  10  class DataWriter:
  11      def __init__(self, output_config: dict, customer_name: str):
  12          self.output_config = output_config
  13          self.customer_name = customer_name
  14          self.bucket = output_config.get("bucket")
  15          self.prefix = output_config.get("prefix").rstrip("/")
  16          self.file_format = output_config.get("file_format", "csv").lower()
  17          self.s3_client = boto3.client('s3')
  18  
  19          self.audit_log_key = f"{self.prefix}/customer={self.customer_name}/sample_audit_log.csv"
  20  
  21      def load_audit_log(self) -> pd.DataFrame:
  22          """Load audit log sample_audit_log.csv from S3 if exists."""
  23          try:
  24              response = self.s3_client.get_object(Bucket=self.bucket, Key=self.audit_log_key)
  25              audit_log = pd.read_csv(io.BytesIO(response['Body'].read()))
  26              audit_log['sampled_on'] = pd.to_datetime(audit_log['sampled_on'])
  27              return audit_log
  28          except self.s3_client.exceptions.NoSuchKey:
  29              # No audit log exists yet
  30              return pd.DataFrame(columns=["base_sku", "sampled_on", "sampling_strategy_type"])
  31          except Exception as e:
  32              raise RuntimeError(f"Failed to load audit log: {str(e)}")
  33  
  34      def exclude_already_sampled(self, candidates_df: pd.DataFrame, audit_log_df: pd.DataFrame) -> pd.DataFrame:
  35          """Exclude base_skus already sampled within last 1 year."""
  36          if audit_log_df.empty:
  37              return candidates_df
  38  
  39          one_year_ago = datetime.now() - timedelta(days=365)
  40          recent_samples = audit_log_df[audit_log_df['sampled_on'] >= one_year_ago]
  41          sampled_skus = set(recent_samples['base_sku'].tolist())
  42  
  43          filtered_df = candidates_df[~candidates_df['base_sku'].isin(sampled_skus)]
  44          return filtered_df
  45  
  46      def write_samples(self, sampled_df: pd.DataFrame, metadata: dict):
  47          """Write sampled data and metadata to S3 partitioned by customer/year/month."""
  48          now = datetime.now()
  49          year = now.strftime("%Y")
  50          month = now.strftime("%m")
  51  
  52          output_prefix = f"{self.prefix}/customer={self.customer_name}/year={year}/month={month}"
  53  
  54          # 1. Write sampled data
  55          sample_key = f"{output_prefix}/sampled_recall_audit.{self.file_format}"
  56          if self.file_format == "csv":
  57              csv_buffer = io.StringIO()
  58              sampled_df.to_csv(csv_buffer, index=False)
  59              self.s3_client.put_object(Bucket=self.bucket, Key=sample_key, Body=csv_buffer.getvalue())
  60          elif self.file_format == "parquet":
  61              parquet_buffer = io.BytesIO()
  62              sampled_df.to_parquet(parquet_buffer, index=False)
  63              self.s3_client.put_object(Bucket=self.bucket, Key=sample_key, Body=parquet_buffer.getvalue())
  64          else:
  65              raise ValueError(f"Unsupported file format: {self.file_format}")
  66  
  67          # 2. Write metadata.json
  68          metadata_key = f"{output_prefix}/metadata.json"
  69          metadata_bytes = json.dumps(metadata, indent=4).encode('utf-8')
  70          self.s3_client.put_object(Bucket=self.bucket, Key=metadata_key, Body=metadata_bytes)
  71  
  72          print(f"[SUCCESS] Sampled data written to S3: s3://{self.bucket}/{sample_key}")
  73          print(f"[SUCCESS] Metadata written to S3: s3://{self.bucket}/{metadata_key}")
  74  
  75      def update_audit_log(self, sampled_skus: list, sampling_strategy_type: str):
  76          """Update or create sample_audit_log.csv in S3."""
  77          now = datetime.now().strftime("%Y-%m-%d")
  78  
  79          new_entries = pd.DataFrame({
  80              "base_sku": sampled_skus,
  81              "sampled_on": now,
  82              "sampling_strategy_type": sampling_strategy_type
  83          })
  84  
  85          try:
  86              # Load existing
  87              response = self.s3_client.get_object(Bucket=self.bucket, Key=self.audit_log_key)
  88              existing_audit_log = pd.read_csv(io.BytesIO(response['Body'].read()))
  89              updated_audit_log = pd.concat([existing_audit_log, new_entries], ignore_index=True)
  90          except self.s3_client.exceptions.NoSuchKey:
  91              # No existing audit log, create new
  92              updated_audit_log = new_entries
  93          except Exception as e:
  94              raise RuntimeError(f"Failed to read existing audit log for update: {str(e)}")
  95  
  96          # Save updated audit log
  97          csv_buffer = io.StringIO()
  98          updated_audit_log.to_csv(csv_buffer, index=False)
  99          self.s3_client.put_object(Bucket=self.bucket, Key=self.audit_log_key, Body=csv_buffer.getvalue())
 100  
 101          print(f"[SUCCESS] Audit log updated at S3: s3://{self.bucket}/{self.audit_log_key}")
 102