db fully handeled by python

This commit is contained in:
2025-11-02 20:31:51 -08:00
parent 982b7a374d
commit 378b740d93
7 changed files with 171 additions and 61 deletions

View File

@ -21,7 +21,7 @@ extra_volumes: ""
# api service vars
api_service_name: "drive_index"
api_service_folder: "{{ service_folder }}"
api_service_exe: "{{ service_folder }}/venv/bin/python {{ service_folder }}/app.py"
api_service_exe: "{{ service_folder }}/venv/bin/python -u {{ service_folder }}/app.py"
# kiosk service vars
kiosk_service_name: "drive_check"
@ -35,5 +35,6 @@ sector_size: "512"
install_kiosk: false
sleep_time: "5"
quick_refresh: false
service_only: false
...

View File

@ -1,17 +1,75 @@
from flask import Flask, jsonify, request
import sqlite3
import json
import os
app = Flask(__name__)
db_path = '/opt/ssd_health/drive_records.db'
# init db function
def init_db():
print("Initializing DB")
db_check = "SELECT name FROM sqlite_master WHERE type='table' AND name='drive_records';"
create_table_command = """
CREATE TABLE drive_records (
id INTEGER PRIMARY KEY,
serial TEXT NOT NULL,
model TEXT NOT NULL,
flavor TEXT NOT NULL,
capacity TEXT NOT NULL,
TBW TEXT NOT NULL,
smart TEXT NOT NULL
);
"""
# this code deletes the db file if 0 bytes
if os.path.exists(db_path) and os.path.getsize(db_path) == 0:
try:
os.remove(db_path)
print("Database is 0 bytes, deleting.")
except Exception as e:
print(f"error during file deletion - 405: {e}")
return jsonify({'error during file deletion': e}), 405
try:
result = bool(query_db(db_check))
print(result)
# Check if any tables were found
if result:
print("drive_records exists - 205")
#return jsonify({'drive_records exists - 205': result}), 205
else:
print("drive_records does not exist, creating")
try:
result_init = query_db(create_table_command)
print("Database created - 201")
#return jsonify({'Database created': True, 'Result': result_init}), 201
except sqlite3.Error as e:
print(f"error during table initialization: {e}")
return jsonify({'error during table initialization - 401': e}), 401
#return len(result) > 0
except sqlite3.Error as e:
print(f"error during table check: {e}")
return jsonify({'error during table check - 400': e}), 400
# sqlite query function
def query_db(sql_query):
conn = sqlite3.connect('drive_records.db')
try:
# Establish a connection to the SQLite database using a context manager
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
print("Executing SQL query:", sql_query)
cursor.execute(sql_query)
rows = cursor.fetchall()
conn.close()
return rows
except sqlite3.Error as e:
# Handle any potential errors that might occur during database operations
print("An error occurred:", e)
return []
#conn = sqlite3.connect(db_path)
#cursor = conn.cursor()
#cursor.execute(sql_query)
#rows = cursor.fetchall()
#conn.close()
#return rows
# Function to return all drive records in database
def get_all_drive_records():
@ -33,12 +91,15 @@ def get_all_drive_records():
# Function to check if a serial number exists in the database
def check_serial_exists(serial):
return bool(query_db("SELECT * FROM drive_records WHERE serial=?", (serial,)))
serial_check = f"SELECT * FROM drive_records WHERE serial='{serial}'"
print(serial_check)
return bool(query_db(serial_check))
# Route to check if a serial number exists in the database
@app.route('/check', methods=['GET'])
def check():
serial_lookup = request.args.get('serial_lookup')
print(f"Serial to check: {serial_lookup}")
if not serial_lookup:
return jsonify({'error': 'No serial number provided'}), 400
@ -50,5 +111,44 @@ def check():
def index():
return get_all_drive_records()
# Route to add drive in database
# serial,model,flavor,capacity,TBW,smart
@app.route('/add_drive', methods=['GET'])
def add_drive():
serial = request.args.get('serial')
model = request.args.get('model')
flavor = request.args.get('flavor')
capacity = request.args.get('capacity')
TBW = request.args.get('TBW')
smart = request.args.get('smart')
if None in [serial, model, flavor, capacity, TBW, smart]:
return jsonify({'error': 'Missing required query parameter(s)'}), 400
add_drive_query = f"INSERT INTO drive_records (serial, model, flavor, capacity, TBW, smart) VALUES ('{serial}', '{model}', '{flavor}', '{capacity}', '{TBW}', '{smart}'); "
print(add_drive_query)
return jsonify(query_db(add_drive_query))
# Route to update drive in database
# serial,TBW,smart
@app.route('/update_drive', methods=['GET'])
def update_drive():
serial = request.args.get('serial')
TBW = request.args.get('TBW')
smart = request.args.get('smart')
if None in [serial, TBW, smart]:
return jsonify({'error': 'Missing required query parameter(s)'}), 400
update_drive_query = f"UPDATE drive_records SET TBW = '{TBW}', smart = '{smart}' WHERE serial = '{serial}';"
print(update_drive_query)
return jsonify(query_db(update_drive_query))
# test route
@app.route('/test', methods=['GET'])
def test():
db_check = "SELECT name FROM sqlite_master WHERE type='table' AND name='drive_records';"
return query_db(db_check)
if __name__ == '__main__':
result=init_db()
print(result)
app.run(debug=True, host='0.0.0.0', port=5000)

View File

@ -3,12 +3,11 @@
# Function to display usage
usage() {
echo "Usage: $0 [-i] [-v] [-x] -d /path/to/drive_records.db [-a 'serial,model,flavor,capacity,TBW,smart' OR -u 'serial,TBW,smart'] "
echo "Options - choose only one of a, u, or i, v and x are optional and not exclusive, and always provide the d"
echo "Usage: $0 [-v] [-x] -d /path/to/drive_records.db [-a 'serial,model,flavor,capacity,TBW,smart' OR -u 'serial,TBW,smart'] "
echo "Options - choose only one of a or u, v and x are optional and not exclusive, and always provide the d"
echo " -d /path/to/drive_records.db Specify path to DB, required"
echo " -a 'serial,model,flavor,capacity,TBW,smart' Add new drive to sqlite db"
echo " -u 'serial,TBW,smart' Update drive data in sqlite db"
echo " -i Initialize database if not present"
echo " -v Output verbose information"
echo " -x Output debug information"
exit 1
@ -60,7 +59,7 @@ CREATE_TABLE="CREATE TABLE drive_records (
);"
# Parse command line options
while getopts ":d:a:u:ivx" opt; do
while getopts ":d:a:u:vx" opt; do
case ${opt} in
v ) # process option v
echo "Be Verbose"
@ -95,13 +94,6 @@ while getopts ":d:a:u:ivx" opt; do
VALID_FLAGS=true
DRIVE_DATA=$OPTARG
;;
i ) # process option i
if [ "$BE_VERBOSE" == "true" ] ; then
echo "initialize database"
fi
VALID_FLAGS=true
init_db
;;
\? ) usage
;;
esac
@ -159,10 +151,8 @@ if [ "$ADD_DRIVE" == "true" ]; then
DRIVE_EXISTS=$(curl -s http://0.0.0.0:5000/check?serial_lookup=${data[0]} | jq .serial_number_exists)
if [ "$DRIVE_EXISTS" == "false" ]; then
# Insert the values into the database
sqlite3 "$DB_FILE" <<EOF
INSERT INTO drive_records (serial, model, flavor, capacity, TBW, smart)
VALUES ('${data[0]}', '${data[1]}', '${data[2]}', '${data[3]}', '${data[4]}', '${data[5]}');
EOF
echo "http://0.0.0.0:5000/add_drive?serial='${data[0]}'&model='${data[1]}'&flavor='${data[2]}'&capacity='${data[3]}'&TBW='${data[4]}'&smart='${data[5]}'"
curl -s "http://0.0.0.0:5000/add_drive?serial='${data[0]}'&model='${data[1]}'&flavor='${data[2]}'&capacity='${data[3]}'&TBW='${data[4]}'&smart='${data[5]}'"
else
if [ "$BE_VERBOSE" == "true" ] ; then
echo "Drive already exists, skipping"
@ -196,11 +186,8 @@ if [ "$UPDATE_DRIVE" == "true" ]; then
DRIVE_EXISTS=$(curl -s http://0.0.0.0:5000/check?serial_lookup=${data[0]} | jq .serial_number_exists)
if [ "$DRIVE_EXISTS" == "true" ]; then
# Update the values in the database
sqlite3 "$DB_FILE" <<EOF
UPDATE drive_records
SET TBW = '${data[1]}', smart = '${data[2]}'
WHERE serial = '${data[0]}';
EOF
echo "http://0.0.0.0:5000/update_drive?serial='${data[0]}'&TBW='${data[1]}'&smart='${data[2]}"
curl -s "http://0.0.0.0:5000/update_drive?serial='${data[0]}'&TBW='${data[1]}'&smart='${data[2]}"
else
if [ "$BE_VERBOSE" == "true" ] ; then
echo "Drive does not exist, skipping"

View File

@ -29,9 +29,6 @@
group: "{{ autologin_user }}"
mode: 0755
- name: Drive Index - initialize db
shell: "{{ service_folder }}/store_drive.sh -i -d {{ db_path }}"
- name: "Drive Index - template drive_check.sh"
template:
src: drive_check.sh
@ -94,7 +91,7 @@
enabled: yes
- name: Drive Index - kiosk mode handler
when: install_kiosk | bool
when: install_kiosk | bool or service_only | bool
block:
- name: Drive Index - set sleep_time to 1

View File

@ -2,7 +2,6 @@
# create and configure user account
- name: Drive health - set up user account
when: not quick_refresh | bool
include_tasks: user_setup.yaml
# create drive index service
@ -15,7 +14,7 @@
# set up autologin
- name: Drive health - configure autologin
when: not install_kiosk | bool
when: not install_kiosk | bool or not service_only | bool
include_tasks: autologin.yaml

View File

@ -1,6 +1,7 @@
---
- name: "User setup - create {{ autologin_user }} user"
when: not quick_refresh | bool
user:
name: "{{ autologin_user }}"
groups: disk
@ -8,6 +9,7 @@
shell: /bin/bash
- name: "User setup - ensure {{ autologin_user }} home folder exists"
when: not quick_refresh | bool
file:
path: "/home/{{ autologin_user }}"
state: directory
@ -15,19 +17,23 @@
group: "{{ autologin_user }}"
mode: '0700'
- name: User setup - update permissions on smartctl & fdisk
- name: User setup - autologin needed stuff
when: not service_only | bool
block:
- name: User setup - update permissions on smartctl & fdisk
shell: |
chmod 755 /usr/sbin/smartctl
chmod u+s /usr/sbin/smartctl
chmod u+s /usr/sbin/fdisk
- name: User setup - create symlink for smartctl & fdisk
- name: User setup - create symlink for smartctl & fdisk
ignore_errors: yes
shell: |
ln /usr/sbin/smartctl /usr/bin/smartctl
ln /usr/sbin/fdisk /usr/bin/fdisk
- name: "User setup - allow {{ autologin_user }} to smartctl"
- name: "User setup - allow {{ autologin_user }} to smartctl"
copy:
dest: /etc/sudoers.d/smartctl
content: |

View File

@ -42,13 +42,23 @@ while true; do
echo "$DISK has $PLR% lifetime remaining"
fi
echo
if [ -x "$TBW"] ; then
TBW="unknown"
fi
# database handler
if [ "$DRIVE_EXISTS" == "false" ] ; then
#echo "{{ service_folder }}/store_drive.sh -a '$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART' -d {{ db_path }}"
{{ service_folder }}/store_drive.sh -a "$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART" -d {{ db_path }}
H_MODEL=$(echo $MODEL | sed 's/ /%20/g')
H_FLAVOR=$(echo $FLAVOR | sed 's/ /%20/g')
H_CAPACITY=$(echo $CAPACITY | sed 's/ /%20/g')
curl -s "http://0.0.0.0:5000/add_drive?serial=$SERIAL&model=$H_MODEL&flavor=$H_FLAVOR&capacity=$H_CAPACITY&TBW=$TBW&smart=$SMART"
# curl -s "http://0.0.0.0:5000/add_drive?serial='${data[0]}'&model='${data[1]}'&flavor='${data[2]}'&capacity='${data[3]}'&TBW='${data[4]}'&smart='${data[5]}'"
# echo "{{ service_folder }}/store_drive.sh -a '$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART' -d {{ db_path }}"
# {{ service_folder }}/store_drive.sh -a "$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART" -d {{ db_path }}
else
#echo "{{ service_folder }}/store_drive.sh -u '$SERIAL,$TBW,$SMART' -d {{ db_path }}"
{{ service_folder }}/store_drive.sh -u "$SERIAL,$TBW,$SMART" -d {{ db_path }}
curl -s "http://0.0.0.0:5000/update_drive?serial=$SERIAL&TBW=$TBW&smart=$SMART"
# curl -s "http://0.0.0.0:5000/update_drive?serial='${data[0]}'&TBW='${data[1]}'&smart='${data[2]}"
# echo "{{ service_folder }}/store_drive.sh -u '$SERIAL,$TBW,$SMART' -d {{ db_path }}"
# {{ service_folder }}/store_drive.sh -u "$SERIAL,$TBW,$SMART" -d {{ db_path }}
fi
# NVMe Logic
elif [ -n "$NVME_CHECK" ] ; then
@ -72,13 +82,23 @@ while true; do
echo "TB Written: $TBW TB"
echo "NAND spare blocks: $AVAIL_SPARE"
echo
if [ -x "$TBW"] ; then
TBW="unknown"
fi
# database handler
if [ "$DRIVE_EXISTS" == "false" ] ; then
#echo "{{ service_folder }}/store_drive.sh -a '$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART' -d {{ db_path }}"
{{ service_folder }}/store_drive.sh -a "$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART" -d {{ db_path }}
H_MODEL=$(echo $MODEL | sed 's/ /%20/g')
H_FLAVOR=$(echo $FLAVOR | sed 's/ /%20/g')
H_CAPACITY=$(echo $CAPACITY | sed 's/ /%20/g')
curl -s "http://0.0.0.0:5000/add_drive?serial=$SERIAL&model=$H_MODEL&flavor=$H_FLAVOR&capacity=$H_CAPACITY&TBW=$TBW&smart=$SMART"
# curl -s "http://0.0.0.0:5000/add_drive?serial='${data[0]}'&model='${data[1]}'&flavor='${data[2]}'&capacity='${data[3]}'&TBW='${data[4]}'&smart='${data[5]}'"
# echo "{{ service_folder }}/store_drive.sh -a '$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART' -d {{ db_path }}"
# {{ service_folder }}/store_drive.sh -a "$SERIAL,$MODEL,$FLAVOR,$CAPACITY,$TBW,$SMART" -d {{ db_path }}
else
#echo "{{ service_folder }}/store_drive.sh -u '$SERIAL,$TBW,$SMART' -d {{ db_path }}"
{{ service_folder }}/store_drive.sh -u "$SERIAL,$TBW,$SMART" -d {{ db_path }}
curl -s "http://0.0.0.0:5000/update_drive?serial=$SERIAL&TBW=$TBW&smart=$SMART"
# curl -s "http://0.0.0.0:5000/update_drive?serial='${data[0]}'&TBW='${data[1]}'&smart='${data[2]}"
# echo "{{ service_folder }}/store_drive.sh -u '$SERIAL,$TBW,$SMART' -d {{ db_path }}"
# {{ service_folder }}/store_drive.sh -u "$SERIAL,$TBW,$SMART" -d {{ db_path }}
fi
fi
else