from typing import Iterator

import pymssql  # type: ignore
from sqlalchemy import text

from ...utils import ExtractionQuery, SqlalchemyClient, uri_encode

SERVER_URI = "{user}:{password}@{host}:{port}/{database}"
MSSQL_URI = f"mssql+pymssql://{SERVER_URI}"
DEFAULT_PORT = 1433

_KEYS = ("user", "password", "host", "port", "database")


def _check_key(credentials: dict) -> None:
    for key in credentials:
        if key not in _KEYS:
            raise KeyError(f"Missing {key} in credentials")


class MSSQLClient(SqlalchemyClient):
    """Microsoft Server SQL client"""

    @staticmethod
    def name() -> str:
        return "MSSQL"

    def _engine_options(self, credentials: dict) -> dict:
        return {}

    def _build_uri(self, credentials: dict) -> str:
        _check_key(credentials)
        uri = MSSQL_URI.format(
            user=credentials["user"],
            password=uri_encode(credentials["password"]),
            host=credentials["host"],
            port=credentials.get("port") or DEFAULT_PORT,
            database=credentials["database"],
        )
        return uri

    def execute(self, query: ExtractionQuery) -> Iterator[dict]:
        """
        Re-implements the SQLAlchemyClient execute function to ensure we consume
        the cursor before calling connection.close() as it wipes out the data
        otherwise
        """
        connection = self.connect()
        try:
            proxy = connection.execute(text(query.statement), query.params)
            results = list(self._process_result(proxy))
            yield from results
        finally:
            self.close()
