From d4edaf6ab646c1da84e9851507c2e06b40c7b577 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleksa=20Siri=C5=A1ki?=
 <31509435+aleksasiriski@users.noreply.github.com>
Date: Fri, 13 Jan 2023 21:15:22 +0100
Subject: [PATCH] Removed duplicate code

Switch to var instead of IF ELSE for creating tables.
---
 rffmpeg | 31 +++++++++++--------------------
 1 file changed, 11 insertions(+), 20 deletions(-)

diff --git a/rffmpeg b/rffmpeg
index 0d3cfe2..c39b4b9 100755
--- a/rffmpeg
+++ b/rffmpeg
@@ -45,6 +45,7 @@ log = logging.getLogger("rffmpeg")
 # Use Postgresql if specified, otherwise use SQLite
 DB_TYPE = "SQLITE"
 SQL_VAR_SIGN = "?"
+SQL_PRIMARY_KEY="INTEGER"
 POSTGRES_DB = os.environ.get("RFFMPEG_POSTGRES_DB", "rffmpeg")
 POSTGRES_USER = os.environ.get("RFFMPEG_POSTGRES_USER")
 POSTGRES_PASS = os.environ.get("RFFMPEG_POSTGRES_PASS")
@@ -54,6 +55,7 @@ POSTGRES_HOST = os.environ.get("RFFMPEG_POSTGRES_HOST", "localhost")
 if POSTGRES_DB != None and POSTGRES_USER != None:
     DB_TYPE = "POSTGRES"
     SQL_VAR_SIGN = "%s"
+    SQL_PRIMARY_KEY="SERIAL"
     POSTGRES_CREDENTIALS = {
         "dbname": POSTGRES_DB,
         "user": POSTGRES_USER,
@@ -655,26 +657,15 @@ def run_control(config):
                 cur.execute("DROP TABLE IF EXISTS hosts")
                 cur.execute("DROP TABLE IF EXISTS processes")
                 cur.execute("DROP TABLE IF EXISTS states")
-                if DB_TYPE == "SQLITE":
-                    cur.execute(
-                        "CREATE TABLE hosts (id INTEGER PRIMARY KEY, hostname TEXT NOT NULL, weight INTEGER DEFAULT 1, servername TEXT NOT NULL)"
-                    )
-                    cur.execute(
-                        "CREATE TABLE processes (id INTEGER PRIMARY KEY, host_id INTEGER, process_id INTEGER, cmd TEXT)"
-                    )
-                    cur.execute(
-                        "CREATE TABLE states (id INTEGER PRIMARY KEY, host_id INTEGER, process_id INTEGER, state TEXT)"
-                    )
-                elif DB_TYPE == "POSTGRES":
-                    cur.execute(
-                        "CREATE TABLE hosts (id SERIAL PRIMARY KEY, hostname TEXT NOT NULL, weight INTEGER DEFAULT 1, servername TEXT NOT NULL)"
-                    )
-                    cur.execute(
-                        "CREATE TABLE processes (id SERIAL PRIMARY KEY, host_id INTEGER, process_id INTEGER, cmd TEXT)"
-                    )
-                    cur.execute(
-                        "CREATE TABLE states (id SERIAL PRIMARY KEY, host_id INTEGER, process_id INTEGER, state TEXT)"
-                    )
+                cur.execute(
+                    f"CREATE TABLE hosts (id {SQL_PRIMARY_KEY} PRIMARY KEY, hostname TEXT NOT NULL, weight INTEGER DEFAULT 1, servername TEXT NOT NULL)"
+                )
+                cur.execute(
+                    f"CREATE TABLE processes (id {SQL_PRIMARY_KEY} PRIMARY KEY, host_id INTEGER, process_id INTEGER, cmd TEXT)"
+                )
+                cur.execute(
+                    f"CREATE TABLE states (id {SQL_PRIMARY_KEY} PRIMARY KEY, host_id INTEGER, process_id INTEGER, state TEXT)"
+                )
         except Exception as e:
             fail("Failed to create database: {}".format(e))