Skip to content

Commit 5a57a41

Browse files
PaliCPaliC
committed
Move over sql commands to sqlalchemy
1 parent c7cf3f3 commit 5a57a41

5 files changed

Lines changed: 646 additions & 570 deletions

File tree

‎requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ yoyo-migrations
1010
ruff
1111
pre-commit
1212
better_profanity
13+
sqlalchemy
1314

1415
# api
1516
fastapi[all] # install all to avoid random bugs

‎scripts/flush_db.py‎

Lines changed: 25 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22

33
import os
44

5-
import psycopg2
65
from dotenv import load_dotenv
7-
from psycopg2 import Error
6+
from sqlalchemy import create_engine, text
7+
from sqlalchemy.exc import SQLAlchemyError
88

99

1010
def flush_database():
@@ -18,35 +18,34 @@ def flush_database():
1818
return
1919

2020
try:
21-
# Connect to database
21+
# Connect to database using SQLAlchemy
2222
print("📡 Connecting to database...")
23-
connection = psycopg2.connect(DATABASE_URL, sslmode="require")
24-
cursor = connection.cursor()
25-
26-
# Drop existing tables
27-
print("🗑️ Dropping existing tables...")
28-
drop_tables_query = """
29-
DROP TABLE IF EXISTS submissions CASCADE;
30-
DROP TABLE IF EXISTS leaderboard CASCADE;
31-
DROP TABLE IF EXISTS runinfo CASCADE;
32-
DROP TABLE IF EXISTS _yoyo_log CASCADE;
33-
DROP TABLE IF EXISTS _yoyo_migration CASCADE;
34-
DROP TABLE IF EXISTS _yoyo_version CASCADE;
35-
DROP TABLE IF EXISTS yoyo_lock CASCADE;
36-
DROP SCHEMA IF EXISTS leaderboard CASCADE;
37-
"""
38-
cursor.execute(drop_tables_query)
39-
# Commit changes
40-
connection.commit()
23+
engine = create_engine(DATABASE_URL)
24+
25+
with engine.connect() as connection:
26+
with connection.begin():
27+
# Drop existing tables
28+
print("🗑️ Dropping existing tables...")
29+
drop_tables_query = text("""
30+
DROP TABLE IF EXISTS submissions CASCADE;
31+
DROP TABLE IF EXISTS leaderboard CASCADE;
32+
DROP TABLE IF EXISTS runinfo CASCADE;
33+
DROP TABLE IF EXISTS _yoyo_log CASCADE;
34+
DROP TABLE IF EXISTS _yoyo_migration CASCADE;
35+
DROP TABLE IF EXISTS _yoyo_version CASCADE;
36+
DROP TABLE IF EXISTS yoyo_lock CASCADE;
37+
DROP SCHEMA IF EXISTS leaderboard CASCADE;
38+
""")
39+
connection.execute(drop_tables_query)
40+
4141
print("✅ Database flushed and recreated successfully!")
4242

43-
except Error as e:
43+
except SQLAlchemyError as e:
4444
print(f"❌ Database error: {e}")
45+
except Exception as e:
46+
print(f"❌ Unexpected error: {e}")
4547
finally:
46-
if "connection" in locals():
47-
cursor.close()
48-
connection.close()
49-
print("🔌 Database connection closed")
48+
print("🔌 Database operation completed")
5049

5150

5251
if __name__ == "__main__":

‎scripts/update_leaderboard.py‎

Lines changed: 101 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22
import os
33
from datetime import datetime
44

5-
import psycopg2
65
import requests
76
from jinja2 import Template
7+
from sqlalchemy import create_engine, text
8+
from sqlalchemy.exc import SQLAlchemyError
89

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

@@ -80,102 +81,108 @@ def get_name_from_id(user_id: str) -> str:
8081
def fetch_leaderboard_data():
8182
print("Fetching data from database...")
8283
try:
83-
with psycopg2.connect(DATABASE_URL) as conn:
84-
with conn.cursor() as cur:
85-
cur.execute(
86-
"""
87-
SELECT id, name, deadline
88-
FROM leaderboard.leaderboard
89-
"""
84+
engine = create_engine(DATABASE_URL)
85+
with engine.connect() as connection:
86+
# Get all leaderboards
87+
leaderboards_query = text("""
88+
SELECT id, name, deadline
89+
FROM leaderboard.leaderboard
90+
""")
91+
92+
leaderboards_result = connection.execute(leaderboards_query)
93+
leaderboards = leaderboards_result.fetchall()
94+
95+
# Get active leaderboards with their GPU types and submission counts
96+
submissions_query = text("""
97+
WITH unique_best_submissions AS (
98+
SELECT DISTINCT ON (s.user_id)
99+
s.file_name,
100+
s.user_id,
101+
s.submission_time,
102+
r.score,
103+
r.runner
104+
FROM leaderboard.runs r
105+
JOIN leaderboard.submission s ON r.submission_id = s.id
106+
JOIN leaderboard.leaderboard l ON s.leaderboard_id = l.id
107+
WHERE l.name = :leaderboard_name AND r.runner = :gpu_type AND NOT r.secret
108+
AND r.score IS NOT NULL AND r.passed
109+
ORDER BY s.user_id, r.score ASC
90110
)
91-
92-
leaderboards = cur.fetchall()
93-
94-
# Get active leaderboards with their GPU types and submission counts
95-
query = """
96-
WITH unique_best_submissions AS (
97-
SELECT DISTINCT ON (s.user_id)
98-
s.file_name,
99-
s.user_id,
100-
s.submission_time,
101-
r.score,
102-
r.runner
103-
FROM leaderboard.runs r
104-
JOIN leaderboard.submission s ON r.submission_id = s.id
105-
JOIN leaderboard.leaderboard l ON s.leaderboard_id = l.id
106-
WHERE l.name = %s AND r.runner = %s AND NOT r.secret
107-
AND r.score IS NOT NULL AND r.passed
108-
ORDER BY s.user_id, r.score ASC
111+
SELECT
112+
file_name,
113+
user_id,
114+
submission_time,
115+
score,
116+
runner,
117+
ROW_NUMBER() OVER (ORDER BY score ASC) as rank
118+
FROM unique_best_submissions
119+
ORDER BY score ASC;
120+
""")
121+
122+
gpu_type_data = {}
123+
for _lb_id, name, deadline in leaderboards:
124+
# Get GPU types for this leaderboard
125+
gpu_types_query = text("""
126+
SELECT gpu_type
127+
FROM leaderboard.gpu_type
128+
WHERE leaderboard_id = :leaderboard_id
129+
""")
130+
131+
gpu_types_result = connection.execute(gpu_types_query, {'leaderboard_id': _lb_id})
132+
gpu_types = [row[0] for row in gpu_types_result.fetchall()]
133+
134+
for gpu_type in gpu_types:
135+
submissions_result = connection.execute(
136+
submissions_query,
137+
{'leaderboard_name': name, 'gpu_type': gpu_type}
109138
)
110-
SELECT
111-
file_name,
112-
user_id,
113-
submission_time,
114-
score,
115-
runner,
116-
ROW_NUMBER() OVER (ORDER BY score ASC) as rank
117-
FROM unique_best_submissions
118-
ORDER BY score ASC;
119-
"""
120-
121-
gpu_type_data = {}
122-
for (
123-
_lb_id,
124-
name,
125-
deadline,
126-
) in leaderboards:
127-
cur.execute(
128-
"SELECT * from leaderboard.gpu_type where leaderboard_id = %s", [_lb_id]
139+
submissions = submissions_result.fetchall()
140+
141+
print(
142+
f"Found {len(submissions)} active submissions in {name} for {gpu_type}"
129143
)
130-
gpu_types = [x[1] for x in cur.fetchall()]
131-
132-
for gpu_type in gpu_types:
133-
args = (name, gpu_type)
134-
cur.execute(query, args)
135-
submissions = cur.fetchall()
136-
137-
print(
138-
f"Found {len(submissions)} active submissions in {name} for {gpu_type}"
139-
)
140-
141-
if len(submissions) > 0:
142-
if gpu_type not in gpu_type_data:
143-
gpu_type_data[gpu_type] = {}
144-
145-
gpu_submissions = []
146-
for lb in submissions:
147-
user_id = lb[1]
148-
time = lb[3]
149-
rank = lb[5]
150-
global_name = get_name_from_id(user_id)
151-
gpu_submissions.append(
152-
{
153-
"user": f"{global_name}",
154-
"time": f"{time:.9f}",
155-
"rank": rank,
156-
}
157-
)
158-
159-
# Sort submissions by time
160-
gpu_submissions.sort(key=lambda x: float(x["time"]))
161-
162-
gpu_type_data[gpu_type][name] = {
163-
"name": name,
164-
"deadline": deadline.strftime("%Y-%m-%d %H:%M"),
165-
"submissions": gpu_submissions,
166-
}
167-
168-
# Convert to final format
169-
formatted_data = {
170-
"gpu_types": [
171-
{"name": gpu_type, "problems": list(problems.values())}
172-
for gpu_type, problems in gpu_type_data.items()
173-
],
174-
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC"),
175-
}
176-
177-
print("Data fetched successfully")
178-
return formatted_data
144+
145+
if len(submissions) > 0:
146+
if gpu_type not in gpu_type_data:
147+
gpu_type_data[gpu_type] = {}
148+
149+
gpu_submissions = []
150+
for lb in submissions:
151+
user_id = lb[1]
152+
time = lb[3]
153+
rank = lb[5]
154+
global_name = get_name_from_id(user_id)
155+
gpu_submissions.append(
156+
{
157+
"user": f"{global_name}",
158+
"time": f"{time:.9f}",
159+
"rank": rank,
160+
}
161+
)
162+
163+
# Sort submissions by time
164+
gpu_submissions.sort(key=lambda x: float(x["time"]))
165+
166+
gpu_type_data[gpu_type][name] = {
167+
"name": name,
168+
"deadline": deadline.strftime("%Y-%m-%d %H:%M"),
169+
"submissions": gpu_submissions,
170+
}
171+
172+
# Convert to final format
173+
formatted_data = {
174+
"gpu_types": [
175+
{"name": gpu_type, "problems": list(problems.values())}
176+
for gpu_type, problems in gpu_type_data.items()
177+
],
178+
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC"),
179+
}
180+
181+
print("Data fetched successfully")
182+
return formatted_data
183+
except SQLAlchemyError as e:
184+
print(f"Database error: {str(e)}")
185+
raise
179186
except Exception as e:
180187
print(f"Error fetching data: {str(e)}")
181188
raise

‎src/discord-cluster-manager/cogs/misc_cog.py‎

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22
from typing import TYPE_CHECKING
33

44
import discord
5-
import psycopg2
65
from discord import app_commands
76
from discord.ext import commands
7+
from sqlalchemy import create_engine, text
8+
from sqlalchemy.exc import SQLAlchemyError
89
from env import DATABASE_URL
910
from utils import send_discord_message, setup_logging
1011

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

3536
try:
36-
with psycopg2.connect(DATABASE_URL, sslmode="require") as conn:
37-
with conn.cursor() as cursor:
38-
cursor.execute("SELECT RANDOM()")
39-
result = cursor.fetchone()
40-
if result:
41-
random_value = result[0]
42-
await send_discord_message(
43-
interaction, f"Your lucky number is {random_value}."
44-
)
45-
else:
46-
await send_discord_message(interaction, "No result returned.")
37+
engine = create_engine(DATABASE_URL)
38+
with engine.connect() as connection:
39+
result = connection.execute(text("SELECT RANDOM()"))
40+
row = result.fetchone()
41+
if row:
42+
random_value = row[0]
43+
await send_discord_message(
44+
interaction, f"Your lucky number is {random_value}."
45+
)
46+
else:
47+
await send_discord_message(interaction, "No result returned.")
48+
except SQLAlchemyError as e:
49+
message = "Database error occurred"
50+
logger.error(f"{message}: {str(e)}", exc_info=True)
51+
await send_discord_message(interaction, f"{message}.")
4752
except Exception as e:
4853
message = "Error interacting with the database"
4954
logger.error(f"{message}: {str(e)}", exc_info=True)

0 commit comments

Comments
 (0)