l3q/libl3q/daemon/sql.py

###############################################################################
#                                                                             #
#    L3q - Light, light, lightweight queue                                    #
#    Copyright (C) 2023  Marcus Pedersén marcus.pedersen@slu.se               #
#                                                                             #
#    This program is free software: you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation, either version 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program is distributed in the hope that it will be useful,          #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
#    GNU General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.    #
#                                                                             #
###############################################################################

# import this with:
# from libparaq import sql

import sqlite3
import os
import pwd
import grp
from typing import Any

'''
This module contains functionality
related to sql and interaction with
sqlite3.
'''


def init_db(conf: Any, testing: bool=False) -> sqlite3.Connection:
    '''
    Checks if database file exists,
    if not it is created with the right
    permissions and owner and group.
    WAL mode is set and Row is configured.
    If required tables does not exist they
    are created.
    Connection object to database is returned.
    '''
    if(not os.path.exists(conf.db_file) and not testing):
        open(conf.db_file, 'a').close()
        uid = pwd.getpwnam(conf.validate_user).pw_uid
        gid = grp.getgrnam(conf.validate_user).gr_gid
        os.chown(conf.db_file, uid, gid)
        os.chmod(conf.db_file, conf.validate_mode)
    
    conn = sqlite3.connect(conf.db_file, isolation_level=None)
    conn.execute('pragma journal_mode=wal;')
    conn.row_factory = sqlite3.Row
    
    conn.execute('''
CREATE TABLE IF NOT EXISTS validate_host
(
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    host_type TEXT NOT NULL,
    ip TEXT NOT NULL,
    port INT,
    hostname TEXT NOT NULL, 
    name TEXT NOT NULL,
    key TEXT NOT NULL,
    l3qd_key TEXT
);''')

    # node.state can have
    # following states:
    # Online
    # Soft Online
    # Offline
    # Register
    # Maintenance
    # Maintenance draining
    conn.execute('''
CREATE TABLE IF NOT EXISTS node
(
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    total_cpus INT,
    used_cpus INT,
    alloc_cpus INT,
    total_memory INT,
    used_memory INT,
    total_swap INT,
    used_swap INT,
    state TEXT NOT NULL,
    validate REFERENCES validate_host(id)
);''')



    # job.state can have
    # Running
    # Queued
    # Terminated
    # Canceled
    # Error
    # Depend
    conn.execute('''
CREATE TABLE IF NOT EXISTS job
(
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name TEXT, 
    type TEXT NOT NULL,
    init_date TEXT NOT NULL,
    start_date TEXT,
    end_date TEXT,
    hosts_alloc INT NOT NULL,
    cores_alloc INT NOT NULL,
    no_tasks_running INT NOT NULL,
    no_tasks_finished INT NOT NULL,
    no_tasks_err INT NOT NULL,
    nodes TEXT,
    depend TEXT,
    user TEXT NOT NULL,
    state TEXT NOT NULL
);''')


    conn.execute('''
CREATE TABLE IF NOT EXISTS task
(
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    command TEXT NOT NULL,
    seqno INT,
    workdir TEXT,
    unitname TEXT,
    start_time TEXT,
    end_time TEXT,
    mono_start_time INT,
    mono_end_time INT,
    active_state TEXT,
    sub_state TEXT,
    exit_code INT,
    main_pid INT,
    memory_peak INT,
    pids_peak INT,
    stdout TEXT,
    stderr TEXT,
    jobid REFERENCES job(id),
    nodeid REFERENCES node(id)
);''')

    
    return conn