Skip to content

Move over sql commands to sqlalchemy #294

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Move over sql commands to sqlalchemy #294

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ yoyo-migrations
ruff
pre-commit
better_profanity
sqlalchemy

# api
fastapi[all] # install all to avoid random bugs
Expand Down
51 changes: 25 additions & 26 deletions scripts/flush_db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@

import os

import psycopg2
from dotenv import load_dotenv
from psycopg2 import Error
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError


def flush_database():
Expand All @@ -18,35 +18,34 @@ def flush_database():
return

try:
# Connect to database
# Connect to database using SQLAlchemy
print("📡 Connecting to database...")
connection = psycopg2.connect(DATABASE_URL, sslmode="require")
cursor = connection.cursor()

# Drop existing tables
print("🗑️ Dropping existing tables...")
drop_tables_query = """
DROP TABLE IF EXISTS submissions CASCADE;
DROP TABLE IF EXISTS leaderboard CASCADE;
DROP TABLE IF EXISTS runinfo CASCADE;
DROP TABLE IF EXISTS _yoyo_log CASCADE;
DROP TABLE IF EXISTS _yoyo_migration CASCADE;
DROP TABLE IF EXISTS _yoyo_version CASCADE;
DROP TABLE IF EXISTS yoyo_lock CASCADE;
DROP SCHEMA IF EXISTS leaderboard CASCADE;
"""
cursor.execute(drop_tables_query)
# Commit changes
connection.commit()
engine = create_engine(DATABASE_URL)

with engine.connect() as connection:
with connection.begin():
# Drop existing tables
print("🗑️ Dropping existing tables...")
drop_tables_query = text("""
DROP TABLE IF EXISTS submissions CASCADE;
DROP TABLE IF EXISTS leaderboard CASCADE;
DROP TABLE IF EXISTS runinfo CASCADE;
DROP TABLE IF EXISTS _yoyo_log CASCADE;
DROP TABLE IF EXISTS _yoyo_migration CASCADE;
DROP TABLE IF EXISTS _yoyo_version CASCADE;
DROP TABLE IF EXISTS yoyo_lock CASCADE;
DROP SCHEMA IF EXISTS leaderboard CASCADE;
""")
connection.execute(drop_tables_query)
print("✅ Database flushed and recreated successfully!")

except Error as e:
except SQLAlchemyError as e:
print(f"❌ Database error: {e}")
except Exception as e:
print(f"❌ Unexpected error: {e}")
finally:
if "connection" in locals():
cursor.close()
connection.close()
print("🔌 Database connection closed")
print("🔌 Database operation completed")


if __name__ == "__main__":
Expand Down
195 changes: 101 additions & 94 deletions scripts/update_leaderboard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@
import os
from datetime import datetime

import psycopg2
import requests
from jinja2 import Template
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError

TOKEN = os.environ.get("DISCORD_DUMMY_TOKEN")

Expand Down Expand Up @@ -80,102 +81,108 @@ def get_name_from_id(user_id: str) -> str:
def fetch_leaderboard_data():
print("Fetching data from database...")
try:
with psycopg2.connect(DATABASE_URL) as conn:
with conn.cursor() as cur:
cur.execute(
"""
SELECT id, name, deadline
FROM leaderboard.leaderboard
"""
engine = create_engine(DATABASE_URL)
with engine.connect() as connection:
# Get all leaderboards
leaderboards_query = text("""
SELECT id, name, deadline
FROM leaderboard.leaderboard
""")

leaderboards_result = connection.execute(leaderboards_query)
leaderboards = leaderboards_result.fetchall()

# Get active leaderboards with their GPU types and submission counts
submissions_query = text("""
WITH unique_best_submissions AS (
SELECT DISTINCT ON (s.user_id)
s.file_name,
s.user_id,
s.submission_time,
r.score,
r.runner
FROM leaderboard.runs r
JOIN leaderboard.submission s ON r.submission_id = s.id
JOIN leaderboard.leaderboard l ON s.leaderboard_id = l.id
WHERE l.name = :leaderboard_name AND r.runner = :gpu_type AND NOT r.secret
AND r.score IS NOT NULL AND r.passed
ORDER BY s.user_id, r.score ASC
)

leaderboards = cur.fetchall()

# Get active leaderboards with their GPU types and submission counts
query = """
WITH unique_best_submissions AS (
SELECT DISTINCT ON (s.user_id)
s.file_name,
s.user_id,
s.submission_time,
r.score,
r.runner
FROM leaderboard.runs r
JOIN leaderboard.submission s ON r.submission_id = s.id
JOIN leaderboard.leaderboard l ON s.leaderboard_id = l.id
WHERE l.name = %s AND r.runner = %s AND NOT r.secret
AND r.score IS NOT NULL AND r.passed
ORDER BY s.user_id, r.score ASC
SELECT
file_name,
user_id,
submission_time,
score,
runner,
ROW_NUMBER() OVER (ORDER BY score ASC) as rank
FROM unique_best_submissions
ORDER BY score ASC;
""")

gpu_type_data = {}
for _lb_id, name, deadline in leaderboards:
# Get GPU types for this leaderboard
gpu_types_query = text("""
SELECT gpu_type
FROM leaderboard.gpu_type
WHERE leaderboard_id = :leaderboard_id
""")

gpu_types_result = connection.execute(gpu_types_query, {'leaderboard_id': _lb_id})
gpu_types = [row[0] for row in gpu_types_result.fetchall()]

for gpu_type in gpu_types:
submissions_result = connection.execute(
submissions_query,
{'leaderboard_name': name, 'gpu_type': gpu_type}
)
SELECT
file_name,
user_id,
submission_time,
score,
runner,
ROW_NUMBER() OVER (ORDER BY score ASC) as rank
FROM unique_best_submissions
ORDER BY score ASC;
"""

gpu_type_data = {}
for (
_lb_id,
name,
deadline,
) in leaderboards:
cur.execute(
"SELECT * from leaderboard.gpu_type where leaderboard_id = %s", [_lb_id]
submissions = submissions_result.fetchall()

print(
f"Found {len(submissions)} active submissions in {name} for {gpu_type}"
)
gpu_types = [x[1] for x in cur.fetchall()]

for gpu_type in gpu_types:
args = (name, gpu_type)
cur.execute(query, args)
submissions = cur.fetchall()

print(
f"Found {len(submissions)} active submissions in {name} for {gpu_type}"
)

if len(submissions) > 0:
if gpu_type not in gpu_type_data:
gpu_type_data[gpu_type] = {}

gpu_submissions = []
for lb in submissions:
user_id = lb[1]
time = lb[3]
rank = lb[5]
global_name = get_name_from_id(user_id)
gpu_submissions.append(
{
"user": f"{global_name}",
"time": f"{time:.9f}",
"rank": rank,
}
)

# Sort submissions by time
gpu_submissions.sort(key=lambda x: float(x["time"]))

gpu_type_data[gpu_type][name] = {
"name": name,
"deadline": deadline.strftime("%Y-%m-%d %H:%M"),
"submissions": gpu_submissions,
}

# Convert to final format
formatted_data = {
"gpu_types": [
{"name": gpu_type, "problems": list(problems.values())}
for gpu_type, problems in gpu_type_data.items()
],
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC"),
}

print("Data fetched successfully")
return formatted_data

if len(submissions) > 0:
if gpu_type not in gpu_type_data:
gpu_type_data[gpu_type] = {}

gpu_submissions = []
for lb in submissions:
user_id = lb[1]
time = lb[3]
rank = lb[5]
global_name = get_name_from_id(user_id)
gpu_submissions.append(
{
"user": f"{global_name}",
"time": f"{time:.9f}",
"rank": rank,
}
)

# Sort submissions by time
gpu_submissions.sort(key=lambda x: float(x["time"]))

gpu_type_data[gpu_type][name] = {
"name": name,
"deadline": deadline.strftime("%Y-%m-%d %H:%M"),
"submissions": gpu_submissions,
}

# Convert to final format
formatted_data = {
"gpu_types": [
{"name": gpu_type, "problems": list(problems.values())}
for gpu_type, problems in gpu_type_data.items()
],
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC"),
}

print("Data fetched successfully")
return formatted_data
except SQLAlchemyError as e:
print(f"Database error: {str(e)}")
raise
except Exception as e:
print(f"Error fetching data: {str(e)}")
raise
Expand Down
29 changes: 17 additions & 12 deletions src/discord-cluster-manager/cogs/misc_cog.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@
from typing import TYPE_CHECKING

import discord
import psycopg2
from discord import app_commands
from discord.ext import commands
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError
from env import DATABASE_URL
from utils import send_discord_message, setup_logging

Expand Down Expand Up @@ -33,17 +34,21 @@ async def verify_db(self, interaction: discord.Interaction):
return

try:
with psycopg2.connect(DATABASE_URL, sslmode="require") as conn:
with conn.cursor() as cursor:
cursor.execute("SELECT RANDOM()")
result = cursor.fetchone()
if result:
random_value = result[0]
await send_discord_message(
interaction, f"Your lucky number is {random_value}."
)
else:
await send_discord_message(interaction, "No result returned.")
engine = create_engine(DATABASE_URL)
with engine.connect() as connection:
result = connection.execute(text("SELECT RANDOM()"))
row = result.fetchone()
if row:
random_value = row[0]
await send_discord_message(
interaction, f"Your lucky number is {random_value}."
)
else:
await send_discord_message(interaction, "No result returned.")
except SQLAlchemyError as e:
message = "Database error occurred"
logger.error(f"{message}: {str(e)}", exc_info=True)
await send_discord_message(interaction, f"{message}.")
except Exception as e:
message = "Error interacting with the database"
logger.error(f"{message}: {str(e)}", exc_info=True)
Expand Down
Loading