Spaces:
Build error
Build error
File size: 3,939 Bytes
ad1f59e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | # database.py
import logging
from supabase import create_client, Client
import config
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# --- DATABASE INITIALIZATION ---
try:
supabase: Client = create_client(config.SUPABASE_URL, config.SUPABASE_KEY)
logger.info("Successfully connected to Supabase.")
except Exception as e:
logger.error(f"Error connecting to Supabase: {e}")
supabase = None
# --- DATABASE FUNCTIONS ---
def get_or_create_user(user_id: int, username: str):
"""
Retrieves a user from the database or creates a new one if they don't exist.
"""
if not supabase:
logger.error("Supabase client not initialized.")
return None
try:
# Check if user exists
response = supabase.table('users').select('*').eq('user_id', user_id).execute()
# response.data is a list. If it's empty, the user doesn't exist.
if not response.data:
logger.info(f"User {user_id} ({username}) not found. Creating new user.")
insert_response = supabase.table('users').insert({
'user_id': user_id,
'username': username
}).execute()
if not insert_response.data:
logger.error(f"Failed to create user {user_id}. Response: {insert_response}")
return None
return insert_response.data[0]
else:
logger.info(f"Found existing user {user_id} ({username}).")
return response.data[0]
except Exception as e:
logger.error(f"Error in get_or_create_user for user {user_id}: {e}")
return None
def get_user_credits(user_id: int) -> int:
"""
Fetches the current credit balance for a given user.
Returns 0 if the user is not found or an error occurs.
"""
if not supabase:
return 0
try:
response = supabase.table('users').select('credits').eq('user_id', user_id).single().execute()
if response.data:
return response.data.get('credits', 0)
return 0
except Exception as e:
logger.error(f"Error getting credits for user {user_id}: {e}")
return 0
def deduct_credit(user_id: int) -> bool:
"""
Deducts one credit from a user. Returns True on success, False on failure.
"""
if not supabase:
return False
current_credits = get_user_credits(user_id)
if current_credits <= 0:
logger.warning(f"User {user_id} has no credits to deduct.")
return False
try:
new_credits = current_credits - 1
response = supabase.table('users').update({'credits': new_credits}).eq('user_id', user_id).execute()
if response.data:
logger.info(f"Successfully deducted 1 credit from user {user_id}. New balance: {new_credits}")
return True
return False
except Exception as e:
logger.error(f"Error deducting credit for user {user_id}: {e}")
return False
def update_user_plan(user_id: int, plan: str, credits_to_add: int):
"""
Updates a user's plan and adds a specified number of credits to their balance.
"""
if not supabase:
return False
try:
current_credits = get_user_credits(user_id)
new_credits = current_credits + credits_to_add
response = supabase.table('users').update({
'plan': plan,
'credits': new_credits
}).eq('user_id', user_id).execute()
if response.data:
logger.info(f"User {user_id} plan updated to '{plan}'. Added {credits_to_add} credits. New balance: {new_credits}")
return True
return False
except Exception as e:
logger.error(f"Error updating plan for user {user_id}: {e}")
return False |