db_helper.py

# db_helper.py
import logging
import asyncpg
from datetime import datetime
from typing import Optional, List
from credential_manager import CredentialManager
from models import UserCreate, User, UserUpdate, UserInDB, OAuth2Client, OAuth2ClientCreate, OAuth2ClientUpdate, \
    OAuth2AuthorizationCode

# Configure logging to output to stdout
logger = logging.getLogger(__name__)

class DBHelper:
    def __init__(self):
        self.credentials = CredentialManager.get_db_credentials()
        self.pool = None  # To be initialized in init_db

    async def connect_create_if_not_exists(self, user, database, password, port, host):
        """
        Connect to the specified database. If it doesn't exist, connect to 'postgres' and create it.
        """
        logger.info(f"Attempting to connect to database '{database}' as user '{user}' at {host}:{port}")
        try:
            self.pool = await asyncpg.create_pool(
                user=user,
                password=password,
                database=database,
                host=host,
                port=port,
                min_size=1,
                max_size=10
            )
            logger.info(f"Successfully connected to database '{database}' as user '{user}'")
        except asyncpg.exceptions.InvalidCatalogNameError:
            logger.warning(f"Database '{database}' does not exist. Attempting to create it.")
            try:
                # Connect to the default 'postgres' database to create the new database
                sys_pool = await asyncpg.create_pool(
                    user=user,
                    password=password,
                    database='postgres',
                    host=host,
                    port=port,
                    min_size=1,
                    max_size=5
                )
                async with sys_pool.acquire() as conn:
                    await conn.execute(f'CREATE DATABASE "{database}" OWNER "{user}"')
                    logger.info(f"Database '{database}' created successfully with owner '{user}'")
                await sys_pool.close()
                # Re-attempt connection to the newly created database
                self.pool = await asyncpg.create_pool(
                    user=user,
                    password=password,
                    database=database,
                    host=host,
                    port=port,
                    min_size=1,
                    max_size=10
                )
                logger.info(f"Successfully connected to newly created database '{database}'")
            except Exception as e:
                logger.error(f"Failed to create database '{database}': {e}")
                raise
        except Exception as e:
            logger.error(f"Failed to connect to database '{database}' as user '{user}': {e}")
            raise

    async def init_db(self):
        """
        Initialize the database by ensuring it exists and creating necessary tables.
        """
        user = self.credentials['user']
        database = self.credentials['database']
        password = self.credentials['password']
        port = self.credentials['port']
        host = self.credentials['host']

        logger.info("Starting database initialization process.")

        # Connect to the database, create if not exists
        await self.connect_create_if_not_exists(
            user=user,
            database=database,
            password=password,
            port=port,
            host=host
        )

        async with self.pool.acquire() as conn:
            try:
                # Begin transaction
                await conn.execute("BEGIN;")

                # Create tables with logging
                tables = {
                    "users": """
                        CREATE TABLE IF NOT EXISTS users (
                            id SERIAL PRIMARY KEY,
                            email VARCHAR UNIQUE NOT NULL,
                            hashed_password VARCHAR NOT NULL,
                            created_at TIMESTAMPTZ DEFAULT NOW()
                        );
                    """,
                    "oauth2_clients": """
                        CREATE TABLE IF NOT EXISTS oauth2_clients (
                            client_id VARCHAR PRIMARY KEY,
                            client_secret VARCHAR NOT NULL,
                            redirect_uris TEXT NOT NULL,
                            created_at TIMESTAMPTZ DEFAULT NOW()
                        );
                    """,
                    "oauth2_authorization_codes": """
                        CREATE TABLE IF NOT EXISTS oauth2_authorization_codes (
                            code VARCHAR PRIMARY KEY,
                            client_id VARCHAR REFERENCES oauth2_clients(client_id),
                            redirect_uri VARCHAR,
                            scope VARCHAR,
                            user_id INTEGER REFERENCES users(id),
                            code_challenge VARCHAR,
                            code_challenge_method VARCHAR,
                            expires_at TIMESTAMPTZ
                        );
                    """,
                    "refresh_tokens": """
                        CREATE TABLE IF NOT EXISTS refresh_tokens (
                            token VARCHAR PRIMARY KEY,
                            user_id INTEGER REFERENCES users(id),
                            expires_at TIMESTAMPTZ
                        );
                    """,
                    "Roles": """
                        CREATE TABLE IF NOT EXISTS Roles (
                            id SERIAL PRIMARY KEY,
                            role_name VARCHAR UNIQUE NOT NULL
                        );
                    """,
                    "Permissions": """
                        CREATE TABLE IF NOT EXISTS Permissions (
                            id SERIAL PRIMARY KEY,
                            permission_name VARCHAR UNIQUE NOT NULL
                        );
                    """,
                    "RolePermissions": """
                        CREATE TABLE IF NOT EXISTS RolePermissions (
                            role_id INTEGER REFERENCES Roles(id),
                            permission_id INTEGER REFERENCES Permissions(id),
                            PRIMARY KEY (role_id, permission_id)
                        );
                    """,
                    "UserRoles": """
                        CREATE TABLE IF NOT EXISTS UserRoles (
                            user_id INTEGER REFERENCES users(id),
                            role_id INTEGER REFERENCES Roles(id),
                            PRIMARY KEY (user_id, role_id)
                        );
                    """
                }

                for table_name, create_stmt in tables.items():
                    logger.info(f"Creating table '{table_name}' if it does not exist.")
                    await conn.execute(create_stmt)
                    logger.info(f"Table '{table_name}' is ready.")

                # Commit the transaction
                await conn.execute("COMMIT;")
                logger.info("Database initialization completed successfully.")

            except Exception as e:
                logger.error(f"Error during database initialization: {e}")
                await conn.execute("ROLLBACK;")
                raise

    async def get_db_connection(self):
        """
        Acquire a connection from the pool.
        """
        try:
            conn = await self.pool.acquire()
            logger.info(f"Acquired connection from the pool.")
            return conn
        except Exception as e:
            logger.error(f"Failed to acquire database connection from pool: {e}")
            raise

    async def close_pool(self):
        """
        Close the connection pool.
        """
        if self.pool:
            await self.pool.close()
            logger.info("Database connection pool closed.")

    # User methods
    async def add_user(self, user_in: UserCreate):
        """
        Add a new user to the database.

        :param user_in: UserCreate object containing email and password
        """
        hashed_password = user_in.password  # Password should already be hashed before calling this method
        email = user_in.email

        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                INSERT INTO users (email, hashed_password) VALUES ($1, $2)
                """, email, hashed_password)
                logger.info(f"Added new user with email: {email}")
            except asyncpg.exceptions.UniqueViolationError:
                logger.warning(f"User with email '{email}' already exists.")
                raise
            except Exception as e:
                logger.error(f"Error adding user '{email}': {e}")
                raise

    async def get_user_by_email(self, email: str) -> Optional[UserInDB]:
        query = "SELECT id, email, hashed_password FROM users WHERE email = $1"
        record = await self.pool.fetchrow(query, email)
        if record:
            return UserInDB(**record)
        return None

    async def get_user_by_id(self, user_id: int) -> Optional[User]:
        """
        Retrieve a user by ID.

        :param user_id: ID of the user
        :return: User object or None if not found
        """
        async with self.pool.acquire() as conn:
            user_record = await conn.fetchrow("SELECT id, email, created_at FROM users WHERE id = $1", user_id)
            logger.info(f"Fetched user by ID '{user_id}': {'Found' if user_record else 'Not Found'}")
            if user_record:
                return User(**dict(user_record))
            return None

    async def get_all_users(self) -> List[User]:
        """
        Retrieve all users.

        :return: List of User objects
        """
        async with self.pool.acquire() as conn:
            user_records = await conn.fetch("SELECT id, email, created_at FROM users")
            users = [User(**dict(user)) for user in user_records]
            logger.info(f"Fetched all users. Total: {len(users)}")
            return users

    async def update_user(self, user_id: int, user_in: UserUpdate):
        """
        Update a user's information.

        :param user_id: ID of the user to update
        :param user_in: UserUpdate object containing updated fields
        """
        async with self.pool.acquire() as conn:
            fields = []
            values = []
            idx = 1
            if user_in.email is not None:
                fields.append(f"email = ${idx}")
                values.append(user_in.email)
                idx += 1
            if user_in.password is not None:
                fields.append(f"hashed_password = ${idx}")
                values.append(user_in.password)  # Password should already be hashed before calling this method
                idx += 1
            if not fields:
                logger.info(f"No fields to update for user ID '{user_id}'")
                return
            values.append(user_id)
            query = f"UPDATE users SET {', '.join(fields)} WHERE id = ${idx}"
            try:
                await conn.execute(query, *values)
                logger.info(f"Updated user ID '{user_id}'")
            except Exception as e:
                logger.error(f"Error updating user ID '{user_id}': {e}")
                raise

    # db_helper.py

    async def delete_user(self, user_id: int):
        """
        Delete a user by ID.

        :param user_id: ID of the user to delete
        """
        async with self.pool.acquire() as conn:
            try:
                # Delete authorization codes related to the user
                await conn.execute("DELETE FROM oauth2_authorization_codes WHERE user_id = $1", user_id)

                # Delete refresh tokens related to the user
                await conn.execute("DELETE FROM refresh_tokens WHERE user_id = $1", user_id)

                # Delete roles associated with the user
                await conn.execute("DELETE FROM UserRoles WHERE user_id = $1", user_id)

                # Delete the user
                await conn.execute("DELETE FROM users WHERE id = $1", user_id)

                logger.info(f"Deleted user ID '{user_id}'")
            except Exception as e:
                logger.error(f"Error deleting user ID '{user_id}': {e}")
                raise

    # Client methods
    async def get_client_by_id(self, client_id: str) -> Optional[OAuth2Client]:
        async with self.pool.acquire() as conn:
            client_record = await conn.fetchrow("SELECT * FROM oauth2_clients WHERE client_id = $1", client_id)
            logger.info(f"Fetched OAuth2 client by ID '{client_id}': {'Found' if client_record else 'Not Found'}")
            logger.info(f"Fetched client: {client_record}")
            if client_record:
                return OAuth2Client(
                    client_id=client_record['client_id'],
                    client_secret=client_record['client_secret'],
                    redirect_uris=client_record['redirect_uris'].split(","),
                    created_at=client_record['created_at']
                )
            return None

    async def get_all_clients(self) -> List[OAuth2Client]:
        """
        Retrieve all OAuth2 clients.

        :return: List of OAuth2Client objects
        """
        async with self.pool.acquire() as conn:
            client_records = await conn.fetch("SELECT * FROM oauth2_clients")
            clients = []
            for record in client_records:
                clients.append(OAuth2Client(
                    client_id=record['client_id'],
                    client_secret=record['client_secret'],
                    redirect_uris=record['redirect_uris'].split(","),
                    created_at=record['created_at']
                ))
            logger.info(f"Fetched all clients. Total: {len(clients)}")
            return clients

    async def add_client(self, client_id: str, client_secret: str, redirect_uris: List[str]):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute('''
                    INSERT INTO oauth2_clients (client_id, client_secret, redirect_uris)
                    VALUES ($1, $2, $3)
                ''', client_id, client_secret, ",".join(redirect_uris))
                logger.info(f"Added OAuth2 client with ID: {client_id}")
            except asyncpg.exceptions.UniqueViolationError:
                logger.warning(f"OAuth2 client with ID '{client_id}' already exists.")
                raise
            except Exception as e:
                logger.error(f"Error adding client '{client_id}': {e}")
                raise

    async def update_client(self, client_id: str, client_data: OAuth2ClientUpdate):
        """
        Update an OAuth2 client's information.

        :param client_id: The client_id of the client to update
        :param client_data: OAuth2ClientUpdate object containing updated fields
        """
        async with self.pool.acquire() as conn:
            fields = []
            values = []
            idx = 1
            if client_data.client_secret is not None:
                fields.append(f"client_secret = ${idx}")
                values.append(client_data.client_secret)
                idx += 1
            if client_data.redirect_uris is not None:
                fields.append(f"redirect_uris = ${idx}")
                values.append(",".join(client_data.redirect_uris))
                idx += 1
            if not fields:
                logger.info(f"No fields to update for client ID '{client_id}'")
                return
            values.append(client_id)
            query = f"UPDATE oauth2_clients SET {', '.join(fields)} WHERE client_id = ${idx}"
            try:
                await conn.execute(query, *values)
                logger.info(f"Updated OAuth2 client with ID '{client_id}'")
            except Exception as e:
                logger.error(f"Error updating OAuth2 client '{client_id}': {e}")
                raise

    async def delete_client(self, client_id: str):
        """
        Delete an OAuth2 client by client_id.

        :param client_id: The client_id of the client to delete
        """
        async with self.pool.acquire() as conn:

            try:
                # Delete authorization codes related to the client
                await conn.execute("DELETE FROM oauth2_authorization_codes WHERE client_id = $1", client_id)

                # Delete the client
                await conn.execute("DELETE FROM oauth2_clients WHERE client_id = $1", client_id)

                logger.info(f"Deleted OAuth2 client with ID '{client_id}'")

            except Exception as e:

                logger.error(f"Error deleting OAuth2 client '{client_id}': {e}")
                raise

    # Authorization code methods
    async def save_authorization_code(self, auth_code: OAuth2AuthorizationCode):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                INSERT INTO oauth2_authorization_codes (
                    code, client_id, redirect_uri, scope, user_id, code_challenge, code_challenge_method, expires_at
                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                """, auth_code.code, auth_code.client_id, auth_code.redirect_uri, auth_code.scope,
                    auth_code.user_id, auth_code.code_challenge, auth_code.code_challenge_method, auth_code.expires_at)
                logger.info(f"Saved authorization code '{auth_code.code}' for client '{auth_code.client_id}' and user '{auth_code.user_id}'")
            except asyncpg.exceptions.UniqueViolationError:
                logger.warning(f"Authorization code '{auth_code.code}' already exists.")
            except Exception as e:
                logger.error(f"Error saving authorization code '{auth_code.code}': {e}")
                raise

    async def get_authorization_code(self, code: str):
        async with self.pool.acquire() as conn:
            auth_code = await conn.fetchrow("SELECT * FROM oauth2_authorization_codes WHERE code = $1", code)
            logger.info(f"Fetched authorization code '{code}': {'Found' if auth_code else 'Not Found'}")
            return dict(auth_code) if auth_code else None

    async def delete_authorization_code(self, code: str):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("DELETE FROM oauth2_authorization_codes WHERE code = $1", code)
                logger.info(f"Deleted authorization code '{code}'")
            except Exception as e:
                logger.error(f"Error deleting authorization code '{code}': {e}")
                raise

    # Refresh token methods
    async def save_refresh_token(self, user_id: int, token: str, expires_at: datetime):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                INSERT INTO refresh_tokens (token, user_id, expires_at) VALUES ($1, $2, $3)
                """, token, user_id, expires_at)
                logger.info(f"Saved refresh token '{token}' for user ID '{user_id}'")
            except asyncpg.exceptions.UniqueViolationError:
                logger.warning(f"Refresh token '{token}' already exists.")
            except Exception as e:
                logger.error(f"Error saving refresh token '{token}': {e}")
                raise

    async def get_refresh_token(self, token: str):
        async with self.pool.acquire() as conn:
            token_data = await conn.fetchrow("SELECT * FROM refresh_tokens WHERE token = $1", token)
            logger.info(f"Fetched refresh token '{token}': {'Found' if token_data else 'Not Found'}")
            return dict(token_data) if token_data else None

    async def delete_refresh_token(self, token: str):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("DELETE FROM refresh_tokens WHERE token = $1", token)
                logger.info(f"Deleted refresh token '{token}'")
            except Exception as e:
                logger.error(f"Error deleting refresh token '{token}': {e}")
                raise

    # Role and Permission methods
    async def create_role(self, role_name: str):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                    INSERT INTO Roles(role_name)
                    VALUES ($1)
                """, role_name)
                logger.info(f"Created role '{role_name}'")
            except asyncpg.exceptions.UniqueViolationError:
                logger.warning(f"Role '{role_name}' already exists")
            except Exception as e:
                logger.error(f"Error creating role '{role_name}': {e}")
                raise

    async def create_permission(self, permission_name: str):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                    INSERT INTO Permissions(permission_name)
                    VALUES ($1)
                """, permission_name)
                logger.info(f"Created permission '{permission_name}'")
            except asyncpg.exceptions.UniqueViolationError:
                logger.warning(f"Permission '{permission_name}' already exists")
            except Exception as e:
                logger.error(f"Error creating permission '{permission_name}': {e}")
                raise

    async def get_role_by_name(self, role_name: str):
        async with self.pool.acquire() as conn:
            role = await conn.fetchrow("SELECT * FROM Roles WHERE role_name = $1", role_name)
            logger.info(f"Fetched role '{role_name}': {'Found' if role else 'Not Found'}")
            return dict(role) if role else None

    async def get_permission_by_name(self, permission_name: str):
        async with self.pool.acquire() as conn:
            permission = await conn.fetchrow("SELECT * FROM Permissions WHERE permission_name = $1", permission_name)
            logger.info(f"Fetched permission '{permission_name}': {'Found' if permission else 'Not Found'}")
            return dict(permission) if permission else None

    async def assign_permission_to_role(self, role_id: int, permission_id: int):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                    INSERT INTO RolePermissions(role_id, permission_id)
                    VALUES ($1, $2)
                    ON CONFLICT DO NOTHING
                """, role_id, permission_id)
                logger.info(f"Assigned permission ID '{permission_id}' to role ID '{role_id}'")
            except Exception as e:
                logger.error(f"Error assigning permission ID '{permission_id}' to role ID '{role_id}': {e}")
                raise

    async def assign_role_to_user(self, user_id: int, role_id: int):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                    INSERT INTO UserRoles(user_id, role_id)
                    VALUES ($1, $2)
                    ON CONFLICT DO NOTHING
                """, user_id, role_id)
                logger.info(f"Assigned role ID '{role_id}' to user ID '{user_id}'")
            except Exception as e:
                logger.error(f"Error assigning role ID '{role_id}' to user ID '{user_id}': {e}")
                raise

    async def get_user_permissions(self, user_id: int):
        async with self.pool.acquire() as conn:
            permissions = await conn.fetch("""
                SELECT P.permission_name
                FROM Permissions P
                INNER JOIN RolePermissions RP ON P.id = RP.permission_id
                INNER JOIN UserRoles UR ON RP.role_id = UR.role_id
                WHERE UR.user_id = $1
            """, user_id)
            permission_names = [perm['permission_name'] for perm in permissions]
            logger.info(f"Retrieved permissions for user ID '{user_id}': {permission_names}")
            return permission_names

    async def get_user_roles(self, user_id: int) -> List[str]:
        async with self.pool.acquire() as conn:
            roles = await conn.fetch("""
                SELECT R.role_name
                FROM Roles R
                INNER JOIN UserRoles UR ON R.id = UR.role_id
                WHERE UR.user_id = $1
            """, user_id)
            return [role['role_name'] for role in roles]


    async def get_all_roles(self) -> List[dict]:
        async with self.pool.acquire() as conn:
            roles = await conn.fetch("SELECT id, role_name FROM Roles")
            return [dict(role) for role in roles]

    async def get_all_permissions(self) -> List[dict]:
        async with self.pool.acquire() as conn:
            permissions = await conn.fetch("SELECT id, permission_name FROM Permissions")
            return [dict(permission) for permission in permissions]

    async def delete_role(self, role_id: int):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("DELETE FROM Roles WHERE id = $1", role_id)
                logger.info(f"Deleted role ID '{role_id}'")
            except Exception as e:
                logger.error(f"Error deleting role ID '{role_id}': {e}")
                raise

    async def delete_permission(self, permission_id: int):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("DELETE FROM Permissions WHERE id = $1", permission_id)
                logger.info(f"Deleted permission ID '{permission_id}'")
            except Exception as e:
                logger.error(f"Error deleting permission ID '{permission_id}': {e}")
                raise

    async def remove_role_from_user(self, user_id: int, role_id: int):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                    DELETE FROM UserRoles
                    WHERE user_id = $1 AND role_id = $2
                """, user_id, role_id)
                logger.info(f"Removed role ID '{role_id}' from user ID '{user_id}'")
            except Exception as e:
                logger.error(f"Error removing role ID '{role_id}' from user ID '{user_id}': {e}")
                raise

    async def get_roles_for_user(self, user_id: int) -> List[dict]:
        async with self.pool.acquire() as conn:
            roles = await conn.fetch("""
                SELECT R.id, R.role_name
                FROM Roles R
                INNER JOIN UserRoles UR ON R.id = UR.role_id
                WHERE UR.user_id = $1
            """, user_id)
            return [dict(role) for role in roles]

    async def remove_permission_from_role(self, role_id: int, permission_id: int):
        async with self.pool.acquire() as conn:
            try:
                await conn.execute("""
                    DELETE FROM RolePermissions
                    WHERE role_id = $1 AND permission_id = $2
                """, role_id, permission_id)
                logger.info(f"Removed permission ID '{permission_id}' from role ID '{role_id}'")
            except Exception as e:
                logger.error(f"Error removing permission ID '{permission_id}' from role ID '{role_id}': {e}")
                raise

    async def get_permissions_for_role(self, role_id: int) -> List[dict]:
        async with self.pool.acquire() as conn:
            permissions = await conn.fetch("""
                SELECT P.id, P.permission_name
                FROM Permissions P
                INNER JOIN RolePermissions RP ON P.id = RP.permission_id
                WHERE RP.role_id = $1
            """, role_id)
            return [dict(permission) for permission in permissions]