From 5a40681ac641158f047d87084702549473026ab2 Mon Sep 17 00:00:00 2001
From: Yasen Pramatarov <yasen@lindeas.com>
Date: Mon, 14 Oct 2024 18:23:09 +0300
Subject: [PATCH] Adds check for table existence

---
 .gitignore |   1 +
 main.go    |  40 ++++++++++++++++++++++++++++++++++++++--
 meet.db    | Bin 12288 -> 0 bytes
 3 files changed, 39 insertions(+), 2 deletions(-)
 delete mode 100644 meet.db

diff --git a/.gitignore b/.gitignore
index 10297ce..0dfcbd2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 jilo-server
+*.db
diff --git a/main.go b/main.go
index 04668ce..b0f057a 100644
--- a/main.go
+++ b/main.go
@@ -2,9 +2,12 @@ package main
 
 import (
     "database/sql"
+    "flag"
+    "fmt"
     "io/ioutil"
     "log"
     "net/http"
+    "os"
     "time"
 
     _ "github.com/mattn/go-sqlite3"
@@ -43,7 +46,7 @@ func readConfig(filePath string) Config {
     return config
 }
 
-func setupDatabase(dbPath string) (*sql.DB, error) {
+func setupDatabase(dbPath string, initDB bool) (*sql.DB, error) {
 
     // Open the database
     db, err := sql.Open("sqlite3", dbPath)
@@ -51,6 +54,20 @@ func setupDatabase(dbPath string) (*sql.DB, error) {
         return nil, err
     }
 
+    // Check if the table exists
+    tableExists := checkTableExists(db)
+    if !tableExists && !initDB {
+        // Ask if we should create the table
+        fmt.Print("Table not found. Do you want to create it? (y/n): ")
+        var response string
+        fmt.Scanln(&response)
+
+        if response != "y" && response != "Y" {
+            log.Println("Exiting because the table is missing, but mandatory.")
+            os.Exit(1)
+        }
+    }
+
     // If the table is not there, initialize it
     createTable := `
     CREATE TABLE IF NOT EXISTS endpoint_data (
@@ -67,6 +84,21 @@ func setupDatabase(dbPath string) (*sql.DB, error) {
     return db, nil
 }
 
+// Check for the table
+func checkTableExists(db *sql.DB) bool {
+    sql := `
+    SELECT name
+        FROM sqlite_master
+        WHERE type='table'
+        AND name='endpoint_data';`
+    row := db.QueryRow(sql)
+
+    var name string
+    err := row.Scan(&name)
+
+    return err == nil && name == "endpoint_data"
+}
+
 func checkEndpoint(endpoint string) (int, int64) {
     log.Println("Sending HTTP get request to Jilo agent:", endpoint)
     start := time.Now()
@@ -94,13 +126,17 @@ func main() {
     // First flush all the logs
     log.SetFlags(log.LstdFlags | log.Lshortfile)
 
+    // Init the DB, "--init-db" creates the table
+    initDB := flag.Bool("init-db", false, "Create database table if not present without prompting")
+    flag.Parse()
+
     // Config file
     log.Println("Reading the config file...")
     config := readConfig("jilo-server.conf")
 
     // Connect to or setup the database
     log.Println("Initializing the database...")
-    db, err := setupDatabase(config.DatabasePath)
+    db, err := setupDatabase(config.DatabasePath, *initDB)
     if err != nil {
         log.Fatal("Failed to initialize the database:", err)
     }
diff --git a/meet.db b/meet.db
deleted file mode 100644
index 110dd8aad978e703c063164d774d4aa03904e150..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 12288
zcmeI&u};E37zgmX0-Bg;)x@wc+y(?N1)K<j6;4A0#8L*PrqVMZ(iUi0+>DE>FX3}I
z_$2PWgI6LF7%(w$G5;^U_O9P|cg=6Rq}}W0Fcj?8pSpI)%4DBVN=_Lggd}8($+ie8
zJMqOpKBMtZks!zSlU-S*?7b58P~H#_fB*y_009U<00Izz00bZafnNx8l&xegNBg0D
zGZMjMBoDJzASN^64c3RTn!$CGGgCina<*2R@oZP*<;?8E9o<f*)AaFZzNYY;u|M=e
z%dtayWuUCAw~}I+9rhJNhc((Jzu*Sz7>$;0^w}lvGrebCHQI887H^yR@9J>qiXgPz
zF{{f@H5)C?>ik^qHBDCQ8HTL2qL*D$Z*`Wl<cx3@SOeb?i_MnFsR+it7YHj_+j4_N
z#eC0fbT*HxPegrJ-_#e`K|lZk5P$##AOHafKmY;|fB*y_@Mi+4c!HD_rKFWk3Pr6@
zEVGhUDV8f*i9XY6Dz+h;Nhuq$)fC;3%|z#as(ui8As_$&2tWV=5P$##AOHafKmY;|
Q_zwaHG?u0-KPtpO0RX6otN;K2