inventory_dofns.py

gistfile1.py

import datetime
import csv
import os
import psycopg2
import apache_beam as beam
from apache_beam import DoFn
from google.protobuf import text_format

class RetrieveInventoryActivityFromPostgreSQLFn(beam.DoFn):
    """
    Connects to PostgreSQL, retrieves inventory activity logs.
    Emits them as string lines.
    """

    def __init__(self):
        pass
    
    def setup(self):
        self.conn = psycopg2.connect(
            dbname=os.getenv("DB_NAME"),
            user=os.getenv("DB_USER"),
            password=os.getenv("DB_PASSWORD"),
            host=os.getenv("DB_HOST"),
            port=int(os.getenv("DB_PORT"))
        )
    
    def process(self, element):
        with self.conn.cursor() as cursor:
            cursor.execute("SELECT product_id, timestamp, action, status, quantity, details FROM inventory_activity")
            for row in cursor.fetchall():
                product_id, timestamp, action, status, quantity, details = row
                timestamp_str = timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
                yield f"{product_id},{timestamp_str},{action},{status},{quantity},{details}"
    
    # cleanup resources when the DoFn is done
    def teardown(self):
        self.conn.close()


class ParseInventoryLogFn(DoFn):
    """
    Parses each line into a dictionary.
    Emits (product_id, parsed_log_dict) tuples.
    """
    def process(self, element):
        try:            
            if not element.startswith("product_id") and len(element.split(",")) == 6:
                line = element.split(",")
                yield (line[0], {
                    "product_id": line[0],
                    "timestamp": datetime.datetime.fromisoformat(line[1].replace('Z', '+00:00')), # Convert ISO 8601 string to datetime object
                    "action": line[2],
                    "status": line[3],
                    "quantity": int(line[4]),
                    "details": line[5],
                    "is_error": line[2] == "ERROR" or line[3] == "FAILED"
                })
        except Exception as e:
            print(f"Debug, Error parsing line: {element}. Error: {e}")


class AggregateInventoryActionsAndWriteFileFn(DoFn):
    """
    Aggregates actions for each product_id and writes the raw logs for that product
    to a separate CSV file. It also calculates summary statistics for the product.
    """
    def __init__(self, output_dir):
        self.output_dir = output_dir

    def process(self, element):
        product_id, logs = element
        
        filename = os.path.join(self.output_dir, f"product_logs_{product_id}.csv")
        logs_to_save = [{k: v for k, v in l.items() if k != "is_error"} for l in logs] # Exclude 'is_error' from saved logs
        headers = logs_to_save[0].keys() if logs_to_save else []
        
        with open(filename, "w", newline="") as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=headers)
            writer.writeheader()
            for log in logs_to_save:
                log["timestamp"] = log["timestamp"].isoformat()
                writer.writerow(log)
        
        earliest = min([l["timestamp"] for l in logs]).isoformat()
        latest = max([l["timestamp"] for l in logs]).isoformat()
        yield {
            "product_id": product_id,
            "filename": filename,
            "total_transactions": len(logs),
            "successful_transactions": sum([1 for l in logs if not l["is_error"]]),
            "total_quantity_processed": sum([l["quantity"] for l in logs if l["action"] in ["RECEIVE", "SHIP"]]),
            "error_transactions": sum([1 for l in logs if l["is_error"]]),
            "min_transaction_timestamp": earliest,
            "max_transaction_timestamp": latest,
        }


class SerializeProductSummaryToProtoFn(beam.DoFn):
    """
    DoFn to serialize a product summary dictionary into a Protobuf ProductSummary message.
    Input is a dictionary representing the product summary.
    """
    def process(self, element):
        from product_summary_pb2 import ProductSummary
        yield ProductSummary(
            product_id=element["product_id"],
            total_transactions=element["total_transactions"],
            successful_transactions=element["successful_transactions"],
            total_quantity_processed=element["total_quantity_processed"],
            error_transactions=element["error_transactions"],
            min_transaction_timestamp=element["min_transaction_timestamp"],
            max_transaction_timestamp=element["max_transaction_timestamp"],
        )


class WriteOverallInventorySummaryFn(beam.DoFn):
    """
    Formats a dictionary representing aggregated product summary data into a Proto string.
    This is used for the *overall* summary file.
    """
    def __init__(self, output_file_path):
        self.output_file_path = output_file_path

    def process(self, element):
        with open(self.output_file_path, "w") as f:
            for summary in element:
                print(f"Writing summary: {text_format.MessageToString(summary)}")
                f.write(text_format.MessageToString(summary) + "\n")