pvc-bootstrap/bootstrap-daemon/pvcbootstrapd/lib/db.py

#!/usr/bin/env python3

# db.py - PVC Cluster Auto-bootstrap database libraries
# Part of the Parallel Virtual Cluster (PVC) system
#
#    Copyright (C) 2018-2021 Joshua M. Boniface <joshua@boniface.me>
#
#    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, version 3.
#
#    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 <https://www.gnu.org/licenses/>.
#
###############################################################################

import os
import sqlite3
import contextlib
import json

from time import sleep
from pvcbootstrapd.lib.dataclasses import Cluster, Node


#
# Database functions
#
@contextlib.contextmanager
def dbconn(db_path):
    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA foreign_keys = 1")
    cur = conn.cursor()
    yield cur
    conn.commit()
    conn.close()


def init_database(config):
    db_path = config["database_path"]
    if not os.path.isfile(db_path):
        # Initializing the database
        with dbconn(db_path) as cur:
            # Table listing all clusters
            cur.execute(
                """CREATE TABLE clusters
                           (id INTEGER PRIMARY KEY AUTOINCREMENT,
                            name TEXT UNIQUE  NOT NULL,
                            state TEXT NOT NULL)"""
            )
            # Table listing all nodes
            # FK: cluster -> clusters.id
            cur.execute(
                """CREATE TABLE nodes
                           (id INTEGER PRIMARY KEY AUTOINCREMENT,
                            cluster INTEGER NOT NULL,
                            state TEXT NOT NULL,
                            name TEXT UNIQUE NOT NULL,
                            nodeid INTEGER NOT NULL,
                            bmc_macaddr TEXT NOT NULL,
                            bmc_ipaddr TEXT NOT NULL,
                            host_macaddr TEXT NOT NULL,
                            host_ipaddr TEXT NOT NULL,
                            CONSTRAINT cluster_col FOREIGN KEY (cluster) REFERENCES clusters(id) ON DELETE CASCADE )"""
            )


#
# Cluster functions
#
def get_cluster(config, cid=None, name=None):
    if cid is None and name is None:
        return None
    elif cid is not None:
        findfield = 'id'
        datafield = cid
    elif name is not None:
        findfield = 'name'
        datafield = name

    with dbconn(config["database_path"]) as cur:
        cur.execute(
            f"""SELECT * FROM clusters WHERE {findfield} = ?""",
            (datafield,)
        )
        rows = cur.fetchall()

    if len(rows) > 0:
        row = rows[0]
    else:
        return None
        
    return Cluster(row[0], row[1], row[2])


def add_cluster(config, name, state):
    with dbconn(config["database_path"]) as cur:
        cur.execute(
            """INSERT INTO clusters
                        (name, state)
                        VALUES
                        (?, ?)""",
            (name, state)
        )
    
    return get_cluster(config, name=name)


def update_cluster_state(config, name, state):
    with dbconn(config["database_path"]) as cur:
        cur.execute(
            """UPDATE clusters
                        SET state = ?
                        WHERE name = ?""",
            (state, name)
        )

    return get_cluster(config, name=name)


#
# Node functions
#
def get_node(config, cluster_name, nid=None, name=None, bmc_macaddr=None):
    cluster = get_cluster(config, name=cluster_name)

    if nid is None and name is None and bmc_macaddr is None:
        return None
    elif nid is not None:
        findfield = 'id'
        datafield = nid
    elif bmc_macaddr is not None:
        findfield = 'bmc_macaddr'
        datafield = bmc_macaddr
    elif name is not None:
        findfield = 'name'
        datafield = name

    with dbconn(config["database_path"]) as cur:
        cur.execute(
            f"""SELECT * FROM nodes WHERE {findfield} = ? AND cluster = ?""",
            (datafield, cluster.id)
        )
        rows = cur.fetchall()


    if len(rows) > 0:
        row = rows[0]
    else:
        return None

    return Node(row[0], cluster.name, row[2], row[3], row[4], row[5], row[6], row[7], row[8])


def get_nodes_in_cluster(config, cluster_name):
    cluster = get_cluster(config, name=cluster_name)

    with dbconn(config["database_path"]) as cur:
        cur.execute(
            """SELECT * FROM nodes WHERE cluster = ?""",
            (cluster.id, )
        )
        rows = cur.fetchall()
    
    node_list = list()
    for row in rows:
        node_list.append(
            Node(row[0], cluster.name, row[2], row[3], row[4], row[5], row[6], row[7], row[8])
        )

    return node_list


def add_node(config, cluster_name, state, name, nodeid, bmc_macaddr, bmc_ipaddr, host_macaddr, host_ipaddr):
    cluster = get_cluster(config, name=cluster_name)

    with dbconn(config["database_path"]) as cur:
        cur.execute(
            """INSERT INTO nodes
                        (cluster, state, name, nodeid, bmc_macaddr, bmc_ipaddr, host_macaddr, host_ipaddr)
                        VALUES
                        (?, ?, ?, ?, ?, ?, ?, ?)""",
            (cluster.id, state, name, nodeid, bmc_macaddr, bmc_ipaddr, host_macaddr, host_ipaddr)
        )

    return get_node(config, cluster_name, name=name)


def update_node_state(config, cluster_name, name, state):
    cluster = get_cluster(config, name=cluster_name)

    with dbconn(config["database_path"]) as cur:
        cur.execute(
            """UPDATE nodes
                        SET state = ?
                        WHERE name = ? AND cluster = ?""",
            (state, name, cluster.id)
        )

    return get_node(config, cluster_name, name=name)


def update_node_addresses(config, cluster_name, name, bmc_macaddr, bmc_ipaddr, host_macaddr, host_ipaddr):
    cluster = get_cluster(config, name=cluster_name)

    with dbconn(config["database_path"]) as cur:
        cur.execute(
            """UPDATE nodes
                        SET bmc_macaddr = ?, bmc_ipaddr = ?, host_macaddr = ?, host_ipaddr = ?
                        WHERE name = ? AND cluster = ?""",
            (bmc_macaddr, bmc_ipaddr, host_macaddr, host_ipaddr, name, cluster.id)
        )

    return get_node(config, cluster_name, name=name)