from typing import Any
from typing import Dict
from typing import List
from typing import Mapping
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Type
from typing import Union

from typing_extensions import Protocol
from typing_extensions import TypedDict

from .base import Connection
from .base import Engine
from .cursor import CursorFetchStrategy
from .cursor import CursorResult
from .url import URL
from ..exc import StatementError
from ..sql.compiler import Compiled as Compiled
from ..sql.compiler import IdentifierPreparer
from ..sql.compiler import TypeCompiler as TypeCompiler
from ..sql.schema import Column
from ..sql.type_api import TypeEngine

class _DBAPIConnectionBase(Protocol):
    def close(self) -> None: ...
    def commit(self) -> None: ...
    def cursor(self) -> _DBAPICursor: ...

class _DBAPIConnectionTransactions(_DBAPIConnectionBase, Protocol):
    def rollback(self) -> None: ...

_DBAPIConnection = Union[_DBAPIConnectionBase, _DBAPIConnectionTransactions]

class _DBAPICursorBase(Protocol):
    @property
    def description(self) -> Any: ...
    @property
    def rowcount(self) -> int: ...
    arraysize: int = ...
    def close(self) -> None: ...
    def execute(self, operation: Any, parameters: Any) -> Any: ...
    def executemany(
        self, operation: Any, parameters: Sequence[Any]
    ) -> Any: ...
    def fetchone(self) -> Optional[Any]: ...
    def fetchmany(self, size: int = ...) -> Sequence[Any]: ...
    def fetchall(self) -> Sequence[Any]: ...
    def setinputsizes(self, sizes: Sequence[Any]) -> None: ...
    def setoutputsize(self, size: Any, column: Any) -> None: ...

class _DBAPICursorStoredProcedures(_DBAPICursorBase, Protocol):
    def callproc(
        self, procname: str, parameters: Sequence[Any] = ...
    ) -> Any: ...

class _DBAPICursorMultiResultSets(_DBAPICursorBase, Protocol):
    def nextset(self) -> Optional[bool]: ...

_DBAPICursor = Union[
    _DBAPICursorBase, _DBAPICursorStoredProcedures, _DBAPICursorMultiResultSets
]

class _OnConnect(Protocol):
    def __call__(self, connection: _DBAPIConnection) -> None: ...

class _PKConstraintDictBase(TypedDict):
    constrained_columns: List[str]

class _PKConstraintDict(_PKConstraintDictBase, total=False):
    name: str

class _ForeignKeyDict(TypedDict):
    name: str
    constrained_columns: List[str]
    referred_schema: str
    referred_table: str
    referred_columns: List[str]

class _IndexDict(TypedDict):
    name: str
    column_names: List[str]
    unique: bool

class _TableCommentDict(TypedDict):
    text: str

class Dialect:
    name: str = ...
    driver: str = ...
    positional: bool = ...
    paramstyle: str = ...
    encoding: str = ...
    statement_compiler: Type[Compiled] = ...
    ddl_compiler: Type[Compiled] = ...
    server_version_info: Optional[Tuple[Any, ...]] = ...
    default_schema_name: Optional[str] = ...
    execution_ctx_cls: Type[ExecutionContext] = ...
    execute_sequence_format: Union[
        Type[Tuple[Any, ...]], Type[List[Any]]
    ] = ...
    preparer: Type[IdentifierPreparer] = ...
    supports_alter: bool = ...
    max_identifier_length: int = ...
    supports_sane_rowcount: bool = ...
    supports_sane_multi_rowcount: bool = ...
    preexecute_autoincrement_sequences: bool = ...
    implicit_returning: bool = ...
    colspecs: Dict[Type[TypeEngine[Any]], Type[TypeEngine[Any]]] = ...
    supports_sequences: bool = ...
    sequences_optional: bool = ...
    supports_native_enum: bool = ...
    supports_native_boolean: bool = ...
    dbapi_exception_translation_map: Dict[str, str] = ...
    supports_statement_cache: bool = ...
    def create_connect_args(
        self, url: URL
    ) -> Tuple[Sequence[Any], Mapping[str, Any]]: ...
    @classmethod
    def type_descriptor(cls, typeobj: Type[TypeEngine[Any]]) -> Any: ...
    def initialize(self, connection: Connection) -> None: ...
    def get_columns(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> List[Dict[str, Any]]: ...
    def get_pk_constraint(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> _PKConstraintDict: ...
    def get_foreign_keys(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> List[_ForeignKeyDict]: ...
    def get_table_names(
        self, connection: Connection, schema: Optional[str] = ..., **kw: Any
    ) -> List[str]: ...
    def get_temp_table_names(
        self, connection: Connection, schema: Optional[str] = ..., **kw: Any
    ) -> List[str]: ...
    def get_view_names(
        self, connection: Connection, schema: Optional[str] = ..., **kw: Any
    ) -> List[str]: ...
    def get_sequence_names(
        self, connection: Connection, schema: Optional[str] = ..., **kw: Any
    ) -> List[str]: ...
    def get_temp_view_names(
        self, connection: Connection, schema: Optional[str] = ..., **kw: Any
    ) -> List[str]: ...
    def get_view_definition(
        self,
        connection: Connection,
        view_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> Any: ...
    def get_indexes(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> List[_IndexDict]: ...
    def get_unique_constraints(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> List[Dict[str, Any]]: ...
    def get_check_constraints(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> List[Dict[str, Any]]: ...
    def get_table_comment(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> _TableCommentDict: ...
    def normalize_name(self, name: Optional[str]) -> Optional[str]: ...
    def denormalize_name(self, name: Optional[str]) -> Optional[str]: ...
    def has_table(
        self,
        connection: Connection,
        table_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> bool: ...
    def has_index(
        self,
        connection: Connection,
        table_name: str,
        index_name: str,
        schema: Optional[str] = ...,
    ) -> bool: ...
    def has_sequence(
        self,
        connection: Connection,
        sequence_name: str,
        schema: Optional[str] = ...,
        **kw: Any,
    ) -> bool: ...
    def do_begin(self, dbapi_connection: _DBAPIConnection) -> None: ...
    def do_rollback(self, dbapi_connection: _DBAPIConnection) -> None: ...
    def do_commit(self, dbapi_connection: _DBAPIConnection) -> None: ...
    def do_close(self, dbapi_connection: _DBAPIConnection) -> None: ...
    def do_set_input_sizes(
        self,
        cursor: _DBAPICursor,
        list_of_tuples: Tuple[str, Any, TypeEngine[Any]],
        context: ExecutionContext,
    ) -> None: ...
    def create_xid(self) -> str: ...
    def do_savepoint(self, connection: Connection, name: str) -> None: ...
    def do_rollback_to_savepoint(
        self, connection: Connection, name: str
    ) -> None: ...
    def do_release_savepoint(
        self, connection: Connection, name: str
    ) -> None: ...
    def do_begin_twophase(self, connection: Connection, xid: str) -> None: ...
    def do_prepare_twophase(
        self, connection: Connection, xid: str
    ) -> None: ...
    def do_rollback_twophase(
        self,
        connection: Connection,
        xid: str,
        is_prepared: bool = ...,
        recover: bool = ...,
    ) -> None: ...
    def do_commit_twophase(
        self,
        connection: Connection,
        xid: str,
        is_prepared: bool = ...,
        recover: bool = ...,
    ) -> None: ...
    def do_recover_twophase(self, connection: Connection) -> None: ...
    def do_executemany(
        self,
        cursor: _DBAPICursor,
        statement: Any,
        parameters: Any,
        context: Optional[ExecutionContext] = ...,
    ) -> None: ...
    def do_execute(
        self,
        cursor: _DBAPICursor,
        statement: Any,
        parameters: Any,
        context: Optional[ExecutionContext] = ...,
    ) -> None: ...
    def do_execute_no_params(
        self,
        cursor: _DBAPICursor,
        statement: Any,
        parameters: Any,
        context: Optional[ExecutionContext] = ...,
    ) -> None: ...
    def is_disconnect(
        self, e: Any, connection: Connection, cursor: _DBAPICursor
    ) -> bool: ...
    def connect(self, *cargs: Any, **cparams: Any) -> _DBAPIConnection: ...
    def on_connect_url(self, url: URL) -> Optional[_OnConnect]: ...
    def on_connect(self) -> Optional[_OnConnect]: ...
    def reset_isolation_level(self, dbapi_conn: _DBAPIConnection) -> None: ...
    def set_isolation_level(
        self, dbapi_conn: _DBAPIConnection, level: Any
    ) -> None: ...
    def get_isolation_level(self, dbapi_conn: _DBAPIConnection) -> Any: ...
    def get_default_isolation_level(
        self, dbapi_conn: _DBAPIConnection
    ) -> Any: ...
    @classmethod
    def get_dialect_cls(cls, url: URL) -> Dialect: ...
    @classmethod
    def load_provisioning(cls) -> None: ...
    @classmethod
    def engine_created(cls, engine: Engine) -> None: ...

class CreateEnginePlugin:
    url: URL = ...
    def __init__(self, url: URL, kwargs: Dict[str, Any]) -> None: ...
    def update_url(self, url: URL) -> URL: ...
    def handle_dialect_kwargs(
        self, dialect_cls: Type[Dialect], dialect_args: Dict[str, Any]
    ) -> None: ...
    def handle_pool_kwargs(
        self, pool_cls: Any, pool_args: Dict[str, Any]
    ) -> None: ...
    def engine_created(self, engine: Engine) -> None: ...

class ExecutionContext:
    connection: Connection = ...
    root_connection: Connection = ...
    dialect: Dialect = ...
    cursor: _DBAPICursor = ...
    compiled: Optional[Compiled] = ...
    statement: Optional[str] = ...
    parameters: Any = ...
    isinsert: bool = ...
    isupdate: bool = ...
    should_autocommit: bool = ...
    prefetch_cols: Sequence[Column[Any]] = ...
    postfetch_cols: Sequence[Column[Any]] = ...
    def create_cursor(self) -> _DBAPICursor: ...
    def pre_exec(self) -> None: ...
    def get_out_parameter_values(
        self, out_param_names: Sequence[str]
    ) -> Sequence[Any]: ...
    def post_exec(self) -> None: ...
    def get_result_cursor_strategy(
        self, result: Any
    ) -> CursorFetchStrategy: ...
    def handle_dbapi_exception(self, e: Any) -> None: ...
    def should_autocommit_text(self, statement: str) -> bool: ...
    def lastrow_has_defaults(self) -> bool: ...
    def get_rowcount(self) -> int: ...

class Connectable:
    def connect(self, **kwargs: Any) -> Connection: ...
    engine: Optional[Engine] = ...
    def execute(
        self, object_: Any, *multiparams: Any, **params: Any
    ) -> CursorResult: ...
    def scalar(
        self, object_: Any, *multiparams: Any, **params: Any
    ) -> Any: ...

class ExceptionContext:
    connection: Optional[Connection] = ...
    engine: Engine = ...
    cursor: Optional[_DBAPICursor] = ...
    statement: Optional[str] = ...
    parameters: Optional[Any] = ...
    original_exception: BaseException = ...
    sqlalchemy_exception: Optional[StatementError] = ...
    chained_exception: Optional[BaseException] = ...
    execution_context: Optional[ExecutionContext] = ...
    is_disconnect: bool = ...
    invalidate_pool_on_disconnect: bool = ...
