# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_database.ipynb.

# %% auto 0
__all__ = ['DatabaseSettings', 'Database']

# %% ../nbs/04_database.ipynb 3
from humble_chuck.models import BaseSettings
from typing import *
from pydantic import SecretStr
from humble_chuck.delegation import delegates
import sqlalchemy
from sqlalchemy import URL
from sqlalchemy import create_engine, Engine
from sqlalchemy.orm import Session
from pydantic import SecretStr, Field
from pydantic_settings import SettingsConfigDict
from typing import Optional,Union
from abc import ABC, abstractproperty,abstractmethod
from contextlib import contextmanager
import pandas as pd
from sqlalchemy import text


# %% ../nbs/04_database.ipynb 4
class DatabaseSettings(BaseSettings):
    """
    Pydantic settings model for arguments passed to SQL Alchemy's URL object.  
    
    Customized to accept variables from a central .yml file according to a key specified by the 'env_prefix' value of the model_config.
    """
    username: Optional[str]=None
    password: Optional[SecretStr]=None
    drivername:str 
    host: Optional[str]=None
    port: Optional[int]=None
    database: Optional[str]=None
    query: dict[str,str]={}

# %% ../nbs/04_database.ipynb 5
class Database(DatabaseSettings):
    """
    Create a Database connection with default functionality. 
    
    Inherits attributes from `DatabaseSettings` to manage credentials. 

    Use the class as-is to quickly create a database connection, or create a subclass to control connections to a specific database.
    
    More on SQL Alchemy engines from [SQLAlchemy docs](https://docs.sqlalchemy.org/en/13/core/engines.html):
    
    - The Engine is the starting point for any SQLAlchemy application. It’s “home base” for the actual database and its DBAPI.  
    - An Engine references both a Dialect and a Pool, which together interpret the DBAPI’s module functions as well as the behavior of the database 
    - Pool object which will establish a DBAPI connection at localhost:5432 when a connection request is first received  
    -- Note that the Engine and its underlying Pool do **not** establish the first actual DBAPI connection until the Engine.connect() method is called, or an operation which is dependent on this method such as Engine.execute() is invoked.
    - In this way, Engine and Pool can be said to have a lazy initialization behavior.  
    - The Engine, once created, can either be used directly to interact with the database, or can be passed to a Session object to work with the ORM.  
    """
    _engine:Engine = None
    _engine_url:URL = None

    def __init__(
        self,
        **kwargs
    ):
        # settings __init__
        super().__init__(**kwargs)
        if hasattr(self.password,'get_secret_value'):
            
            password = self.password.get_secret_value()
            
        else:
            password = self.password
        url = URL.create(
            drivername=self.drivername,
            username=self.username,
            password=password,
            host=self.host,
            port=self.port,
            database=self.database,
            query=self.query
        )
        self._engine_url=url
        self._engine=create_engine(url)

    
    model_config = SettingsConfigDict(
        #allows for attributes of `database settings` to be set as defaults in subclasses without type annotation
        ignored_types=(int,str,dict),
        arbitrary_types_allowed=True
    )

    def query_to_records(
        self,
        query_string:str,
    ):
        with self._engine.connect() as conn:
            results = [row for row in conn.execute(text(query_string)).mappings()]
        return results

    @delegates(pd.read_sql_query)
    def query_to_df(
        self,
        query_string,
        **kwargs
    ):
        f"""{pd.read_sql_query.__doc__}"""
        with self._engine.connect() as conn:
            df = pd.read_sql_query(
                text(query_string),conn,**kwargs)
            conn.close()
        return df

    @contextmanager
    def engine_scope(self,**kwargs):
        with self._engine.connect() as conn: 
            try:
                yield conn
                conn.commit()
            except:
                raise
            finally:
                conn.close()
                
    @contextmanager
    def session_scope(self,bind=None,**kwargs):
        """Provide a transactional scope around a series of operations."""

        session = Session(bind=self._engine,**kwargs)
        try:
            yield session
            session.commit()
        except:
            session.rollback()
            raise
        finally:
            session.close()

