<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

##############################################################################
#  Copyright Kitware Inc.
#
#  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.
##############################################################################

import six
import sqlalchemy
import sqlalchemy.engine.reflection
import sqlalchemy.orm
import time

from six.moves import range

from girder import logger as log

from . import base
from .base import DatabaseConnectorException


MAX_SCHEMAS_IN_TABLE_LIST = 25

DatabaseOperators = {
    'eq': '=',
    'ne': '!=',
    'gte': '>=',
    'gt': '>',
    'lte': '<=',
    'lt': '<',
}


_enginePool = {}
_enginePoolMaxSize = 5


def adjustDBUrl(url):
    """
    Adjust a url to match the form sqlalchemy requires.  In general, the url is
    of the form dialect+driver://username:password@host:port/database.

    :param url: the url to adjust.
    :returns: the adjusted url
    """
    # The below code is disabled until a test can be made where it works.
    # sqlalchemy doesn't seem to permit this without some additional module
    # If the prefix is sql://, use the default generic-sql dialect
    # if url.startswith('sql://'):
    #     url = url.split('sql://', 1)[1]
    return url


def getEngine(url, **kwargs):
    """
    Get a sqlalchemy engine from a pool in case we use the same parameters for
    multiple connections.
    """
    key = (url, frozenset(six.viewitems(kwargs)))
    engine = _enginePool.get(key)
    if engine is None:
        engine = sqlalchemy.create_engine(url, **kwargs)
        if len(_enginePool) >= _enginePoolMaxSize:
            _enginePoolMaxSize.clear()
        _enginePool[key] = engine
    return engine


class SQLAlchemyConnector(base.DatabaseConnector):
    name = 'sqlalchemy'

    def __init__(self, *args, **kwargs):
        super(SQLAlchemyConnector, self).__init__(*args, **kwargs)
        self.table = kwargs.get('table')
        self.schema = kwargs.get('schema')
        self.dbEngine = None
        self.sessions = {}
        # dbparams can include values in http://www.postgresql.org/docs/
        #   current/static/libpq-connect.html#LIBPQ-PARAMKEYWORDS
        self.dbparams = kwargs.get('dbparams', {})
        self.databaseUrl = adjustDBUrl(kwargs.get('url'))

        # Additional parameters:
        #  idletime: seconds after which a connection is considered idle
        #  abandontime: seconds after which a connection will be abandoned
        self.dbIdleTime = float(kwargs.get('idletime', 300))
        self.dbAbandonTime = float(kwargs.get('abandontime',
                                   self.dbIdleTime * 5))
        self.databaseOperators = DatabaseOperators
        self.fields = None
        self.allowFieldFunctions = True
        self.allowSortFunctions = True
        self.allowFilterFunctions = True
        self.initialized = True
        self.types = {type: getattr(sqlalchemy, type) for type in dir(sqlalchemy)
                      if isinstance(getattr(sqlalchemy, type),
                                    sqlalchemy.sql.visitors.VisitableType)}

        class Table(object):
            """
            This is used to handle table properties from SQLAlchemy.
            """
            pass

        self.tableClass = Table
        self._allowedFunctions = {
            'cast': True,
            'count': True,
            'distinct': True,
        }

    def _addFilter(self, filterList, filter):
        """
        Add a filter to a list of SQLAlchemy filters.

        :param filterList: a list of SQLAlchemy filters which is modified.
        :param filter: information on the filter.
        :return: the modified list.
        """
        if 'group' in filter:
            sublist = []
            for subfilter in filter['value']:
                sublist = self._addFilter(sublist, subfilter)
            if filter['group'] == 'and':
                filterList.append(sqlalchemy.and_(*sublist))
            elif filter['group'] == 'or':
                filterList.append(sqlalchemy.or_(*sublist))
            return filterList
        operator = filter['operator']
        operator = base.FilterOperators.get(operator, operator)
        operator = self.databaseOperators.get(operator, operator)
        field = self._convertFieldOrFunction(filter['field'])
        negate = False
        if operator.startswith('not_'):
            negate = True
            operator = operator.split('not_', 1)[1]
        if operator == 'in':
            values = filter['value']
            if not isinstance(values, (list, tuple)):
                values = [values]
            values = [self._convertFieldOrFunction(value, True)
                      for value in values]
            opfunc = field.in_(values)
        elif operator == 'is':
            value = self._convertFieldOrFunction(filter['value'], True)
            opfunc = field.is_(value)
        else:
            value = self._convertFieldOrFunction(filter['value'], True)
            opfunc = field.op(operator)(value)
        if negate:
            opfunc = sqlalchemy.not_(opfunc)
        filterList.append(opfunc)
        return filterList

    def _convertFieldOrFunction(self, fieldOrFunction, preferValue=False):
        """
        Convert a string to a column reference, or a dictionary to a column or
        function reference.  If a function is passed, this should be a
        canonical function reference ('func' and 'param' are both populated).

        :param fieldOrFunction: a string with a column name or a dictionary
                                with either a field, function, or value.
        :param preferValue: if True then if fieldOrFunction is not a
                            dictionary, return it unchanged.
        :returns: a constructed column or function object, or a bare value.
        """
        if not isinstance(fieldOrFunction, dict):
            if preferValue:
                return fieldOrFunction
            return getattr(self.tableClass, fieldOrFunction)
        if 'field' in fieldOrFunction:
            return getattr(self.tableClass, fieldOrFunction['field'])
        if 'value' in fieldOrFunction:
            if not preferValue:
                return sqlalchemy.sql.elements.literal(
                    fieldOrFunction['value'])
            return fieldOrFunction['value']
        fieldOrFunction = self.isFunction(fieldOrFunction)
        if fieldOrFunction is False:
            raise DatabaseConnectorException('Not a function')
        if not self._isFunctionAllowed(fieldOrFunction['func']):
            raise DatabaseConnectorException('Function %s is not allowed' %
                                             fieldOrFunction['func'])
        param = fieldOrFunction.get('param', fieldOrFunction.get('params', []))
        # Determine the function we need to call to apply the function
        if fieldOrFunction['func'] in ('distinct', 'cast'):
            if (fieldOrFunction['func'] == 'cast' and len(param) == 2 and
                    isinstance(param[1], dict) and 'value' in param[1]):
                param[1]['value'] = self.types.get(param[1]['value'], param[1]['value'])
            funcfunc = getattr(sqlalchemy, fieldOrFunction['func'])
        else:
            funcfunc = getattr(sqlalchemy.func, fieldOrFunction['func'])
        return funcfunc(
            *[self._convertFieldOrFunction(entry, True) for entry in param])

    def _isFunctionAllowed(self, proname):
        """
        Check if the specified function is allowed.  Currently, only
        non-volatile functions are allowed, even though there are volatile
        functions that are harmless.

        :param proname: name of the function to check.
        :returns: True is allowed, False is not.
        """
        return self._allowedFunctions.get(proname, False)

    def connect(self, client=None):
        """
        Connect to the database.

        :param client: if None, use a new session.  If specified, if this
                       client is currently marked in use, cancel the client's
                       existing query and return a connection from the pool fo
                       r the client to use.
        :return: a SQLAlchemny session object.
        """
        if not self.dbEngine:
            self.dbEngine = getEngine(self.databaseUrl, **self.dbparams)
            metadata = sqlalchemy.MetaData(self.dbEngine)
            table = sqlalchemy.Table(self.table, metadata, schema=self.schema,
                                     autoload=True)

            # The orm.mapper is used to refer to our columns.  If the table or
            # view we are connecting to does not have any primary keys, the
            # mapper will fail.  Use the first column as a fallback; this is
            # only safe because we DON'T alter data; we have no guarantee we
            # can refer to a specific row (but we don't need to).
            fallbackPrimaryCol = None
            for col in table.c:
                if col.primary_key:
                    fallbackPrimaryCol = None
                    break
                if fallbackPrimaryCol is None:
                    fallbackPrimaryCol = col

            sqlalchemy.orm.mapper(
                self.tableClass, table, primary_key=fallbackPrimaryCol)
        # If we are asking for a specific client, clean up defunct clients
        curtime = time.time()
        if client:
            for oldsess in list(self.sessions):
                idle = curtime - self.sessions[oldsess]['last']
                if ((idle > self.dbIdleTime and
                        not self.sessions[oldsess]['used']) or
                        idle > self.dbAbandonTime):
                    # Close the session.  sqlalchemy keeps them too long
                    # otherwise
                    self.sessions[oldsess]['session'].close()<commit_msg>If we fail to connect or map a table, reconnect next time.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

##############################################################################
#  Copyright Kitware Inc.
#
#  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.
##############################################################################

import six
import sqlalchemy
import sqlalchemy.engine.reflection
import sqlalchemy.orm
import time

from six.moves import range

from girder import logger as log

from . import base
from .base import DatabaseConnectorException


MAX_SCHEMAS_IN_TABLE_LIST = 25

DatabaseOperators = {
    'eq': '=',
    'ne': '!=',
    'gte': '>=',
    'gt': '>',
    'lte': '<=',
    'lt': '<',
}


_enginePool = {}
_enginePoolMaxSize = 5


def adjustDBUrl(url):
    """
    Adjust a url to match the form sqlalchemy requires.  In general, the url is
    of the form dialect+driver://username:password@host:port/database.

    :param url: the url to adjust.
    :returns: the adjusted url
    """
    # The below code is disabled until a test can be made where it works.
    # sqlalchemy doesn't seem to permit this without some additional module
    # If the prefix is sql://, use the default generic-sql dialect
    # if url.startswith('sql://'):
    #     url = url.split('sql://', 1)[1]
    return url


def getEngine(url, **kwargs):
    """
    Get a sqlalchemy engine from a pool in case we use the same parameters for
    multiple connections.
    """
    key = (url, frozenset(six.viewitems(kwargs)))
    engine = _enginePool.get(key)
    if engine is None:
        engine = sqlalchemy.create_engine(url, **kwargs)
        if len(_enginePool) >= _enginePoolMaxSize:
            _enginePoolMaxSize.clear()
        _enginePool[key] = engine
    return engine


class SQLAlchemyConnector(base.DatabaseConnector):
    name = 'sqlalchemy'

    def __init__(self, *args, **kwargs):
        super(SQLAlchemyConnector, self).__init__(*args, **kwargs)
        self.table = kwargs.get('table')
        self.schema = kwargs.get('schema')
        self.dbEngine = None
        self.sessions = {}
        # dbparams can include values in http://www.postgresql.org/docs/
        #   current/static/libpq-connect.html#LIBPQ-PARAMKEYWORDS
        self.dbparams = kwargs.get('dbparams', {})
        self.databaseUrl = adjustDBUrl(kwargs.get('url'))

        # Additional parameters:
        #  idletime: seconds after which a connection is considered idle
        #  abandontime: seconds after which a connection will be abandoned
        self.dbIdleTime = float(kwargs.get('idletime', 300))
        self.dbAbandonTime = float(kwargs.get('abandontime',
                                   self.dbIdleTime * 5))
        self.databaseOperators = DatabaseOperators
        self.fields = None
        self.allowFieldFunctions = True
        self.allowSortFunctions = True
        self.allowFilterFunctions = True
        self.initialized = True
        self.types = {type: getattr(sqlalchemy, type) for type in dir(sqlalchemy)
                      if isinstance(getattr(sqlalchemy, type),
                                    sqlalchemy.sql.visitors.VisitableType)}

        class Table(object):
            """
            This is used to handle table properties from SQLAlchemy.
            """
            pass

        self.tableClass = Table
        self._allowedFunctions = {
            'cast': True,
            'count': True,
            'distinct': True,
        }

    def _addFilter(self, filterList, filter):
        """
        Add a filter to a list of SQLAlchemy filters.

        :param filterList: a list of SQLAlchemy filters which is modified.
        :param filter: information on the filter.
        :return: the modified list.
        """
        if 'group' in filter:
            sublist = []
            for subfilter in filter['value']:
                sublist = self._addFilter(sublist, subfilter)
            if filter['group'] == 'and':
                filterList.append(sqlalchemy.and_(*sublist))
            elif filter['group'] == 'or':
                filterList.append(sqlalchemy.or_(*sublist))
            return filterList
        operator = filter['operator']
        operator = base.FilterOperators.get(operator, operator)
        operator = self.databaseOperators.get(operator, operator)
        field = self._convertFieldOrFunction(filter['field'])
        negate = False
        if operator.startswith('not_'):
            negate = True
            operator = operator.split('not_', 1)[1]
        if operator == 'in':
            values = filter['value']
            if not isinstance(values, (list, tuple)):
                values = [values]
            values = [self._convertFieldOrFunction(value, True)
                      for value in values]
            opfunc = field.in_(values)
        elif operator == 'is':
            value = self._convertFieldOrFunction(filter['value'], True)
            opfunc = field.is_(value)
        else:
            value = self._convertFieldOrFunction(filter['value'], True)
            opfunc = field.op(operator)(value)
        if negate:
            opfunc = sqlalchemy.not_(opfunc)
        filterList.append(opfunc)
        return filterList

    def _convertFieldOrFunction(self, fieldOrFunction, preferValue=False):
        """
        Convert a string to a column reference, or a dictionary to a column or
        function reference.  If a function is passed, this should be a
        canonical function reference ('func' and 'param' are both populated).

        :param fieldOrFunction: a string with a column name or a dictionary
                                with either a field, function, or value.
        :param preferValue: if True then if fieldOrFunction is not a
                            dictionary, return it unchanged.
        :returns: a constructed column or function object, or a bare value.
        """
        if not isinstance(fieldOrFunction, dict):
            if preferValue:
                return fieldOrFunction
            return getattr(self.tableClass, fieldOrFunction)
        if 'field' in fieldOrFunction:
            return getattr(self.tableClass, fieldOrFunction['field'])
        if 'value' in fieldOrFunction:
            if not preferValue:
                return sqlalchemy.sql.elements.literal(
                    fieldOrFunction['value'])
            return fieldOrFunction['value']
        fieldOrFunction = self.isFunction(fieldOrFunction)
        if fieldOrFunction is False:
            raise DatabaseConnectorException('Not a function')
        if not self._isFunctionAllowed(fieldOrFunction['func']):
            raise DatabaseConnectorException('Function %s is not allowed' %
                                             fieldOrFunction['func'])
        param = fieldOrFunction.get('param', fieldOrFunction.get('params', []))
        # Determine the function we need to call to apply the function
        if fieldOrFunction['func'] in ('distinct', 'cast'):
            if (fieldOrFunction['func'] == 'cast' and len(param) == 2 and
                    isinstance(param[1], dict) and 'value' in param[1]):
                param[1]['value'] = self.types.get(param[1]['value'], param[1]['value'])
            funcfunc = getattr(sqlalchemy, fieldOrFunction['func'])
        else:
            funcfunc = getattr(sqlalchemy.func, fieldOrFunction['func'])
        return funcfunc(
            *[self._convertFieldOrFunction(entry, True) for entry in param])

    def _isFunctionAllowed(self, proname):
        """
        Check if the specified function is allowed.  Currently, only
        non-volatile functions are allowed, even though there are volatile
        functions that are harmless.

        :param proname: name of the function to check.
        :returns: True is allowed, False is not.
        """
        return self._allowedFunctions.get(proname, False)

    def connect(self, client=None):
        """
        Connect to the database.

        :param client: if None, use a new session.  If specified, if this
                       client is currently marked in use, cancel the client's
                       existing query and return a connection from the pool fo
                       r the client to use.
        :return: a SQLAlchemny session object.
        """
        if not self.dbEngine:
            engine = getEngine(self.databaseUri, **self.dbparams)
            metadata = sqlalchemy.MetaData(engine)
            table = sqlalchemy.Table(self.table, metadata, schema=self.schema,
                                     autoload=True)

            # The orm.mapper is used to refer to our columns.  If the table or
            # view we are connecting to does not have any primary keys, the
            # mapper will fail.  Use the first column as a fallback; this is
            # only safe because we DON'T alter data; we have no guarantee we
            # can refer to a specific row (but we don't need to).
            fallbackPrimaryCol = None
            for col in table.c:
                if col.primary_key:
                    fallbackPrimaryCol = None
                    break
                if fallbackPrimaryCol is None:
                    fallbackPrimaryCol = col

            sqlalchemy.orm.mapper(
                self.tableClass, table, primary_key=fallbackPrimaryCol)
            self.dbEngine = engine
        # If we are asking for a specific client, clean up defunct clients
        curtime = time.time()
        if client:
            for oldsess in list(self.sessions):
                idle = curtime - self.sessions[oldsess]['last']
                if ((idle > self.dbIdleTime and
                        not self.sessions[oldsess]['used']) or
                        idle > self.dbAbandonTime):
                    # Close the session.  sqlalchemy keeps them too long
                    # otherwise
                    self.sessions[oldsess]['session'].close()
{'fraction_non_alphanumeric': 0.10378171373863092, 'fraction_numerical': 0.0021062709430349448, 'mean_word_length': 2.1448140900195694, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4502299', 'n_tokens_mistral': 5464, 'n_tokens_neox': 5297, 'n_words': 1732}