database.py

import os
import logging
import asyncpg
from typing import Optional

# Configure logging for production observability
logging.basicConfig(
 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
 level=logging.INFO
)
logger = logging.getLogger(__name__)

# Global connection pool
_pool: Optional[asyncpg.Pool] = None

async def get_pool() -> asyncpg.Pool:
 """
 Initializes and returns the asyncpg connection pool.
 Ensures a singleton pool is used across the application.
 """
 global _pool
 if _pool is None:
 db_url = os.getenv("DATABASE_URL")
 if not db_url:
 raise ValueError("Critical Error: DATABASE_URL environment variable is not set.")
 
 try:
 # Render managed PostgreSQL typically requires SSL for external connections.
 # The sslmode can be configured via the DATABASE_URL DSN.
 _pool = await asyncpg.create_pool(
 dsn=db_url,
 min_size=2,
 max_size=20,
 command_timeout=60,
 server_settings={'application_name': 'web3_tg_bot'}
 )
 logger.info("PostgreSQL connection pool created successfully.")
 except Exception as e:
 logger.error(f"Failed to initialize database pool: {e}")
 raise
 return _pool

async def init_db() -> None:
 """
 Connects to the database and initializes required tables atomically.
 """
 pool = await get_pool()
 try:
 async with pool.acquire() as conn:
 async with conn.transaction():
 await conn.execute("""
 CREATE TABLE IF NOT EXISTS users (
 telegram_id BIGINT PRIMARY KEY,
 sol_wallet TEXT,
 xp INT DEFAULT 0,
 level INT DEFAULT 1,
 balance_lamports BIGINT DEFAULT 0,
 trust_score INT DEFAULT 100,
 streak_count INT DEFAULT 0,
 last_check_in TIMESTAMP
 );
 """)
 
 await conn.execute("""
 CREATE TABLE IF NOT EXISTS agencies (
 agency_id BIGINT PRIMARY KEY,
 company_name TEXT,
 balance_lamports BIGINT DEFAULT 0,
 api_key TEXT,
 status TEXT DEFAULT 'pending'
 );
 """)
 
 await conn.execute("""
 CREATE TABLE IF NOT EXISTS tasks (
 task_id SERIAL PRIMARY KEY,
 owner_id BIGINT,
 prompt_text TEXT,
 reward_lamports BIGINT,
 required_consensus INT,
 status TEXT DEFAULT 'active'
 );
 """)
 
 await conn.execute("""
 CREATE TABLE IF NOT EXISTS submissions (
 submission_id SERIAL PRIMARY KEY,
 task_id INT REFERENCES tasks(task_id) ON DELETE CASCADE,
 user_id BIGINT REFERENCES users(telegram_id) ON DELETE CASCADE,
 user_answer TEXT,
 timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
 );
 """)
 
 await conn.execute("""
 CREATE TABLE IF NOT EXISTS payouts (
 payout_id SERIAL PRIMARY KEY,
 user_id BIGINT REFERENCES users(telegram_id) ON DELETE CASCADE,
 destination_wallet TEXT,
 amount_lamports BIGINT,
 tx_signature TEXT,
 status TEXT
 );
 """)
 logger.info("Database schema initialized successfully.")
 except Exception as e:
 logger.error(f"Error initializing database schema: {e}")
 raise

async def register_user(telegram_id: int, wallet: str) -> Optional[asyncpg.Record]:
 """
 Registers a new user or updates their Solana wallet if they already exist.
 """
 pool = await get_pool()
 try:
 async with pool.acquire() as conn:
 query = """
 INSERT INTO users (telegram_id, sol_wallet)
 VALUES ($1, $2)
 ON CONFLICT (telegram_id) DO UPDATE
 SET sol_wallet = EXCLUDED.sol_wallet
 RETURNING *;
 """
 return await conn.fetchrow(query, telegram_id, wallet)
 except Exception as e:
 logger.error(f"Error registering user {telegram_id}: {e}")
 raise

async def get_user(telegram_id: int) -> Optional[asyncpg.Record]:
 """
 Retrieves a user's record by their Telegram ID.
 """
 pool = await get_pool()
 try:
 async with pool.acquire() as conn:
 query = "SELECT * FROM users WHERE telegram_id = $1;"
 return await conn.fetchrow(query, telegram_id)
 except Exception as e:
 logger.error(f"Error fetching user {telegram_id}: {e}")
 raise

async def update_balance(telegram_id: int, amount_lamports: int) -> Optional[int]:
 """
 Atomically updates a user's balance (supports both positive and negative amounts).
 Returns the new balance.
 """
 pool = await get_pool()
 try:
 async with pool.acquire() as conn:
 query = """
 UPDATE users
 SET balance_lamports = balance_lamports + $2
 WHERE telegram_id = $1
 RETURNING balance_lamports;
 """
 return await conn.fetchval(query, telegram_id, amount_lamports)
 except Exception as e:
 logger.error(f"Error updating balance for user {telegram_id}: {e}")
 raise

async def add_task(owner_id: int, prompt_text: str, reward_lamports: int, required_consensus: int) -> Optional[int]:
 """
 Creates a new task and returns the generated task_id.
 """
 pool = await get_pool()
 try:
 async with pool.acquire() as conn:
 query = """
 INSERT INTO tasks (owner_id, prompt_text, reward_lamports, required_consensus)
 VALUES ($1, $2, $3, $4)
 RETURNING task_id;
 """
 return await conn.fetchval(query, owner_id, prompt_text, reward_lamports, required_consensus)
 except Exception as e:
 logger.error(f"Error adding task for owner {owner_id}: {e}")
 raise

async def get_available_task(user_id: int) -> Optional[asyncpg.Record]:
 """
 Fetches an active task that the user has not yet submitted an answer for.
 """
 pool = await get_pool()
 try:
 async with pool.acquire() as conn:
 query = """
 SELECT * FROM tasks
 WHERE status = 'active'
 AND task_id NOT IN (
 SELECT task_id FROM submissions WHERE user_id = $1
 )
 ORDER BY task_id ASC
 LIMIT 1;
 """
 return await conn.fetchrow(query, user_id)
 except Exception as e:
 logger.error(f"Error fetching available task for user {user_id}: {e}")
 raise

async def submit_work(task_id: int, user_id: int, user_answer: str) -> Optional[int]:
 """
 Records a user's submission for a task atomically.
 Returns the generated submission_id.
 """
 pool = await get_pool()
 try:
 async with pool.acquire() as conn:
 async with conn.transaction():
 query = """
 INSERT INTO submissions (task_id, user_id, user_answer)
 VALUES ($1, $2, $3)
 RETURNING submission_id;
 """
 return await conn.fetchval(query, task_id, user_id, user_answer)
 except Exception as e:
 logger.error(f"Error submitting work for task {task_id} by user {user_id}: {e}")
 raise

async def close_db() -> None:
 """
 Gracefully closes the database connection pool.
 """
 global _pool
 if _pool is not None:
 await _pool.close()
 _pool = None
 logger.info("Database connection pool closed.")