db fully handeled by python

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)
+    result=init_db()
+    print(result)
+    app.run(debug=True, host='0.0.0.0', port=5000)
+
+