# copyright 2003-2016 LOGILAB S.A. (Paris, FRANCE), all rights reserved.

# contact http://www.logilab.fr/ -- mailto:contact@logilab.fr

#

# This file is part of CubicWeb.

#

# CubicWeb is free software: you can redistribute it and/or modify it under the

# terms of the GNU Lesser General Public License as published by the Free

# Software Foundation, either version 2.1 of the License, or (at your option)

# any later version.

#

# CubicWeb is distributed in the hope that it will be useful, but WITHOUT

# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS

# FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more

# details.

#

# You should have received a copy of the GNU Lesser General Public License along

# with CubicWeb.  If not, see <http://www.gnu.org/licenses/>.

"""SQL utilities functions and classes."""

import os

import sys

import re

import subprocess

from os.path import abspath

from logging import getLogger

from datetime import time, datetime, timedelta

from pytz import utc

from logilab import database as db, common as lgc

from logilab.common.shellutils import ProgressBar, DummyProgressBar

from logilab.common.logging_ext import set_log_methods

from logilab.common.date import utctime, utcdatetime, strptime

from logilab.database.sqlgen import SQLGenerator

from cubicweb import Binary, ConfigurationError

from cubicweb.uilib import remove_html_tags

from cubicweb.schema import PURE_VIRTUAL_RTYPES

from cubicweb.server import SQL_CONNECT_HOOKS

from cubicweb.server.utils import crypt_password

lgc.USE_MX_DATETIME = False

SQL_PREFIX = 'cw_'

def _run_command(cmd, extra_env=None):

    env = os.environ.copy()

    for key, value in (extra_env or {}).items():

        env.setdefault(key, value)

        print(cmd)

        return subprocess.call(cmd, shell=True, env=env)

    else:

        print(' '.join(cmd))

        return subprocess.call(cmd, env=env)

def sqlexec(sqlstmts, cursor_or_execute, withpb=True,

            pbtitle='', delimiter=';', cnx=None):

    """execute sql statements ignoring DROP/ CREATE GROUP or USER statements

    error.

    :sqlstmts_as_string: a string or a list of sql statements.

    :cursor_or_execute: sql cursor or a callback used to execute statements

    :cnx: if given, commit/rollback at each statement.

    :withpb: if True, display a progresse bar

    :pbtitle: a string displayed as the progress bar title (if `withpb=True`)

    :delimiter: a string used to split sqlstmts (if it is a string)

    Return the failed statements (same type as sqlstmts)

    """

    if hasattr(cursor_or_execute, 'execute'):

        execute = cursor_or_execute.execute

    else:

        execute = cursor_or_execute

    sqlstmts_as_string = False

        sqlstmts_as_string = True

        sqlstmts = sqlstmts.split(delimiter)

    if withpb:

        if sys.stdout.isatty():

            pb = ProgressBar(len(sqlstmts), title=pbtitle)

        else:

            pb = DummyProgressBar()

    failed = []

    for sql in sqlstmts:

        sql = sql.strip()

        if withpb:

            pb.update()

        if not sql:

            continue

        try:

            # some dbapi modules doesn't accept unicode for sql string

            execute(str(sql))

        except Exception as ex:

            print(ex, file=sys.stderr)

            if cnx:

                cnx.rollback()

            failed.append(sql)

        else:

            if cnx:

                cnx.commit()

    if withpb:

        print()

    if sqlstmts_as_string:

        failed = delimiter.join(failed)

    return failed

def sqlgrants(schema, driver, user,

              text_index=True, set_owner=True,

              skip_relations=(), skip_entities=()):

    """Return a list of SQL statements to give all access privileges to the given user on the

    database.

    """

    from cubicweb.server.schema2sql import grant_schema

    from cubicweb.server.sources import native

    stmts = list(native.grant_schema(user, set_owner))

    if text_index:

        dbhelper = db.get_db_helper(driver)

        # XXX should return a list of sql statements rather than ';' joined statements

        stmts += dbhelper.sql_grant_user_on_fti(user).split(';')

    stmts += grant_schema(schema, user, set_owner, skip_entities=skip_entities, prefix=SQL_PREFIX)

    return stmts

def sqlschema(schema, driver, text_index=True,

              user=None, set_owner=False,

              skip_relations=PURE_VIRTUAL_RTYPES, skip_entities=()):

    """Return the database SQL schema as a list of SQL statements, according to the given parameters.

    """

    from cubicweb.server.schema2sql import schema2sql

    from cubicweb.server.sources import native

    if set_owner:

        assert user, 'user is argument required when set_owner is true'

    stmts = list(native.sql_schema(driver))

    dbhelper = db.get_db_helper(driver)

    if text_index:

        stmts += dbhelper.sql_init_fti().split(';')  # XXX

    stmts += schema2sql(dbhelper, schema, prefix=SQL_PREFIX,

                        skip_entities=skip_entities,

                        skip_relations=skip_relations)

    if dbhelper.users_support and user:

        stmts += sqlgrants(schema, driver, user, text_index, set_owner,

                           skip_relations, skip_entities)

    return stmts

_SQL_DROP_ALL_USER_TABLES_FILTER_FUNCTION = re.compile('^(?!(sql|pg)_)').match

def sql_drop_all_user_tables(driver_or_helper, sqlcursor):

    """Return ths sql to drop all tables found in the database system."""

    if not getattr(driver_or_helper, 'list_tables', None):

        dbhelper = db.get_db_helper(driver_or_helper)

    else:

        dbhelper = driver_or_helper

    stmts = [dbhelper.sql_drop_sequence('entities_id_seq')]

    # for mssql, we need to drop views before tables

    if hasattr(dbhelper, 'list_views'):

        stmts += ['DROP VIEW %s;' % name

                  for name in filter(_SQL_DROP_ALL_USER_TABLES_FILTER_FUNCTION,

                                     dbhelper.list_views(sqlcursor))]

    stmts += ['DROP TABLE %s;' % name

              for name in filter(_SQL_DROP_ALL_USER_TABLES_FILTER_FUNCTION,

                                 dbhelper.list_tables(sqlcursor))]

    return stmts

class ConnectionWrapper(object):

    """Wrap a connection to the system source's database, attempting to handle

    automatic reconnection.

    """

    # since 3.19, we only have to manage the system source connection

    def __init__(self, system_source):

        # dictionary of (source, connection), indexed by sources'uri

        self._source = system_source

        self.cnx = system_source.get_connection()

        self.cu = self.cnx.cursor()

    def commit(self):

        """commit the current transaction for this user"""

        # let exception propagates

        self.cnx.commit()

    def rollback(self):

        """rollback the current transaction for this user"""

        # catch exceptions, rollback other sources anyway

        try:

            self.cnx.rollback()

        except Exception:

            self._source.critical('rollback error', exc_info=sys.exc_info())

            # error on rollback, the connection is much probably in a really

            # bad state. Replace it by a new one.

            self.reconnect()

    def close(self, i_know_what_i_do=False):

        """close all connections in the set"""

        if i_know_what_i_do is not True:  # unexpected closing safety belt

            raise RuntimeError('connections set shouldn\'t be closed')

        try:

            self.cu.close()

            self.cu = None

        except Exception:

            pass

        try:

            self.cnx.close()

            self.cnx = None

        except Exception:

            pass

    # internals ###############################################################

    def cnxset_freed(self):

        """connections set is being freed from a session"""

        pass  # no nothing by default

    def reconnect(self):

        """reopen a connection for this source or all sources if none specified

        """

        try:

            # properly close existing connection if any

            self.cnx.close()

        except Exception:

            pass

        self._source.info('trying to reconnect')

        self.cnx = self._source.get_connection()

        self.cu = self.cnx.cursor()

class SqliteConnectionWrapper(ConnectionWrapper):

    """Sqlite specific connection wrapper: close the connection each time it's

    freed (and reopen it later when needed)

    """

    def __init__(self, system_source):

        # don't call parent's __init__, we don't want to initiate the connection

        self._source = system_source

    _cnx = None

    def cnxset_freed(self):

        self.cu.close()

        self.cnx.close()

        self.cnx = self.cu = None

    @property

    def cnx(self):

        if self._cnx is None:

            self._cnx = self._source.get_connection()

            self._cu = self._cnx.cursor()

        return self._cnx

    @cnx.setter

    def cnx(self, value):

        self._cnx = value

    @property

    def cu(self):

        if self._cnx is None:

            self._cnx = self._source.get_connection()

            self._cu = self._cnx.cursor()

        return self._cu

    @cu.setter

    def cu(self, value):

        self._cu = value

class SQLAdapterMixIn(object):

    """Mixin for SQL data sources, getting a connection from a configuration

    dictionary and handling connection locking

    """

    cnx_wrap = ConnectionWrapper

    def __init__(self, source_config, repairing=False):

        try:

            self.dbdriver = source_config['db-driver'].lower()

            dbname = source_config['db-name']

        except KeyError:

            raise ConfigurationError('missing some expected entries in sources file')

        dbhost = source_config.get('db-host')

        port = source_config.get('db-port')

        dbport = port and int(port) or None

        dbuser = source_config.get('db-user')

        dbpassword = source_config.get('db-password')

        dbencoding = source_config.get('db-encoding', 'UTF-8')

        dbextraargs = source_config.get('db-extra-arguments')

        dbnamespace = source_config.get('db-namespace')

        self.dbhelper = db.get_db_helper(self.dbdriver)

        self.dbhelper.record_connection_info(dbname, dbhost, dbport, dbuser,

                                             dbpassword, dbextraargs,

                                             dbencoding, dbnamespace)

        self.sqlgen = SQLGenerator()

        # copy back some commonly accessed attributes

        dbapi_module = self.dbhelper.dbapi_module

        self.OperationalError = dbapi_module.OperationalError

        self.InterfaceError = dbapi_module.InterfaceError

        self.DbapiError = dbapi_module.Error

        self._binary = self.dbhelper.binary_value

        self._process_value = dbapi_module.process_value

        self._dbencoding = dbencoding

        if self.dbdriver == 'sqlite':

            self.cnx_wrap = SqliteConnectionWrapper

            self.dbhelper.dbname = abspath(self.dbhelper.dbname)

        if not repairing:

            statement_timeout = int(source_config.get('db-statement-timeout', 0))

            if statement_timeout > 0:

                def set_postgres_timeout(cnx):

                    cnx.cursor().execute('SET statement_timeout to %d' % statement_timeout)

                    cnx.commit()

                postgres_hooks = SQL_CONNECT_HOOKS['postgres']

                postgres_hooks.append(set_postgres_timeout)

    def wrapped_connection(self):

        """open and return a connection to the database, wrapped into a class

        handling reconnection and all

        """

        return self.cnx_wrap(self)

    def get_connection(self):

        """open and return a connection to the database"""

        return self.dbhelper.get_connection()

    def _backup_restore_env(self):

        if (self.config['db-driver'] == 'postgres'

                and self.config['db-password'] is not None):

            return {'PGPASSWORD': self.config['db-password']}

    def backup_to_file(self, backupfile, confirm):

        extra_env = self._backup_restore_env()

        for cmd in self.dbhelper.backup_commands(backupfile,

                                                 keepownership=False):

            if _run_command(cmd, extra_env=extra_env):

                if not confirm('   [Failed] Continue anyway?', default='n'):

                    raise Exception('Failed command: %s' % cmd)

    def restore_from_file(self, backupfile, confirm, drop=True):

        extra_env = self._backup_restore_env()

        for cmd in self.dbhelper.restore_commands(backupfile,

                                                  keepownership=False,

                                                  drop=drop):

            if _run_command(cmd, extra_env=extra_env):

                if not confirm('   [Failed] Continue anyway?', default='n'):

                    raise Exception('Failed command: %s' % cmd)

    def merge_args(self, args, query_args):

        if args is not None:

            newargs = {}

            for key, val in args.items():

                # convert cubicweb binary into db binary

                if isinstance(val, Binary):

                    val = self._binary(val.getvalue())

                # convert timestamp to utc.

                # expect SET TiME ZONE to UTC at connection opening time.

                # This shouldn't change anything for datetime without TZ.

                elif isinstance(val, datetime) and val.tzinfo is not None:

                    val = utcdatetime(val)

                elif isinstance(val, time) and val.tzinfo is not None:

                    val = utctime(val)

                newargs[key] = val

            # should not collide

            assert not (frozenset(newargs) & frozenset(query_args)), \

                'unexpected collision: %s' % (frozenset(newargs) & frozenset(query_args))

            newargs.update(query_args)

            return newargs

        return query_args

    def process_result(self, cursor, cnx=None, column_callbacks=None):

        """return a list of CubicWeb compliant values from data in the given cursor

        """

        return list(self.iter_process_result(cursor, cnx, column_callbacks))

    def iter_process_result(self, cursor, cnx, column_callbacks=None):

        """return a iterator on tuples of CubicWeb compliant values from data

        in the given cursor

        """

        # use two different implementations to avoid paying the price of