from fansifter_common.adapters.db.base import Database
from snowflake.connector import SnowflakeConnection
from snowflake.snowpark import (
    Session as SnowparkSession,  # type: ignore[reportMissingTypeStubs]
)
from snowflake.snowpark._internal.server_connection import ServerConnection


class DB(Database):
    alias = "snowflake"

    def create_snowpark_session(self) -> SnowparkSession:
        return SnowparkSession(
            conn=ServerConnection(options={}, conn=self._get_snowflake_connection())
        )

    def _get_snowflake_connection(self) -> SnowflakeConnection:
        # Use snowflake active connection
        proxied = self.session.connection().connection
        if not isinstance(proxied.dbapi_connection, SnowflakeConnection):
            raise ValueError(
                "dbapi_connection must be an instance of SnowflakeConnection"
            )
        return proxied.dbapi_connection


class PgDB(Database):
    alias = "pg"