snapraid_smart_logging/smart_logger.py

#!/usr/bin/env python3
"""
SMART Logger - Collect SMART data from Scrutiny InfluxDB,
calculate BackBlaze failure probability, and log to PostgreSQL.
"""

import os
import sys
from datetime import datetime, timezone
from typing import Dict, List, Optional
import json

from influxdb_client import InfluxDBClient
import psycopg
from psycopg.rows import dict_row
from dotenv import load_dotenv

from backblaze_tables import calculate_afr

load_dotenv()


class SmartLogger:
    def __init__(self):
        # InfluxDB configuration
        self.influx_url = os.getenv('INFLUXDB_URL')
        self.influx_token = os.getenv('INFLUXDB_TOKEN')
        self.influx_org = os.getenv('INFLUXDB_ORG', 'scrutiny')
        self.influx_bucket = os.getenv('INFLUXDB_BUCKET', 'metrics')

        # PostgreSQL configuration
        self.pg_host = os.getenv('POSTGRES_HOST', 'localhost')
        self.pg_port = int(os.getenv('POSTGRES_PORT', '5432'))
        self.pg_db = os.getenv('POSTGRES_DB', 'smart_monitoring')
        self.pg_user = os.getenv('POSTGRES_USER', 'postgres')
        self.pg_password = os.getenv('POSTGRES_PASSWORD')

        self.influx_client = None
        self.pg_conn = None

    def connect_influxdb(self):
        """Connect to InfluxDB."""
        print(f"Connecting to InfluxDB at {self.influx_url}...")
        self.influx_client = InfluxDBClient(
            url=self.influx_url,
            token=self.influx_token,
            org=self.influx_org
        )
        print("✓ Connected to InfluxDB")

    def connect_postgres(self):
        """Connect to PostgreSQL."""
        print(f"Connecting to PostgreSQL at {self.pg_host}:{self.pg_port}...")
        self.pg_conn = psycopg.connect(
            host=self.pg_host,
            port=self.pg_port,
            dbname=self.pg_db,
            user=self.pg_user,
            password=self.pg_password,
            row_factory=dict_row
        )
        print("✓ Connected to PostgreSQL")

    def fetch_smart_data(self) -> List[Dict]:
        """
        Fetch latest SMART data from InfluxDB.

        NOTE: This query needs to be customized based on your Scrutiny schema.
        Run explore_influx.py first to understand the measurement and field names.

        Returns:
            List of device data dictionaries
        """
        query_api = self.influx_client.query_api()

        # TODO: Customize this query based on your Scrutiny InfluxDB schema
        # This is a template - you'll need to adjust field names and tags
        query = f'''
        from(bucket: "{self.influx_bucket}")
            |> range(start: -5m)
            |> filter(fn: (r) => r._measurement == "smart")
            |> last()
            |> pivot(rowKey:["_time"], columnKey: ["_field"], valueColumn: "_value")
        '''

        print("Querying InfluxDB for SMART data...")
        result = query_api.query(query)

        devices = []
        for table in result:
            for record in table.records:
                # Extract device information
                # TODO: Adjust these field names based on your schema
                device_data = {
                    'timestamp': record.get_time(),
                    'device_path': record.values.get('device', 'unknown'),
                    'serial_number': record.values.get('serial', 'unknown'),
                    'model': record.values.get('model', 'unknown'),
                    'temperature': record.values.get('temp', None),
                    'power_on_hours': record.values.get('power_on_hours', None),
                    'smart_attributes': {}
                }

                # Collect SMART attributes (IDs 5, 187, 188, 197, 198)
                # TODO: Map your InfluxDB field names to SMART attribute IDs
                for attr_id in [5, 187, 188, 197, 198]:
                    field_name = f'attr_{attr_id}'  # Adjust based on your schema
                    if field_name in record.values:
                        device_data['smart_attributes'][attr_id] = record.values[field_name]

                devices.append(device_data)

        print(f"✓ Found {len(devices)} device(s)")
        return devices

    def process_device_data(self, devices: List[Dict]) -> List[Dict]:
        """
        Process device data and calculate failure probabilities.

        Args:
            devices: Raw device data from InfluxDB

        Returns:
            Processed device data with calculated AFR
        """
        processed = []

        for device in devices:
            # Determine if this is a Seagate disk
            model = device.get('model', '').lower()
            is_seagate = 'seagate' in model

            # Calculate failure probability
            smart_attrs = device.get('smart_attributes', {})
            afr = calculate_afr(smart_attrs, is_seagate=is_seagate)
            failure_probability_pct = afr * 100  # Convert to percentage

            # Calculate power-on days
            power_on_hours = device.get('power_on_hours')
            power_on_days = power_on_hours // 24 if power_on_hours else None

            # Get error count (SMART attribute 199 or similar)
            error_count = smart_attrs.get(199, 0)  # TODO: Adjust based on your needs

            processed_device = {
                'timestamp': device.get('timestamp', datetime.now(timezone.utc)),
                'device_path': device.get('device_path'),
                'serial_number': device.get('serial_number'),
                'model': device.get('model'),
                'disk_role': None,  # TODO: Map device to snapraid role if needed
                'temperature_celsius': device.get('temperature'),
                'power_on_days': power_on_days,
                'error_count': error_count,
                'size_tb': None,  # TODO: Get disk size if available
                'failure_probability_pct': round(failure_probability_pct, 2),
                'smart_attributes': json.dumps(smart_attrs)
            }

            processed.append(processed_device)

            # Print summary
            print(f"\nDevice: {processed_device['device_path']}")
            print(f"  Serial: {processed_device['serial_number']}")
            print(f"  Model: {processed_device['model']}")
            print(f"  Temp: {processed_device['temperature_celsius']}°C")
            print(f"  Power-On: {processed_device['power_on_days']} days")
            print(f"  Failure Probability: {processed_device['failure_probability_pct']:.2f}%")

        return processed

    def save_to_postgres(self, devices: List[Dict]):
        """
        Save processed device data to PostgreSQL.

        Args:
            devices: Processed device data
        """
        if not devices:
            print("No devices to save")
            return

        print(f"\nSaving {len(devices)} device record(s) to PostgreSQL...")

        insert_query = '''
        INSERT INTO smart_metrics (
            timestamp, device_path, serial_number, model, disk_role,
            temperature_celsius, power_on_days, error_count, size_tb,
            failure_probability_pct, smart_attributes
        ) VALUES (
            %(timestamp)s, %(device_path)s, %(serial_number)s, %(model)s, %(disk_role)s,
            %(temperature_celsius)s, %(power_on_days)s, %(error_count)s, %(size_tb)s,
            %(failure_probability_pct)s, %(smart_attributes)s::jsonb
        )
        ON CONFLICT (device_path, timestamp) DO UPDATE SET
            temperature_celsius = EXCLUDED.temperature_celsius,
            power_on_days = EXCLUDED.power_on_days,
            error_count = EXCLUDED.error_count,
            failure_probability_pct = EXCLUDED.failure_probability_pct,
            smart_attributes = EXCLUDED.smart_attributes
        '''

        with self.pg_conn.cursor() as cur:
            for device in devices:
                cur.execute(insert_query, device)

        self.pg_conn.commit()
        print("✓ Data saved to PostgreSQL")

    def run(self):
        """Main execution flow."""
        try:
            self.connect_influxdb()
            self.connect_postgres()

            devices = self.fetch_smart_data()
            processed = self.process_device_data(devices)
            self.save_to_postgres(processed)

            print("\n" + "=" * 80)
            print("Summary:")
            print("=" * 80)
            for device in processed:
                fp = device['failure_probability_pct']
                status = "🔴 HIGH RISK" if fp > 50 else "🟡 MEDIUM" if fp > 20 else "🟢 OK"
                print(f"{status} {device['device_path']}: {fp:.2f}% failure probability")

        except Exception as e:
            print(f"ERROR: {e}")
            import traceback
            traceback.print_exc()
            sys.exit(1)

        finally:
            if self.influx_client:
                self.influx_client.close()
            if self.pg_conn:
                self.pg_conn.close()


def main():
    # Validate environment variables
    required_vars = ['INFLUXDB_URL', 'INFLUXDB_TOKEN', 'POSTGRES_PASSWORD']
    missing = [var for var in required_vars if not os.getenv(var)]

    if missing:
        print(f"ERROR: Missing required environment variables: {', '.join(missing)}")
        print("Please create a .env file based on .env.example")
        sys.exit(1)

    logger = SmartLogger()
    logger.run()


if __name__ == '__main__':
    main()