smartmontoolstopsql/attrlogimport.py

#!/usr/bin/env python3
import os
import re
import sys
from datetime import datetime
from zoneinfo import ZoneInfo
from pytz import timezone

import psycopg

drive_name_re = re.compile(r'attrlog\.(.*).ata.csv')

conn = psycopg.connect("dbname=smartmontools")


def create_tables():
    with conn.cursor() as cur:
        cur.execute("START TRANSACTION")
        cur.execute("""
            CREATE TABLE device (
                id serial primary key,
                name text
            )""")
        cur.execute("""
            CREATE TABLE attrlog (
                time timestamptz not null,
                device_id integer references device(id) not null,
                id smallint not null,
                norm smallint,
                raw bigint,
                unique (time, device_id, id)
            )
        """)
        cur.execute("COMMIT")


def create_or_find_device_id(name):
    with conn.cursor() as cur:
        cur.execute("SELECT id FROM device WHERE name = %s", (name,))
        row = cur.fetchone()
        if row is None:
            cur.execute("INSERT INTO device (name) VALUES (%s) RETURNING id", (name,))
            row = cur.fetchone()
        return row[0]


def parse_attrlog_file(filename, device_id=None):
    tz = timezone("Europe/Tallinn")
    utc = timezone("UTC")
    fd = open(filename)
    file_size = fd.seek(0, os.SEEK_END)
    fd.seek(0)
    while fd.tell() != file_size:
        line = fd.readline()
        line_parts = [p for p in line.strip().split(";") if p.strip()]
        #dt = line_parts.pop(0)+"+0200"
        dt = datetime.strptime(line_parts.pop(0), "%Y-%m-%d %H:%M:%S")#.replace(tzinfo=tz)
        dtz = dt.replace(tzinfo=tz)
        dtu = dt.astimezone(utc)
        print(dt, dtz, dtu, sep='; ')
        while line_parts:
            id = int(line_parts.pop(0))
            norm = int(line_parts.pop(0))
            raw = int(line_parts.pop(0))
            yield dtu, id, norm, raw, device_id
            if fd.tell() % 100 == 0:
                print(f"{int((fd.tell() / file_size)*100):>5}%", end='\r')
    print()


def import_attrlog_file(filename):
    drive_name = drive_name_re.search(filename).group(1)
    print(drive_name)
    device_id = create_or_find_device_id(drive_name)
    with conn.cursor() as cur:
        cur.execute("START TRANSACTION")
        for row in parse_attrlog_file(filename, device_id):
            cur.execute("""
            INSERT INTO attrlog (time, id, norm, raw, device_id)
            VALUES (%s, %s, %s, %s, %s)
            """, row)
        # with cur.copy("copy attrlog (time, id, norm, raw, device_id) FROM STDIN") as copy:
        #     for row in parse_attrlog_file(filename, device_id):
        #         #print(row)
        #         copy.write_row(row)
        cur.execute("COMMIT")


if __name__ == '__main__':
    if len(sys.argv) < 2:
        create_tables()
    else:
        import_attrlog_file(sys.argv[1])