#!/usr/bin/env python import os import re import socket import sys import time from typing import List import django from django.contrib.auth import get_user_model from django.core.management import call_command, execute_from_command_line DEFAULT_ENVS = { "DJANGO_SETTINGS_MODULE": "pydis_site.settings", "SUPER_USERNAME": "admin", "SUPER_PASSWORD": "admin", "DEFAULT_BOT_API_KEY": "badbot13m0n8f570f942013fc818f234916ca531", } try: import dotenv dotenv.load_dotenv() except ModuleNotFoundError: pass for key, value in DEFAULT_ENVS.items(): os.environ.setdefault(key, value) class SiteManager: """ Manages the preparation and serving of the website. Handles both development and production environments. Usage: manage.py run [option]... Options: --debug Runs a development server with debug mode enabled. --silent Sets minimal console output. --verbose Sets verbose console output. """ def __init__(self, args: List[str]): self.debug = "--debug" in args self.silent = "--silent" in args if self.silent: self.verbosity = 0 else: self.verbosity = 2 if "--verbose" in args else 1 if self.debug: os.environ.setdefault("DEBUG", "true") print("Starting in debug mode.") @staticmethod def create_superuser() -> None: """Create a default django admin super user in development environments.""" print("Creating a superuser.") name = os.environ["SUPER_USERNAME"] password = os.environ["SUPER_PASSWORD"] bot_token = os.environ["DEFAULT_BOT_API_KEY"] user = get_user_model() # Get or create admin superuser. if user.objects.filter(username=name).exists(): user = user.objects.get(username=name) print('Admin superuser already exists.') else: user = user.objects.create_superuser(name, '', password) print('Admin superuser created.') # Setup a default bot token to connect with site API from rest_framework.authtoken.models import Token token, is_new = Token.objects.update_or_create(user=user) if token.key != bot_token: token.delete() token, is_new = Token.objects.update_or_create(user=user, key=bot_token) if is_new: print(f"New bot token created: {token}") else: print(f"Existing bot token found: {token}") @staticmethod def wait_for_postgres() -> None: """Wait for the PostgreSQL database specified in DATABASE_URL.""" print("Waiting for PostgreSQL database.") # Get database URL based on environmental variable passed in compose database_url = os.environ["DATABASE_URL"] match = re.search(r"@([\w.]+):(\d+)/", database_url) if not match: raise OSError("Valid DATABASE_URL environmental variable not found.") domain = match.group(1) port = int(match.group(2)) # Attempt to connect to the database socket s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) attempts_left = 10 while attempts_left: try: # Ignore 'incomplete startup packet' s.connect((domain, port)) s.shutdown(socket.SHUT_RDWR) print("Database is ready.") break except socket.error: attempts_left -= 1 print("Not ready yet, retrying.") time.sleep(0.5) else: print("Database could not be found, exiting.") sys.exit(1) @staticmethod def set_dev_site_name() -> None: """Set the development site domain in admin from default example.""" # import Site model now after django setup from django.contrib.sites.models import Site query = Site.objects.filter(id=1) site = query.get() if site.domain == "example.com": query.update( domain="pythondiscord.local:8000", name="pythondiscord.local:8000" ) @staticmethod def run_metricity_init() -> None: """ Initilise metricity relations and populate with some testing data. This is done at run time since other projects, like Python bot, rely on the site initilising it's own db, since they do not have access to the init.sql file to mount a docker-compose volume. """ import psycopg2 from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT from urllib.parse import urlsplit # The database URL has already been validated in `wait_for_postgres()` db_url_parts = urlsplit(os.environ["DATABASE_URL"]) db_connection_kwargs = { "host": db_url_parts.hostname, "port": db_url_parts.port, "user": db_url_parts.username, "password": db_url_parts.password, } # Connect to pysite first to create metricity db conn = psycopg2.connect( database=db_url_parts.path[1:], **db_connection_kwargs ) conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) with conn.cursor() as cursor: cursor.execute("SELECT 1 FROM pg_catalog.pg_database WHERE datname = 'metricity'") exists = cursor.fetchone() if exists: # Assume metricity is already populated if it exists return print("Creating metricity relations and populating with some data.") cursor.execute("CREATE DATABASE metricity") # Switch connection to metricity and initialise some data conn = psycopg2.connect( database="metricity", **db_connection_kwargs ) with conn.cursor() as cursor, open("postgres/init.sql", encoding="utf-8") as f: cursor.execute(f.read()) def prepare_server(self) -> None: """Perform preparation tasks before running the server.""" self.wait_for_postgres() if self.debug: self.run_metricity_init() django.setup() print("Applying migrations.") call_command("migrate", verbosity=self.verbosity) if self.debug: # In Production, collectstatic is ran in the Docker image print("Collecting static files.") call_command( "collectstatic", interactive=False, clear=True, verbosity=self.verbosity - 1 ) self.set_dev_site_name() self.create_superuser() def run_server(self) -> None: """Prepare and run the web server.""" in_reloader = os.environ.get('RUN_MAIN') == 'true' # Prevent preparing twice when in dev mode due to reloader if not self.debug or in_reloader: self.prepare_server() print("Starting server.") # Run the development server if self.debug: call_command("runserver", "0.0.0.0:8000") return # Import gunicorn only if we aren't in debug mode. import gunicorn.app.wsgiapp # Patch the arguments for gunicorn sys.argv = [ "gunicorn", "--preload", "-b", "0.0.0.0:8000", "pydis_site.wsgi:application", "-w", "2", "--statsd-host", "graphite.default.svc.cluster.local:8125", "--statsd-prefix", "site", "--config", "file:gunicorn.conf.py" ] # Run gunicorn for the production server. gunicorn.app.wsgiapp.run() def main() -> None: """Entry point for Django management script.""" # Always run metricity init in CI in_ci = os.environ.get("CI", "false").lower() == "true" if in_ci: SiteManager.wait_for_postgres() SiteManager.run_metricity_init() # Use the custom site manager for launching the server if len(sys.argv) > 1 and sys.argv[1] == "run": SiteManager(sys.argv).run_server() # Pass any others directly to standard management commands else: execute_from_command_line(sys.argv) if __name__ == '__main__': main()