# -*- test-case-name: ranger-ims-server.ext.test.test_sqlite -*-
"""
SQLite utilities
"""

from pathlib import Path
from sqlite3 import (
    Cursor, Error as SQLiteError, Row as BaseRow, connect as sqliteConnect
)
from typing import Any, Iterable, Mapping, Optional, Tuple, TypeVar
from typing.io import TextIO

from attr import Factory, attrib, attrs
from attr.validators import instance_of, optional

from twisted.logger import Logger


__all__ = (
    "Connection",
    "QueryPlanExplanation",
    "createDB",
    "explainQueryPlans",
    "openDB",
    "printSchema",
)


TConnection = TypeVar("TConnection", bound="Connection")



class Row(BaseRow):
    """
    Subclass of :class:`sqlite3.Row` that has a :class:`dict`-like ``get``
    method.
    """

    def get(self, key: str, default: Any = None) -> Any:
        if key in self.keys():
            return self[key]
        else:
            return default



@attrs(frozen=True)
class Connection(object):
    """
    (Incomplete) stand-in for :class:`sqlite3.Connection`.

    This class adds logging of SQL statements for debugging purposes.
    """

    # This would probably be a lot simpler if it were possible to subclass
    # sqlite3.Connection.

    _log = Logger()

    @attrs(frozen=False)
    class _State(object):
        """
        Internal mutable state for :class:`Connection`.
        """

        db = attrib(init=False)

    _database = attrib(
        validator=instance_of(str)
    )  # type: str

    _state = attrib(default=Factory(_State), init=False)  # type: _State


    def __attrs_post_init__(self) -> None:
        db = sqliteConnect(self._database)
        db.row_factory = Row
        db.execute("pragma foreign_keys = ON")

        self._state.db = db


    def commit(self) -> None:
        """
        See :meth:`sqlite3.Connection.commit`.
        """
        self._log.debug("COMMIT")
        self._state.db.commit()


    def executescript(self, sql_script: str) -> Cursor:
        """
        See :meth:`sqlite3.Connection.executescript`.
        """
        self._log.debug("EXECUTE SCRIPT:\n{script}", script=sql_script)
        return self._state.db.executescript(sql_script)


    def execute(self, sql: str, parameters: Mapping = None) -> Cursor:
        """
        See :meth:`sqlite3.Connection.execute`.
        """
        if parameters is None:
            parameters = {}
        self._log.debug(
            "EXECUTE: {sql} <- {parameters}", sql=sql, parameters=parameters
        )
        return self._state.db.execute(sql, parameters)


    def __enter__(self: TConnection) -> TConnection:
        self._log.debug("---------- ENTER ----------")
        self._state.db.__enter__()
        return self


    def __exit__(
        self, exc_type: type, exc_val: BaseException, exc_tb: Any
    ) -> bool:
        self._log.debug("---------- EXIT ----------")
        return self._state.db.__exit__(exc_type, exc_val, exc_tb)



def connect(path: Optional[Path]) -> Connection:
    """
    Open the database at the given path and configure it.
    """
    if path is None:
        endpoint = ":memory:"
    else:
        endpoint = str(path)

    return Connection(endpoint)


def createDB(path: Optional[Path], schema: str) -> Connection:
    """
    Create a new database at the given path.
    """
    db = connect(path)

    db.executescript(schema)
    db.commit()

    return db


def openDB(path: Path, schema: str = None) -> Connection:
    """
    Open an SQLite DB with the schema for this application.
    """
    if path.exists():
        return connect(path)

    if schema is not None:
        return createDB(path, schema)

    raise SQLiteError("Database does not exist: {}".format(path))


def printSchema(db: Connection, out: TextIO) -> None:
    """
    Print the database schema.
    """
    for (tableName,) in db.execute(
        """
        select NAME from SQLITE_MASTER where TYPE='table' order by NAME;
        """
    ):
        print("{}:".format(tableName), file=out)
        for (
            rowNumber, colName, colType, colNotNull, colDefault, colPK
        ) in db.execute("pragma table_info('{}');".format(tableName)):
            print(
                "  {n}: {name}({type}){null}{default}{pk}".format(
                    n=rowNumber,
                    name=colName,
                    type=colType,
                    null=" not null" if colNotNull else "",
                    default=" [{}]".format(colDefault) if colDefault else "",
                    pk=" *{}".format(colPK) if colPK else "",
                ),
                file=out,
            )



@attrs(frozen=True)
class QueryPlanExplanation(object):
    """
    Container for information about a query plan.
    """

    @attrs(frozen=True)
    class Line(object):
        """
        A line of information about a query plan.
        """

        nestingOrder = attrib(
            validator=optional(instance_of(int))
        )  # type: Optional[int]
        selectFrom = attrib(
            validator=optional(instance_of(int))
        )  # type: Optional[int]
        details = attrib(validator=instance_of(str))  # type: str

        def __str__(self) -> str:
            return "[{},{}] {}".format(
                self.nestingOrder, self.selectFrom, self.details
            )


    name  = attrib(validator=instance_of(str))    # type: str
    query = attrib(validator=instance_of(str))    # type: str
    lines = attrib(validator=instance_of(tuple))  # type: Tuple[Line]


    def __str__(self) -> str:
        text = ["{}:".format(self.name), "", "  -- query --", ""]

        text.extend(
            "    {}".format(l)
            for l in self.query.strip().split("\n")
        )

        if self.lines:
            text.extend(("", "  -- query plan --", ""))
            text.extend("    {}".format(line) for line in self.lines)

        return "\n".join(text)



def explainQueryPlans(
    db: Connection, queries: Iterable[Tuple[str, str]]
) -> Iterable[QueryPlanExplanation]:
    """
    Explain query plans for the given queries.
    """
    for query, name in queries:
        params = dict((x, x) for x in range(query.count(":")))  # Dummy params
        try:
            lines = (
                QueryPlanExplanation.Line(nestingOrder, selectFrom, details)
                for n, nestingOrder, selectFrom, details in (
                    db.execute("explain query plan {}".format(query), params)
                )
            )  # type: Iterable[QueryPlanExplanation.Line]
        except SQLiteError as e:
            lines = (QueryPlanExplanation.Line(None, None, "{}".format(e),),)

        yield QueryPlanExplanation(name, query, lines)

@@ -275,17 +275,24 @@
  TypeVar
+, Union
 %0A
-
 from typ
@@ -466,24 +466,41 @@
 onnection%22,%0A
+    %22Parameter%22,%0A
     %22QueryPl
@@ -511,24 +511,54 @@
 planation%22,%0A
+    %22Row%22,%0A    %22SQLiteError%22,%0A
     %22createD
@@ -681,16 +681,68 @@
 tion%22)%0A%0A
+Parameter = Optional%5BUnion%5Bbytes, str, int, float%5D%5D%0A
 %0A%0Aclass 

Add Parameter type. Export Parameter, Row, SQLiteError