# 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