"""Proxy for snowflake requests in order to log them in New Relic."""

from ddtrace import tracer
from snowflake.connector.connection import SnowflakeConnection
from snowflake_connector import snowflake_conn

from socials.utils.cache import cache_in_redis


def pass_rollback(self) -> None:
    """We're only doing SELECT queries. Rollback is not necessary.

    Original implementation in the connector:
        def rollback(self) -> None:
            self.cursor().execute("ROLLBACK")
    """
    pass


SnowflakeConnection.rollback = pass_rollback


DEFAULT_CONFIG = {
    'pool_pre_ping': False,
    'pool_reset_on_return': None,
    'commit_before_close': False
}


@cache_in_redis()
@tracer.wrap(name='snowflake_fetchone')
def fetchone(*args, **kwargs):
    """Passthrough to Snowflake to fetchone."""
    kwargs.update(DEFAULT_CONFIG)
    return snowflake_conn.fetchone(*args, **kwargs)


@cache_in_redis()
@tracer.wrap(name='snowflake_fetchall')
def fetchall(*args, **kwargs):
    """Passthrough to Snowflake to fetchall."""
    kwargs.update(DEFAULT_CONFIG)
    return snowflake_conn.fetchall(*args, **kwargs)


def SQLLoader(path):
    """Create a SQL loader in given path."""
    return snowflake_conn.SQLLoader(path)