database.py

import sqlite3
import os
from flask import g

def get_db_path():
    if os.path.exists('/data'):
        return '/data/users.db'
    elif os.path.exists('/home/user_data'):
        return '/home/user_data/users.db'
    return 'users.db'

def get_db():
    if 'db' not in g:
        db_path = get_db_path()
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        g.db = sqlite3.connect(db_path)
        g.db.row_factory = sqlite3.Row
    return g.db

def close_db(e=None):
    db = g.pop('db', None)
    if db is not None:
        db.close()

def init_db(app):
    with app.app_context():
        db = get_db()
        cursor = db.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                udid TEXT NOT NULL UNIQUE,
                name TEXT DEFAULT NULL,
                premium_account BOOLEAN DEFAULT FALSE,
                usage_limit INTEGER DEFAULT 3,
                invoice TEXT DEFAULT NULL
            )
        ''')
        
        cursor.execute('SELECT COUNT(*) FROM users')
        if cursor.fetchone()[0] == 0:
            cursor.executemany('INSERT INTO users (udid, name, premium_account, usage_limit, invoice) VALUES (?, ?, ?, ?, ?)',
                             [('test_udid_1', 'Test User 1', True, 4, 'test_invoice_1'), 
                              ('test_udid_2', 'Test User 2', False, 4, None)])
        
        db.commit()