From 1c58d5e215837a02fc850fa25563931ede953c93 Mon Sep 17 00:00:00 2001 From: Parker Date: Thu, 3 Apr 2025 17:28:53 -0500 Subject: [PATCH] Update config --- config.py | 100 +++++++++++++++++++++++++------------------- config.yaml.example | 21 ++++++++++ database.py | 2 +- 3 files changed, 79 insertions(+), 44 deletions(-) create mode 100644 config.yaml.example 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"] - 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"] + 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"] + + 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)