<commit_before>from __future__ import absolute_import
from __future__ import print_function

import inspect
import logging
import sys
import os
import getpass
import six
import uuid
from collections import defaultdict
import functools

import dataset
import pandas as pd

import socket
import os.path
import yaml
from six.moves import input

try:
    from traitlets import (
        HasTraits, CUnicode, List, CInt, Instance, Enum,
        CFloat, CBool)
except ImportError:
    from IPython.utils.traitlets import (
        HasTraits, CUnicode, List, CInt, Instance, Enum,
        CFloat, CBool)


logger = logging.getLogger(__name__)

# Whether we are running from NERSC
ON_NERSC = 'METATLAS_LOCAL' not in os.environ
logger.info('NERSC=%s', ON_NERSC)

# Observable List from
# http://stackoverflow.com/a/13259435

def callback_method(func):
    def notify(self, *args, **kwargs):
        if not hasattr(self, '_callbacks'):
            return func(self, *args, **kwargs)
        for _, callback in self._callbacks:
            callback()
        return func(self, *args, **kwargs)
    return notify


class NotifyList(list):
    extend = callback_method(list.extend)
    append = callback_method(list.append)
    remove = callback_method(list.remove)
    pop = callback_method(list.pop)
    __delitem__ = callback_method(list.__delitem__)
    __setitem__ = callback_method(list.__setitem__)
    __iadd__ = callback_method(list.__iadd__)
    __imul__ = callback_method(list.__imul__)

    def __getitem__(self, item):
        if isinstance(item, slice):
            return self.__class__(list.__getitem__(self, item))
        else:
            return list.__getitem__(self, item)

    def __init__(self, *args):
        list.__init__(self, *args)
        self._callbacks = []
        self._callback_cntr = 0

    def register_callback(self, cb):
        self._callbacks.append((self._callback_cntr, cb))
        self._callback_cntr += 1
        return self._callback_cntr - 1

    def unregister_callback(self, cbid):
        for idx, (i, cb) in enumerate(self._callbacks):
            if i == cbid:
                self._callbacks.pop(idx)
                return cb
        else:
            return None

def set_docstring(cls):
    """Set the docstring for a MetatlasObject object"""
    doc = cls.__doc__
    if not doc:
        doc = cls.__name__ + ' object.'
    doc += '\n\nParameters\n----------\n'
    for (tname, trait) in sorted(cls.class_traits().items()):
        if tname.startswith('_'):
            continue
        descr = trait.__class__.__name__.lower()
        if descr.startswith('c'):
            descr = descr[1:]
        elif descr == 'enum':
            descr = '{' + ', '.join(trait.values) + '}'
        doc += '%s: %s\n' % (tname, descr)
        help_text = trait.help#get_metadata('help')
        if not help_text:
            help_text = '%s value.' % tname
        help_text = help_text.strip()
        if help_text.endswith('.'):
            help_text = help_text[:-1]
        if trait.metadata.get('readonly', False):
            help_text += ' (read only)'
        help_text += '.'
        doc += '    %s\n' % help_text
    cls.__doc__ = doc
    return cls


def _get_subclasses(cls):
    return cls.__subclasses__() + [g for s in cls.__subclasses__()
                                   for g in _get_subclasses(s)]

class Workspace(object):
    instance = None

    def __init__(self):
        # get metatlas directory since notebooks and scripts could be launched
        # from other locations
        # this directory contains the config files
        metatlas_dir = os.path.dirname(sys.modules[self.__class__.__module__].__file__)
        if ON_NERSC:
            with open(os.path.join(metatlas_dir, 'nersc_config', 'nersc.yml')) as fid:
                nersc_info = yaml.safe_load(fid)

            with open(nersc_info['db_passwd_file']) as fid:
                pw = fid.read().strip()
                self.path = 'mysql+pymysql://meta_atlas_admin:%<EMAIL>/%s' % (pw, nersc_info['db_name'])
        else:
            local_config_file = os.path.join(metatlas_dir, 'local_config', 'local.yml')
            if os.path.isfile(local_config_file):
                with open(local_config_file) as fid:
                    local_info = yaml.load(fid)
                hostname = 'localhost' if 'db_hostname' not in local_info else local_info['db_hostname']
                login = ''
                if 'db_username' in local_info:
                    if 'db_password' in local_info:
                        login = f"{local_info['db_username']}:{local_info['db_password']}@"
                    else:
                        login = f"{local_info['db_username']}@"
                self.path = f"mysql+pymysql://{login}{hostname}/{local_info['db_name']}"
            else:
                filename = f"{getpass.getuser()}_workspace.db"
                self.path = f"sqlite:///{filename}"
                if os.path.exists(filename):
                    os.chmod(filename, 0o775)
        logger.debug('Using database at: %s', self.path)

        self.tablename_lut = dict()
        self.subclass_lut = dict()
        from .metatlas_objects import MetatlasObject
        for klass in _get_subclasses(MetatlasObject):
            name = klass.__name__.lower()
            self.subclass_lut[name] = klass
            if name.endswith('s'):
                self.subclass_lut[name + 'es'] = klass
                self.tablename_lut[klass] = name + 'es'
            else:
                self.subclass_lut[name + 's'] = klass
                self.tablename_lut[klass] = name + 's'
        # handle circular references
        self.seen = dict()
        Workspace.instance = self
        self.db = None

    @classmethod
    def get_instance(cls):
        """Returns a existing instance of Workspace or creates a new one"""
        if Workspace.instance is None:
            return Workspace()
        return Workspace.instance

    def get_connection(self):
        """
        Get a re-useable connection to the database.
        Each activity that queries the database needs to have this function preceeding it.
        """
        if self.db is None:
            self.db = dataset.connect(self.path)
        else:
            self.db.begin()
            try:
                self.db.query('SELECT 1')
                self.db.commit()
            except Exception:
                self.db.rollback()
                self.db = dataset.connect(self.path)
        assert self.db is not None

    def close_connection(self):
        """close database connections"""
        if self.db is not None:
            try:
                self.db.close()
            except AttributeError:
                logger.debug('Could not close DB. %s object contains %s', self.db.__class__, dir(self.db))
            self.db = None

    def convert_to_double(self, table, entry):
        """Convert a table column to double type."""
        self.get_connection()
        self.db.begin()
        try:
            self.db.query('alter table `%s` modify `%s` double' % (table, entry))
            self.db.commit()
        except Exception as err:
            rollback_and_log(self.db, err)

    def save_objects(self, objects, _override=False):
        """Save objects to the database"""
        logger.debug('Entering Workspace.save_objects')
        if not isinstance(objects, (list, set)):
            objects = [objects]<commit_msg>set mysql pool lifetime (#507)<commit_after>from __future__ import absolute_import
from __future__ import print_function

import inspect
import logging
import sys
import os
import getpass
import six
import uuid
from collections import defaultdict
import functools

import dataset
import pandas as pd

import socket
import os.path
import yaml
from six.moves import input

try:
    from traitlets import (
        HasTraits, CUnicode, List, CInt, Instance, Enum,
        CFloat, CBool)
except ImportError:
    from IPython.utils.traitlets import (
        HasTraits, CUnicode, List, CInt, Instance, Enum,
        CFloat, CBool)


logger = logging.getLogger(__name__)

# Whether we are running from NERSC
ON_NERSC = 'METATLAS_LOCAL' not in os.environ
logger.info('NERSC=%s', ON_NERSC)

# Observable List from
# http://stackoverflow.com/a/13259435

def callback_method(func):
    def notify(self, *args, **kwargs):
        if not hasattr(self, '_callbacks'):
            return func(self, *args, **kwargs)
        for _, callback in self._callbacks:
            callback()
        return func(self, *args, **kwargs)
    return notify


class NotifyList(list):
    extend = callback_method(list.extend)
    append = callback_method(list.append)
    remove = callback_method(list.remove)
    pop = callback_method(list.pop)
    __delitem__ = callback_method(list.__delitem__)
    __setitem__ = callback_method(list.__setitem__)
    __iadd__ = callback_method(list.__iadd__)
    __imul__ = callback_method(list.__imul__)

    def __getitem__(self, item):
        if isinstance(item, slice):
            return self.__class__(list.__getitem__(self, item))
        else:
            return list.__getitem__(self, item)

    def __init__(self, *args):
        list.__init__(self, *args)
        self._callbacks = []
        self._callback_cntr = 0

    def register_callback(self, cb):
        self._callbacks.append((self._callback_cntr, cb))
        self._callback_cntr += 1
        return self._callback_cntr - 1

    def unregister_callback(self, cbid):
        for idx, (i, cb) in enumerate(self._callbacks):
            if i == cbid:
                self._callbacks.pop(idx)
                return cb
        else:
            return None

def set_docstring(cls):
    """Set the docstring for a MetatlasObject object"""
    doc = cls.__doc__
    if not doc:
        doc = cls.__name__ + ' object.'
    doc += '\n\nParameters\n----------\n'
    for (tname, trait) in sorted(cls.class_traits().items()):
        if tname.startswith('_'):
            continue
        descr = trait.__class__.__name__.lower()
        if descr.startswith('c'):
            descr = descr[1:]
        elif descr == 'enum':
            descr = '{' + ', '.join(trait.values) + '}'
        doc += '%s: %s\n' % (tname, descr)
        help_text = trait.help#get_metadata('help')
        if not help_text:
            help_text = '%s value.' % tname
        help_text = help_text.strip()
        if help_text.endswith('.'):
            help_text = help_text[:-1]
        if trait.metadata.get('readonly', False):
            help_text += ' (read only)'
        help_text += '.'
        doc += '    %s\n' % help_text
    cls.__doc__ = doc
    return cls


def _get_subclasses(cls):
    return cls.__subclasses__() + [g for s in cls.__subclasses__()
                                   for g in _get_subclasses(s)]

class Workspace(object):
    instance = None

    def __init__(self):
        # get metatlas directory since notebooks and scripts could be launched
        # from other locations
        # this directory contains the config files
        metatlas_dir = os.path.dirname(sys.modules[self.__class__.__module__].__file__)
        if ON_NERSC:
            with open(os.path.join(metatlas_dir, 'nersc_config', 'nersc.yml')) as fid:
                nersc_info = yaml.safe_load(fid)

            with open(nersc_info['db_passwd_file']) as fid:
                pw = fid.read().strip()
                self.path = 'mysql+pymysql://meta_atlas_admin:%s@nerscdb04.nersc.gov/%s' % (pw, nersc_info['db_name'])
        else:
            local_config_file = os.path.join(metatlas_dir, 'local_config', 'local.yml')
            if os.path.isfile(local_config_file):
                with open(local_config_file) as fid:
                    local_info = yaml.load(fid)
                hostname = 'localhost' if 'db_hostname' not in local_info else local_info['db_hostname']
                login = ''
                if 'db_username' in local_info:
                    if 'db_password' in local_info:
                        login = f"{local_info['db_username']}:{local_info['db_password']}@"
                    else:
                        login = f"{local_info['db_username']}@"
                self.path = f"mysql+pymysql://{login}{hostname}/{local_info['db_name']}"
            else:
                filename = f"{getpass.getuser()}_workspace.db"
                self.path = f"sqlite:///{filename}"
                if os.path.exists(filename):
                    os.chmod(filename, 0o775)
        logger.debug('Using database at: %s', self.path)
        self.engine_kwargs = {} if self.path.startswith("sqlite") else {"pool_recycle": 3600}

        self.tablename_lut = dict()
        self.subclass_lut = dict()
        from .metatlas_objects import MetatlasObject
        for klass in _get_subclasses(MetatlasObject):
            name = klass.__name__.lower()
            self.subclass_lut[name] = klass
            if name.endswith('s'):
                self.subclass_lut[name + 'es'] = klass
                self.tablename_lut[klass] = name + 'es'
            else:
                self.subclass_lut[name + 's'] = klass
                self.tablename_lut[klass] = name + 's'
        # handle circular references
        self.seen = dict()
        Workspace.instance = self
        self.db = None

    @classmethod
    def get_instance(cls):
        """Returns a existing instance of Workspace or creates a new one"""
        if Workspace.instance is None:
            return Workspace()
        return Workspace.instance

    def get_connection(self):
        """
        Get a re-useable connection to the database.
        Each activity that queries the database needs to have this function preceeding it.
        """
        if self.db is None:
            self.db = dataset.connect(self.path, engine_kwargs=self.engine_kwargs)
        else:
            self.db.begin()
            try:
                self.db.query('SELECT 1')
                self.db.commit()
            except Exception:
                self.db.rollback()
                self.db = dataset.connect(self.path, engine_kwargs=self.engine_kwargs)
        assert self.db is not None

    def close_connection(self):
        """close database connections"""
        if self.db is not None:
            try:
                self.db.close()
            except AttributeError:
                logger.debug('Could not close DB. %s object contains %s', self.db.__class__, dir(self.db))
            self.db = None

    def convert_to_double(self, table, entry):
        """Convert a table column to double type."""
        self.get_connection()
        self.db.begin()
        try:
            self.db.query('alter table `%s` modify `%s` double' % (table, entry))
            self.db.commit()
        except Exception as err:
            rollback_and_log(self.db, err)

    def save_objects(self, objects, _override=False):
        """Save objects to the database"""
        logger.debug('Entering Workspace.save_objects')
        if not isinstance(objects, (list, set)):
            objects = [objects]
{'fraction_non_alphanumeric': 0.10910658619088678, 'fraction_numerical': 0.0029846786495987265, 'mean_word_length': 2.110148514851485, 'pattern_counts': {'":': 1, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5383478', 'n_tokens_mistral': 4513, 'n_tokens_neox': 4394, 'n_words': 1097}