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

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

# %% ../nbs/04_database.ipynb 1
from archetypon.delegates import delegates
import os
import pandas as pd
from typing import Any,Type,Union,Callable,Optional
from snowflake.sqlalchemy import URL as SnowflakeURL
from pydantic import SecretStr,Field,FilePath
from typing import Optional
from pathlib import Path
import configparser
from sqlalchemy.engine import create_engine,Engine
from sqlalchemy import MetaData
from contextlib import contextmanager
from sqlalchemy.orm import Session
from typing import Literal,Any
from pydantic import validator,root_validator

from sqlalchemy import Sequence
from abc import ABC, abstractproperty,abstractmethod
from sqlalchemy.engine.url import URL
from pydantic import BaseSettings

# %% ../nbs/04_database.ipynb 3
class AbstractDatabaseClass(ABC):

    """
    Abstract Base Class for all Database Connections.

    <br><br>
    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 <br><br>
    > 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.

    """

    def query_to_records(
        self,
        query_string:str,
        model:Union[Type,Callable]=dict
    ):
        with self.engine.connect() as conn:
            results = [model(row) for row in conn.execute(query_string).fetchall()]
        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(query_string,conn,**kwargs)
        return df

    @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()
    
    @property
    def tables(self):
        return self.metadata.tables

    def __repr__(self):
        return (
            "<Database: " +"".join(
                [f"{k}={v}, " for k, v in
                 self.engine.url.__dict__.items()
                 if k in ('drivername','host','database')
                ]).strip(' ,')+'>'
        )

# %% ../nbs/04_database.ipynb 5
class DatabaseCredentialsBase(BaseSettings):
    username: str = None
    password: SecretStr = ''


# %% ../nbs/04_database.ipynb 6
class DatabaseCredentials(DatabaseCredentialsBase):
    """ Passed to regular SQLAlchemy URL constructor"""
    drivername: str
    host:str = None
    port:int = None
    query:str = None
    database:str = None

# %% ../nbs/04_database.ipynb 7
class SnowflakeCredentials(DatabaseCredentialsBase):
    """Passed to special Snowflake URL constructor"""
    account: str
    database: Optional[str]
    snowflake_schema: Optional[str] = Field(
        ...,
        env='snowflake_schema',
        alias='schema',
        description="""passed to the SQLAlchemy URL constructor as 'schema', 
        but we need to alias it to avoid clashing with 'BaseSettings'""")
    warehouse: Optional[str]
    role: Optional[str]
    numpy: Optional[bool]
        


# %% ../nbs/04_database.ipynb 8
class Database(DatabaseCredentials,AbstractDatabaseClass):
    f"""{URL.__doc__}"""

    engine:Engine = None
    engine_url:URL = None
    metadata:MetaData = None

    def __init__(
        self,
        **kwargs
    ):
        # settings __init__
        super().__init__(**kwargs)
        url = URL(
            drivername=self.drivername,
            username=self.username,
            password=self.password.get_secret_value(),
            host=self.host,
            port=self.port,
            database=self.database,
            query=self.query
        )
        self.engine_url=url
        self.engine=create_engine(url)
        self.metadata = MetaData(bind=self.engine)

    class Config:
        arbitrary_types_allowed=True

    def __repr__(self):
        return (
            "<Database: " +"".join(
                [f"{k}={v}, " for k, v in
                 self.engine.url.__dict__.items()
                 if k in ('drivername','host','database', 'username')
                ]).strip(' ,')+'>'
        )

# %% ../nbs/04_database.ipynb 9
class SnowflakeDatabase(SnowflakeCredentials,AbstractDatabaseClass):
    engine_url: URL = None
    engine: Engine = None
    metadata: MetaData = None


    def __init__(
        self,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.engine_url = SnowflakeURL(
            user = self.username,
            password = self.password.get_secret_value(),
            account=self.account,
            database = self.database,
            schema = self.snowflake_schema,
            warehouse = self.warehouse,
            role = self.role,
            numpy = self.numpy
        )

        self.engine=create_engine(self.engine_url)
        self.metadata = MetaData(bind=self.engine)

    class Config:
        arbitrary_types_allowed=True

    def __repr__(self):
        return (
            "<Snowflake Database: " +"".join(
                [f"{k}={v}, " for k, v in
                 self.dict().items()
                 if k in ('account','database','snowflake_schema','warehouse')
                ]).strip(' ,')+'>'
        )
