db.py•2.87 kB
from datetime import datetime
import sqlite3
from typing import Optional, Dict
def init_db() -> None:
"""Initialize the SQLite database with required tables"""
conn = sqlite3.connect('spotify_users.db')
c = conn.cursor()
c.execute('''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
spotify_id TEXT UNIQUE,
email TEXT,
access_token TEXT,
refresh_token TEXT,
token_expiry TIMESTAMP,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
''')
conn.commit()
conn.close()
def get_user_tokens(spotify_id: str) -> Optional[Dict]:
"""Retrieve user tokens from the database"""
conn = sqlite3.connect('spotify_users.db')
c = conn.cursor()
c.execute('SELECT access_token, refresh_token, token_expiry FROM users WHERE spotify_id = ?', (spotify_id,))
result = c.fetchone()
conn.close()
if result:
return {
'access_token': result[0],
'refresh_token': result[1],
'token_expiry': result[2]
}
return None
def save_user_tokens(spotify_id: str, email: str, access_token: str, refresh_token: str) -> None:
"""Save or update user tokens in the database"""
conn = sqlite3.connect('spotify_users.db')
c = conn.cursor()
c.execute('''
INSERT OR REPLACE INTO users (spotify_id, email, access_token, refresh_token, token_expiry)
VALUES (?, ?, ?, ?, ?)
''', (spotify_id, email, access_token, refresh_token,
datetime.now().timestamp() + 3600)) # Tokens typically expire in 1 hour
conn.commit()
conn.close()
def update_access_token(spotify_id: str, access_token: str, expires_in: int = 3600) -> None:
"""
Update the access token and its expiry time for a user after token refresh.
Args:
spotify_id: The Spotify user ID
access_token: The new access token
expires_in: Token expiry duration in seconds (default 1 hour)
"""
conn = sqlite3.connect('spotify_users.db')
c = conn.cursor()
c.execute('''
UPDATE users
SET access_token = ?, token_expiry = ?
WHERE spotify_id = ?
''', (access_token, datetime.now().timestamp() + expires_in, spotify_id))
conn.commit()
conn.close()
def get_first_user() -> Optional[Dict]:
"""Get the first user from the database"""
conn = sqlite3.connect('spotify_users.db')
c = conn.cursor()
c.execute('SELECT spotify_id, access_token, refresh_token, token_expiry FROM users ORDER BY created_at ASC LIMIT 1')
result = c.fetchone()
conn.close()
if result:
return {
'spotify_id': result[0],
'access_token': result[1],
'refresh_token': result[2],
'token_expiry': result[3]
}
return None