From 378b740d93c623fc0f02881900749feb73e633c4 Mon Sep 17 00:00:00 2001 From: Matt Date: Sun, 2 Nov 2025 20:31:51 -0800 Subject: [PATCH] db fully handeled by python --- defaults/main.yaml | 3 +- files/scripts/app.py | 118 ++++++++++++++++++++++++++++++++--- files/scripts/store_drive.sh | 29 +++------ tasks/drive_index.yaml | 5 +- tasks/main.yaml | 3 +- tasks/user_setup.yaml | 38 ++++++----- templates/drive_check.sh | 36 ++++++++--- 7 files changed, 171 insertions(+), 61 deletions(-) diff --git a/defaults/main.yaml b/defaults/main.yaml index 4deada4..313bd2a 100644 --- a/defaults/main.yaml +++ b/defaults/main.yaml @@ -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 ... \ No newline at end of file diff --git a/files/scripts/app.py b/files/scripts/app.py index b838767..eae7b00 100644 --- a/files/scripts/app.py +++ b/files/scripts/app.py @@ -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') - cursor = conn.cursor() - cursor.execute(sql_query) - rows = cursor.fetchall() - conn.close() - return rows + 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() + 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__': - app.run(debug=True, host='0.0.0.0', port=5000) \ No newline at end of file + result=init_db() + print(result) + app.run(debug=True, host='0.0.0.0', port=5000) + + \ No newline at end of file diff --git a/files/scripts/store_drive.sh b/files/scripts/store_drive.sh index 597bfd8..e4fda02 100644 --- a/files/scripts/store_drive.sh +++ b/files/scripts/store_drive.sh @@ -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 @@ -158,11 +150,9 @@ if [ "$ADD_DRIVE" == "true" ]; then fi 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" <