Update config

3 files changed, 79 insertions, 44 deletions

diff --git a/config.py b/config.py
index 720ecdd..5448110 100644
--- a/config.py
+++ b/config.py
@@ -27,6 +27,7 @@ LOG.addHandler(stream)
 TOKEN = None
 BOT_COLOR = None
 
+SQLITE_NAME = "disarchive"
 DB_NAME = None
 DB_ENGINE = None
 DB_HOST = None
@@ -46,11 +47,34 @@ schema = {
             },
             "required": ["token"],
         },
-        "database": {
+        "sqlite": {  # TODO - turn SQLITE path and name into global variables
+            "type": "object",
+            "properties": {
+                "name": {"type": "string", "default": "disarchive"},
+            },
+            "required": ["name"],
+        },
+        "mysql": {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string", "default": "disarchive"},
+                "host": {"type": "string"},
+                "port": {"type": "integer"},
+                "user": {"type": "string"},
+                "password": {"type": "string"},
+            },
+            "required": [
+                "name",
+                "host",
+                "port",
+                "user",
+                "password",
+            ],
+        },
+        "postgresql": {
             "type": "object",
             "properties": {
                 "name": {"type": "string", "default": "disarchive"},
-                "engine": {"type": "string"},
                 "host": {"type": "string"},
                 "port": {"type": "integer"},
                 "user": {"type": "string"},
@@ -58,7 +82,6 @@ schema = {
             },
             "required": [
                 "name",
-                "engine",
                 "host",
                 "port",
                 "user",
@@ -66,7 +89,7 @@ schema = {
             ],
         },
     },
-    "required": ["bot_info", "database"],
+    "required": ["bot_info"],
 }
 
 
@@ -83,35 +106,17 @@ def load_config():
             validate_config(file_contents)
 
     except FileNotFoundError:
-        # Create a new config.yaml file with the template
-        with open(file_path, "w") as f:
-            f.write(
-                """
-bot_info:
-    token:
-    bot_color:
-
-database:
-    name:
-    engine:
-    host:
-    port:
-    user:
-    password: """
-            )
-
         sys.exit(
             LOG.critical(
-                "Configuration file `config.yaml` has been generated. Please"
-                " fill out all of the necessary information. Refer to the docs"
-                " for information on what a specific configuration option is."
+                "config.yaml file not found. Please use the template to"
+                " configure the bot."
             )
         )
 
 
 # Thouroughly validate all of the options in the config.yaml file
 def validate_config(file_contents):
-    global TOKEN, BOT_COLOR, DB_NAME, DB_ENGINE, DB_HOST, DB_PORT, DB_USER, DB_PASSWORD
+    global TOKEN, BOT_COLOR, SQLITE_NAME, DB_NAME, DB_ENGINE, DB_HOST, DB_PORT, DB_USER, DB_PASSWORD
     config = yaml.safe_load(file_contents)
 
     try:
@@ -123,37 +128,46 @@ def validate_config(file_contents):
     hex_pattern_one = "^#([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})$"
     hex_pattern_two = "^([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})$"
 
+    # Check if the bot_color is a valid hex color
     if "bot_color" in config["bot_info"]:
         if not bool(
             re.match(hex_pattern_one, config["bot_info"]["bot_color"])
         ) and not bool(
             re.match(hex_pattern_two, config["bot_info"]["bot_color"])
         ):
-            LOG.critical(
-                "Error in config.yaml file: bot_color is not a valid hex color"
+            LOG.warn(
+                "bot_color is not a valid hex color... defaulting to #26dfc9"
             )
         else:
             BOT_COLOR = discord.Color(
                 int((config["bot_info"]["bot_color"]).replace("#", ""), 16)
             )
 
-    if config["database"]["engine"] not in [
-        "sqlite",
-        "mysql",
-        "postgresql",
-    ]:
-        LOG.error(
-            "database_engine must be either 'sqlite', 'mysql', or 'postgresql'"
-        )
-        return False
-    else:
-        DB_ENGINE = config["database"]["engine"]
+    # Assign database variables
+    if "sqlite" in config:
+        DB_ENGINE = "sqlite"
+        if "name" in config["sqlite"]:
+            SQLITE_NAME = config["sqlite"]["name"]
+
+    elif "mysql" in config:
+        DB_ENGINE = "mysql"
+        DB_NAME = config["mysql"]["name"]
+        DB_HOST = config["mysql"]["host"]
+        DB_PORT = config["mysql"]["port"]
+        DB_USER = config["mysql"]["user"]
+        DB_PASSWORD = config["mysql"]["password"]
+
+    elif "postgresql" in config:
+        DB_ENGINE = "postgresql"
+        DB_NAME = config["postgresql"]["name"]
+        DB_HOST = config["postgresql"]["host"]
+        DB_PORT = config["postgresql"]["port"]
+        DB_USER = config["postgresql"]["user"]
+        DB_PASSWORD = config["postgresql"]["password"]
 
-    DB_NAME = config["database"]["name"]
-    DB_HOST = config["database"]["host"]
-    DB_PORT = config["database"]["port"]
-    DB_USER = config["database"]["user"]
-    DB_PASSWORD = config["database"]["password"]
+    else:
+        LOG.warn("No database engine specified. Defaulting to SQLite.")
+        DB_ENGINE = "sqlite"
 
     TOKEN = config["bot_info"]["token"]
 
diff --git a/config.yaml.example b/config.yaml.example
new file mode 100644
index 0000000..911e52f
--- /dev/null
+++ b/config.yaml.example
@@ -0,0 +1,21 @@
+bot_info:
+  token: "BOT TOKEN"
+  bot_color: 26dfc9 #optional - default is 26dfc9
+
+# Comment out if not using SQLite, and uncomment preferred database
+sqlite:
+  name: disarchive
+
+# mysql:
+#   name: disarchive
+#   host: 127.0.0.1
+#   port: 3306
+#   user: username
+#   password: pass123
+
+# postgresql:
+#   name: disarchive
+#   host: 127.0.0.1
+#   port: 5432
+#   user: username
+#   password: pass123
\ No newline at end of file
diff --git a/database.py b/database.py
index 1258e42..212d423 100644
--- a/database.py
+++ b/database.py
@@ -14,7 +14,7 @@ elif config.DB_ENGINE == "postgresql":
 else:
     if not os.path.exists("data"):
         os.makedirs("data")
-    database_url = "sqlite:///data/data.db"
+    database_url = f"sqlite:///data/{config.SQLITE_NAME}.db"
 
 engine = create_engine(database_url)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)