spacesavers2_pdq_update_db

#!/usr/bin/env python3
# pqd = pretty darn quick
# This script appends records from the supplied TSV file into the supplied sqlite3 db file

from src.VersionCheck import version_check
from src.VersionCheck import __version__
from src.utils import *

version_check()

# import required modules
import sqlite3
import textwrap
import argparse
from pathlib import Path
import subprocess
import pandas
import sys

# data access functions


def get_full_name(uid):
    cmd = "getent passwd {}".format(uid)
    results = subprocess.run(cmd,
                             shell=True,
                             text=True,
                             capture_output=True)
    if results.returncode != 0:
        return "Unknown Unknown"
    else:
        x = results.stdout
        x = x.split(":")
        if len(x) != 7:
            return "Unknown Unknown"
        full_name = x[4]
    return full_name

def convert_date_int_to_date_str(date_int):
    # Convert integer to string
    date_str = str(date_int)

    # Insert hyphens at appropriate positions
    date_text = f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:]}"

    return date_text

def count_column_equals_value(cursor,column_name, value, table_name):

    # Execute a SELECT query to count the occurrences of the value in the column
    query = f"SELECT COUNT(*) FROM {table_name} WHERE {column_name} = ?"
    cursor.execute(query, (value,))
    result = cursor.fetchone()[0]

    return result

def count_two_columns_equals_two_values(cursor,column_name1, value1, column_name2, value2, table_name):

    # Execute a SELECT query to count the occurrences of the value in the column
    query = f"SELECT COUNT(*) FROM {table_name} WHERE {column_name1} = ? AND {column_name2} = ?"
    cursor.execute(query, (value1, value2,))
    result = cursor.fetchone()[0]

    return result

def check_value_and_return_primary_key(cursor, value, column_name, table_name, primary_key_column):

    # Execute a SELECT query to check if the value exists in the column and return the primary key
    query = f"SELECT {primary_key_column}, COUNT(*) FROM {table_name} WHERE {column_name} = ?"
    cursor.execute(query, (value,))
    result = cursor.fetchone()

    # If the count is greater than 0, the value exists in the column
    if result[1] > 0:
        return result[0]  # Return the primary key
    else:
        return None  # Value not found

def add_user(cursor,user_id,username):

    full_name = get_full_name(user_id)
    full_name = full_name.split(" ")
    last_name = full_name.pop(-1)
    first_name = full_name.pop(0)
    
    # Execute INSERT statement
    query = f"INSERT INTO users (user_id, username, first_name, last_name) VALUES (?, ?, ?, ?)"
    cursor.execute(query, (user_id, username, first_name, last_name))

def add_date(cursor,date):

    date_text = convert_date_int_to_date_str(date)
    
    # Execute INSERT statement
    query = f"INSERT INTO dates (date_int, date_text) VALUES (?, ?)"
    cursor.execute(query, (int(date), date_text))

def add_datamount(cursor,datamount):

    # Execute INSERT statement
    query = f"INSERT INTO datamounts (datamount_name) VALUES (?)"
    cursor.execute(query,(datamount,))

    # query = f"INSERT INTO datamounts (datamount_name) VALUES (\"{datamount}\")"
    # cursor.execute(query)
def add_datapoint(cursor,datamount_id,date,user_id,ninodes,nbytes):

    query = f"INSERT INTO datapoints (date_int, datamount_id, user_id, ninodes, nbytes) VALUES (?, ?, ?, ?, ?)"
    cursor.execute(query, (int(date),int(datamount_id),int(user_id),int(ninodes),int(nbytes)))

##########

def main():
    elog = textwrap.dedent(
        """\
    Version:
        {}
    Example:
        > spacesavers2_pdq_update_db -t /path/to/tsv -o /path/to/db -m datamount_name -d date
        """.format(
            __version__
        )
    )
    parser = argparse.ArgumentParser(
        description="spacesavers2_pdq_create_db: update/append date from TSV to DB file",
        epilog=elog,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "-t",
        "--tsv",
        dest="tsv",
        required=True,
        type=str,
        help="spacesavers2_pdq output TSV file",
    )
    parser.add_argument(
        "-o",
        "--database",
        dest="database",
        required=True,
        help="database file path (use spacesavers2_pdb_create_db to create if it does not exists.)",
    )
    parser.add_argument(
        "-m",
        "--datamount",
        dest="datamount",
        required=True,
        type=str,
        help="name of the datamount eg. CCBR or CCBR_Pipeliner",
    )
    parser.add_argument(
        "-d",
        "--date",
        dest="date",
        required=True,
        type=int,
        help="date in YYYYMMDD integer format",
    )
    parser.add_argument("-v", "--version", action="version", version=__version__)

    global args
    args = parser.parse_args()

    # check TSV file
    tsv = args.tsv
    tsv = Path(tsv).absolute()
    if not os.access(tsv, os.R_OK):
        exit("ERROR: {} file exists but cannot be read from".format(tsv))
    
    # check db file
    db = args.database
    db = Path(db).absolute()
    if not os.path.exists(db):
        exit("ERROR: {} does not exist. Create it using spacesavers2_pdq_create_db".format(db))
    if not os.access(db, os.W_OK):
        exit("ERROR: {} file exists but cannot be written to".format(db))

    # check date format
    if len(str(args.date)) != 8:
        exit("ERROR: date {} needs to be in format YYYYMMDD!".format(args.date))

    # Connect to the SQLite database (or create it if it doesn't exist)
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    # read TSV
    df = pandas.read_csv(args.tsv,
                        header=0,
                        sep="\t")
    
    date = int(args.date)
    datamount = args.datamount

    # check if datamount is new or known
    ndatamounts = count_column_equals_value(cursor=cursor,
                                            column_name="datamount_name",
                                            value=datamount,
                                            table_name="datamounts")
    if ndatamounts == 0: # this datamount is new... so add it
        add_datamount(cursor=cursor,datamount=datamount)
        conn.commit()
    datamount_id = check_value_and_return_primary_key(cursor=cursor,
                                                    column_name="datamount_name",
                                                    table_name="datamounts",
                                                    primary_key_column="datamount_id",
                                                    value=datamount)

    # check if date is new or known
    ndate = count_column_equals_value(cursor=cursor,column_name="date_int",table_name="dates",value=date)
    if ndate == 0: # date is new ... so add it
        add_date(cursor=cursor,date=date)
        conn.commit()
    else: # date is known ... check if datapoints exist for this date+datamount combination
        ndatapoints = count_two_columns_equals_two_values(cursor=cursor,
                                                        column_name1="date_int",
                                                        value1=date,
                                                        column_name2="datamount_id",
                                                        value2=datamount_id,
                                                        table_name="datapoints")
        if ndatapoints != 0: # data already entered for this data+datamount combo ... nothing to do
            print("WARNING: db already contains {} data points for date {} and datamount {}".format(ndatapoints,date,datamount))
            exit(0)

    count = 0
    # Iterate over rows and append data
    for row in df.itertuples(index=False):
        # ignore all_users rows
        if row.uid == 0: continue 
        
        # check if user is new or known
        nuid = count_column_equals_value(cursor=cursor,column_name="user_id",table_name="users",value=row.uid)
        if nuid == 0: # user does not exist so add user
            add_user(cursor=cursor,
                     user_id=row.uid,
                     username=row.username)
            conn.commit()
        
        # add datapoint
        add_datapoint(cursor=cursor,
                    date=date,
                    datamount_id=datamount_id,
                    user_id=row.uid,
                    ninodes=row.ninodes,
                    nbytes=row.nbytes)
        conn.commit()
        count += 1

    # Commit changes and close the connection
    conn.commit()
    conn.close()
    print(f"{count} new datapoints appended to database:{args.database}")

if __name__ == "__main__":
    main()