# Copyright 2020 Catalyst Cloud
#
#    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
#
#        http://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 oslo_log import log as logging
from pymysql import err as pymysql_err
from sqlalchemy import exc
from sqlalchemy.sql.expression import text

from trove.guestagent.common import sql_query

LOG = logging.getLogger(__name__)
FLUSH = text(sql_query.FLUSH)


class SqlClient(object):
    """A sqlalchemy wrapper to manage transactions."""

    def __init__(self, engine, use_flush=False):
        self.engine = engine
        self.use_flush = use_flush

    def __enter__(self):
        self.conn = self.engine.connect()
        self.trans = self.conn.begin()
        return self

    def __exit__(self, type, value, traceback):
        if self.trans:
            if type is not None:
                self.trans.rollback()
            else:
                if self.use_flush:
                    self.conn.execute(FLUSH)
                self.trans.commit()
        self.conn.close()

    def execute(self, t, **kwargs):
        LOG.debug('Execute SQL: %s', t)
        if isinstance(t, str):
            t = text(t)
        try:
            return self.conn.execute(t, kwargs)
        except Exception as err:
            LOG.error(f'Failed to execute SQL {t}, error: {err}')
            self.trans.rollback()
            self.trans = None
            raise err


def connection_checkout(dbapi_con, con_record, con_proxy):
    """
    A connection pool listener that ensures live connections are returned
    from the connection pool at checkout. This alleviates the problem of
    MySQL connections timing out.
    """
    try:
        try:
            dbapi_con.ping(False)
        except TypeError:
            dbapi_con.ping()
    except dbapi_con.OperationalError as ex:
        if ex.args[0] in (2006, 2013, 2014, 2045, 2055):
            raise exc.DisconnectionError()
        else:
            raise
    # MariaDB seems to timeout the client in a different
    # way than MySQL and PXC
    except pymysql_err.InternalError as ex:
        if "Packet sequence number wrong" in str(ex):
            raise exc.DisconnectionError()
        elif 'Connection was killed' in str(ex):
            raise exc.DisconnectionError()
        else:
            raise
