reform the data generator script

Author: HarshCasper

bin/data_generator.py

@@ -3,125 +3,135 @@
 import csv
 import os
 import argparse
+import glob
+import re
 
-def generate_sensor_data(num_records=100, machines=None, anomaly_probability=0.15):
-    """
-    Generate synthetic sensor data with occasional anomalies
-    
-    Parameters:
-    - num_records: Number of records to generate
-    - machines: List of machine IDs (defaults to M001-M010)
-    - anomaly_probability: Chance of generating anomalous readings
-    
-    Returns:
-    - List of dictionaries containing sensor data
-    """
+def get_latest_batch_info(output_dir="data"):
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        return 1, None
+
+    batch_files = glob.glob(os.path.join(output_dir, "sensor_data_batch_*.csv"))
+    if not batch_files:
+        return 1, None
+
+    latest_batch_num = 0
+    latest_file = None
+
+    for f in batch_files:
+        match = re.search(r"sensor_data_batch_(\d+).csv", os.path.basename(f))
+        if match:
+            batch_num = int(match.group(1))
+            if batch_num > latest_batch_num:
+                latest_batch_num = batch_num
+                latest_file = f
+    
+    if latest_file:
+        try:
+            with open(latest_file, 'r', newline='') as csvfile:
+                reader = csv.reader(csvfile)
+                header = next(reader)
+                last_line = None
+                for row in reader:
+                    if row:
+                        last_line = row
+                
+                if last_line:
+                    last_timestamp_str = last_line[1]
+                    last_timestamp_dt = datetime.datetime.fromisoformat(last_timestamp_str.replace('Z', '+00:00'))
+                    return latest_batch_num + 1, last_timestamp_dt
+        except Exception as e:
+            print(f"Warning: Could not read or parse last timestamp from {latest_file}: {e}")
+            # Fallback if reading fails, start fresh for the next batch number
+            return latest_batch_num + 1, datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=1) 
+
+    return latest_batch_num + 1, None # Should not happen if latest_file was found, but as a fallback
+
+
+def generate_sensor_data(
+    num_records_per_machine=50, 
+    machines=None, 
+    anomaly_probability=0.15,
+    start_timestamp=None
+):
     if machines is None:
         machines = [f"M{str(i).zfill(3)}" for i in range(1, 11)]
 
-    # Base timestamp (current time minus some random offset)
-    base_time = datetime.datetime.now() - datetime.timedelta(hours=random.randint(0, 24))
-    
-    # Normal operating ranges for each sensor
+    num_total_records = num_records_per_machine * len(machines)
+
+    if start_timestamp is None:
+        current_time = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=1)
+    else:
+        current_time = start_timestamp
+
     normal_ranges = {
-        "temperature": (65.0, 80.0),  # degrees Celsius
-        "vibration": (0.01, 0.2),     # mm/s
-        "pressure": (98.0, 102.0)     # kPa
+        "temperature": (65.0, 80.0),
+        "vibration": (0.01, 0.2),
+        "pressure": (98.0, 102.0)
     }
 
-    # Warning thresholds
     warning_thresholds = {
-        "temperature": 85.0,  # above this is WARNING
-        "vibration": 0.5,     # above this is WARNING
-        "pressure_high": 105.0,  # above this is WARNING
-        "pressure_low": 95.0     # below this is WARNING
+        "temperature": 85.0,
+        "vibration": 0.5,
+        "pressure_high": 105.0,
+        "pressure_low": 95.0
     }
 
-    # Critical thresholds
     critical_thresholds = {
-        "temperature": 95.0,  # above this is CRITICAL
-        "vibration": 0.8,     # above this is CRITICAL
-        "pressure_high": 110.0,  # above this is CRITICAL
-        "pressure_low": 92.0     # below this is CRITICAL
+        "temperature": 95.0,
+        "vibration": 0.8,
+        "pressure_high": 110.0,
+        "pressure_low": 92.0
     }
 
     data = []
+    machine_record_counts = {machine_id: 0 for machine_id in machines}
 
-    for i in range(num_records):
-        # Select a random machine
-        machine_id = random.choice(machines)
+    available_machines = list(machines)
+
+    for i in range(num_total_records):
+        if not available_machines:
+            break 
 
-        # Generate timestamp with some randomness
-        timestamp = base_time + datetime.timedelta(
-            minutes=random.randint(0, 60),
-            seconds=random.randint(0, 59)
-        )
+        machine_id = random.choice(available_machines)
 
-        # Decide if this reading should be anomalous
+        current_time += datetime.timedelta(seconds=random.randint(1, 15 + i//10))
+
         is_anomaly = random.random() < anomaly_probability
-        is_critical = is_anomaly and random.random() < 0.3  # 30% of anomalies are critical
+        is_critical = is_anomaly and random.random() < 0.3
 
-        # Generate sensor readings
+        status = "AOK"
+        temperature = random.uniform(*normal_ranges["temperature"])
+        vibration = random.uniform(*normal_ranges["vibration"])
+        pressure = random.uniform(*normal_ranges["pressure"])
+
         if is_anomaly:
+            anomaly_type = random.choice(["temperature", "vibration", "pressure"])
             if is_critical:
-                # Critical anomaly
-                anomaly_type = random.choice(["temperature", "vibration", "pressure"])
-                
+                status = "CRIT"
                 if anomaly_type == "temperature":
                     temperature = random.uniform(critical_thresholds["temperature"], critical_thresholds["temperature"] + 10)
-                    vibration = random.uniform(*normal_ranges["vibration"])
-                    pressure = random.uniform(*normal_ranges["pressure"])
-                    status = "CRIT"
                 elif anomaly_type == "vibration":
-                    temperature = random.uniform(*normal_ranges["temperature"])
                     vibration = random.uniform(critical_thresholds["vibration"], critical_thresholds["vibration"] + 0.5)
-                    pressure = random.uniform(*normal_ranges["pressure"])
-                    status = "CRIT"
-                else:  # pressure anomaly
-                    temperature = random.uniform(*normal_ranges["temperature"])
-                    vibration = random.uniform(*normal_ranges["vibration"])
-                    
-                    # Either too high or too low pressure
+                else: # pressure
                     if random.random() < 0.5:
                         pressure = random.uniform(critical_thresholds["pressure_high"], critical_thresholds["pressure_high"] + 5)
                     else:
                         pressure = random.uniform(critical_thresholds["pressure_low"] - 5, critical_thresholds["pressure_low"])
-                    status = "CRIT"
-            else:
-                # Warning anomaly
-                anomaly_type = random.choice(["temperature", "vibration", "pressure"])
-                
+            else: # Warning
+                status = "WARN"
                 if anomaly_type == "temperature":
-                    temperature = random.uniform(warning_thresholds["temperature"], critical_thresholds["temperature"])
-                    vibration = random.uniform(*normal_ranges["vibration"])
-                    pressure = random.uniform(*normal_ranges["pressure"])
-                    status = "WARN"
+                    temperature = random.uniform(warning_thresholds["temperature"], critical_thresholds["temperature"] -0.1) # ensure it's below critical
                 elif anomaly_type == "vibration":
-                    temperature = random.uniform(*normal_ranges["temperature"])
-                    vibration = random.uniform(warning_thresholds["vibration"], critical_thresholds["vibration"])
-                    pressure = random.uniform(*normal_ranges["pressure"])
-                    status = "WARN"
-                else:  # pressure anomaly
-                    temperature = random.uniform(*normal_ranges["temperature"])
-                    vibration = random.uniform(*normal_ranges["vibration"])
-                    
-                    # Either too high or too low pressure
+                    vibration = random.uniform(warning_thresholds["vibration"], critical_thresholds["vibration"]-0.01)
+                else: # pressure
                     if random.random() < 0.5:
-                        pressure = random.uniform(warning_thresholds["pressure_high"], critical_thresholds["pressure_high"])
+                        pressure = random.uniform(warning_thresholds["pressure_high"], critical_thresholds["pressure_high"]-0.1)
                     else:
-                        pressure = random.uniform(critical_thresholds["pressure_low"], warning_thresholds["pressure_low"])
-                    status = "WARN"
-        else:
-            # Normal readings
-            temperature = random.uniform(*normal_ranges["temperature"])
-            vibration = random.uniform(*normal_ranges["vibration"])
-            pressure = random.uniform(*normal_ranges["pressure"])
-            status = "AOK"
+                        pressure = random.uniform(critical_thresholds["pressure_low"]+0.1, warning_thresholds["pressure_low"])
 
-        # Format timestamp as ISO format
-        timestamp_str = timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
+        timestamp_str = current_time.strftime("%Y-%m-%dT%H:%M:%SZ")
 
-        # Create data record
         record = {
             "machine_id": machine_id,
             "timestamp": timestamp_str,
@@ -130,10 +140,12 @@ def generate_sensor_data(num_records=100, machines=None, anomaly_probability=0.1
             "pressure": round(pressure, 1),
             "status_code": status
         }
-        
         data.append(record)
-    
-    # Sort by timestamp
+
+        machine_record_counts[machine_id] += 1
+        if machine_record_counts[machine_id] >= num_records_per_machine:
+            available_machines.remove(machine_id)
+            
     data.sort(key=lambda x: x["timestamp"])
 
     return data
@@ -150,25 +162,32 @@ def write_csv(data, filename):
         writer.writerows(data)
 
     print(f"Generated {len(data)} records and saved to {filename}")
-    
     return filename
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Generate synthetic sensor data")
-    parser.add_argument("--records", type=int, default=100, help="Number of records to generate (default: 100)")
-    parser.add_argument("--output", type=str, default="data/generated_sensor_data.csv", help="Output CSV file path")
+    parser = argparse.ArgumentParser(description="Generate synthetic sensor data in batches.")
+    parser.add_argument("--records_per_machine", type=int, default=10, help="Number of records per machine (default: 50)")
+    parser.add_argument("--output_dir", type=str, default="data", help="Output directory for CSV files (default: data)")
     parser.add_argument("--anomalies", type=float, default=0.15, help="Probability of anomalies (0-1, default: 0.15)")
 
     args = parser.parse_args()
-    
-    # Generate sensor data
-    data = generate_sensor_data(args.records, anomaly_probability=args.anomalies)
-    
-    # Write to CSV
-    output_file = write_csv(data, args.output)
-    
-    print(f"\nNext steps:")
-    print(f"1. Upload the file to S3 to trigger Snowpipe:")
-    print(f"   python setup/03_upload_file.py --file {output_file}")
-    print(f"2. Check the pipeline status:")
-    print(f"   python setup/check_pipeline_status.py") 
+
+    next_batch_num, last_timestamp = get_latest_batch_info(args.output_dir)
+    
+    if last_timestamp:
+        print(f"Last timestamp from batch {next_batch_num-1}: {last_timestamp.strftime('%Y-%m-%dT%H:%M:%SZ')}")
+        start_timestamp_for_new_batch = last_timestamp + datetime.timedelta(seconds=random.randint(10,30)) 
+    else:
+        print("No previous batch found or unable to read last timestamp. Starting fresh.")
+        start_timestamp_for_new_batch = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=1)
+
+    print(f"Generating data for batch {next_batch_num}, starting after ~{start_timestamp_for_new_batch.strftime('%Y-%m-%dT%H:%M:%SZ')}")
+
+    data = generate_sensor_data(
+        num_records_per_machine=args.records_per_machine,
+        anomaly_probability=args.anomalies,
+        start_timestamp=start_timestamp_for_new_batch
+    )
+    
+    output_filename = os.path.join(args.output_dir, f"sensor_data_batch_{next_batch_num}.csv")
+    output_file = write_csv(data, output_filename)