#! /usr/bin/python3
# Copyright 2021-present StarRocks, Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https:#www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from sqlalchemy import exc, log
from sqlalchemy.dialects.mysql.mysqldb import MySQLDialect_mysqldb
from sqlalchemy.engine import reflection


@log.class_logger
class StarRocksDialect(MySQLDialect_mysqldb):
    # Caching
    # Warnings are generated by SQLAlchmey if this flag is not explicitly set
    # and tests are needed before being enabled
    supports_statement_cache = False

    name = 'starrocks'

    def __init__(self, *args, **kw):
        super(StarRocksDialect, self).__init__(*args, **kw)

    @reflection.cache
    def _get_default_schema_name(self, connection):
        return connection.exec_driver_sql("SELECT DATABASE()").scalar()

    @reflection.cache
    def has_table(self, connection, table_name, schema=None, **kw):
        self._ensure_has_table_connection(connection)

        if schema is None:
            schema = self.default_schema_name

        assert schema is not None

        full_name = ".".join(
            self.identifier_preparer._quote_free_identifiers(
                schema, table_name
            )
        )

        # DESCRIBE *must* be used because there is no information schema
        # table that returns information on temp tables that is consistently
        # available on MariaDB / MySQL / engine-agnostic etc.
        # therefore we have no choice but to use DESCRIBE and an error catch
        # to detect "False".  See issue #9058

        try:
            with connection.exec_driver_sql(
                    f"DESCRIBE {full_name}",
                    execution_options={"skip_user_error_events": True},
            ) as rs:
                return rs.fetchone() is not None
        except exc.DBAPIError as e:
            # https://dev.mysql.com/doc/mysql-errors/8.0/en/server-error-reference.html  # noqa: E501
            # there are a lot of codes that *may* pop up here at some point
            # but we continue to be fairly conservative.  We include:
            # 1146: Table '%s.%s' doesn't exist - what every MySQL has emitted
            # for decades
            #
            # mysql 8 suddenly started emitting:
            # 1049: Unknown database '%s'  - for nonexistent schema
            #
            # also added:
            # 1051: Unknown table '%s' - not known to emit
            #
            # there's more "doesn't exist" kinds of messages but they are
            # less clear if mysql 8 would suddenly start using one of those
            if self._extract_error_code(e.orig) in (1146, 1049, 1051):
                return False
            raise

    @reflection.cache
    def get_schema_names(self, connection, **kw):
        rp = connection.exec_driver_sql("SHOW schemas")
        return [r[0] for r in rp]

    @reflection.cache
    def get_table_names(self, connection, schema=None, **kw):
        """Return a Unicode SHOW TABLES from a given schema."""
        if schema is not None:
            current_schema = schema
        else:
            current_schema = self.default_schema_name

        charset = self._connection_charset

        rp = connection.exec_driver_sql(
            "SHOW FULL TABLES FROM %s"
            % self.identifier_preparer.quote_identifier(current_schema)
        )

        return [
            row[0]
            for row in self._compat_fetchall(rp, charset=charset)
            if row[1] == "BASE TABLE"
        ]

    @reflection.cache
    def get_view_names(self, connection, schema=None, **kw):
        if schema is None:
            schema = self.default_schema_name
        charset = self._connection_charset
        rp = connection.exec_driver_sql(
            "SHOW FULL TABLES FROM %s"
            % self.identifier_preparer.quote_identifier(schema)
        )
        return [
            row[0]
            for row in self._compat_fetchall(rp, charset=charset)
            if row[1] in ("VIEW", "SYSTEM VIEW")
        ]

    def _fetch_setting(self, connection, setting_name):
        charset = self._connection_charset

        if self.server_version_info and self.server_version_info < (5, 6):
            sql = "SHOW VARIABLES LIKE '%s'" % setting_name
            fetch_col = 1
        else:
            sql = "SELECT @@%s" % setting_name
            fetch_col = 0

        show_var = connection.exec_driver_sql(sql)
        row = self._compat_first(show_var, charset=charset)
        if not row:
            return None
        else:
            return row[fetch_col]

    def _detect_collations(self, connection):
        """Pull the active COLLATIONS list from the server.

        Cached per-connection.
        """

        collations = {}
        charset = self._connection_charset
        rs = connection.exec_driver_sql("SHOW COLLATION")
        for row in self._compat_fetchall(rs, charset):
            collations[row[0]] = row[1]
        return collations

    def _show_create_table(
            self, connection, table, charset=None, full_name=None
    ):
        """Run SHOW CREATE TABLE for a ``Table``."""

        if full_name is None:
            full_name = self.identifier_preparer.format_table(table)
        st = "SHOW CREATE TABLE %s" % full_name

        rp = None
        try:
            rp = connection.execution_options(
                skip_user_error_events=True
            ).exec_driver_sql(st)
        except exc.DBAPIError as e:
            if self._extract_error_code(e.orig) == 1146:
                raise exc.NoSuchTableError(full_name) from e
            else:
                raise
        row = self._compat_first(rp, charset=charset)
        if not row:
            raise exc.NoSuchTableError(full_name)
        return row[1].strip()

    def _describe_table(self, connection, table, charset=None, full_name=None):
        """Run DESCRIBE for a ``Table`` and return processed rows."""

        if full_name is None:
            full_name = self.identifier_preparer.format_table(table)
        st = "DESCRIBE %s" % full_name

        rp, rows = None, None
        try:
            try:
                rp = connection.execution_options(
                    skip_user_error_events=True
                ).exec_driver_sql(st)
            except exc.DBAPIError as e:
                code = self._extract_error_code(e.orig)
                if code == 1146:
                    raise exc.NoSuchTableError(full_name) from e

                elif code == 1356:
                    raise exc.UnreflectableTableError(
                        "Table or view named %s could not be "
                        "reflected: %s" % (full_name, e)
                    ) from e

                else:
                    raise
            rows = self._compat_fetchall(rp, charset=charset)
        finally:
            if rp:
                rp.close()
        return rows
