<commit_before>"""
Queue database methods
"""

import logging
from datetime import datetime
from functools import partial,reduce
import operator
from collections import OrderedDict, Iterable
import math

from iceprod.core.util import Node_Resources
from iceprod.core.dataclasses import Number,String
from iceprod.core import serialization
from iceprod.core.jsonUtil import json_encode,json_decode,json_compressor

from iceprod.server import GlobalID
from iceprod.server.dbmethods import dbmethod,_Methods_Base,datetime2str,str2datetime,nowstr

logger = logging.getLogger('dbmethods.queue')

class queue(_Methods_Base):
    """
    The Queue DB methods.

    Takes a handle to a subclass of iceprod.server.modules.db.DBAPI
    as an argument.
    """

    def _queue_get_task_from_ret(self,ret):
        tasks = OrderedDict()
        try:
            for row in ret:
                dict_row = self._list_to_dict('task',row)
                dict_row['status_changed'] = str2datetime(dict_row['status_changed'])
                tasks[row[0]] = dict_row
        except:
            return {}
        else:
            return tasks

    @dbmethod
    def queue_get_site_id(self,callback=None):
        """Get the current site_id"""
        sql = 'select site_id from setting'
        bindings = tuple()
        cb = partial(self._queue_get_site_id_callback,callback=callback)
        self.db.sql_read_task(sql,bindings,callback=cb)
    def _queue_get_site_id_callback(self,ret,callback=None):
        if callback:
            if isinstance(ret,Exception):
                callback(ret)
            elif ret is None or len(ret) < 1 or len(ret[0]) < 1:
                callback(Exception('no site id'))
            else:
                callback(ret[0][0])

    @dbmethod
    def queue_set_site_queues(self,site_id,queues,callback=None):
        """Set the site queues"""
        cb = partial(self._queue_set_site_queues_blocking,site_id,queues,
                     callback=callback)
        self.db.blocking_task('queue',cb)
    def _queue_set_site_queues_blocking(self,site_id,queues,callback=None):
        conn,archive_conn = self.db._dbsetup()
        sql = 'select * from site where site_id = ?'
        bindings = (site_id,)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        elif len(ret) > 0 and len(ret[0]) > 0:
            # already a site entry, so just update
            try:
                old_site = self._list_to_dict('site',ret[0])
                old_queues = json_decode(old_site['queues'])
                for k in set(queues) & set(old_queues):
                    try:
                        for kk in set(old_queues[k]['resources']) - set(queues[k]['resources']):
                            queues[k]['resources'][kk] = old_queues[k]['resources'][kk]
                    except:
                        try:
                            queues[k]['resources'] = old_queues[k]['resources']
                        except:
                            queues[k]['resources'] = {}
                queues = json_encode(queues)
            except Exception as e:
                logger.warn('set_site_queues(): cannot encode queues to json')
                callback(e)
                return
            sql = 'update site set queues = ? where site_id = ?'
            bindings = (queues,site_id)
        else:
            # add a new site entry
            try:
                queues = json_encode(queues)
            except Exception as e:
                logger.warn('set_site_queues(): cannot encode queues to json')
                callback(e)
                return
            sql = 'insert into site (site_id,queues) values (?,?)'
            bindings = (site_id,queues)
        try:
            ret = self.db._db_write(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
        else:
            if self._is_master():
                sql3 = 'replace into master_update_history (table_name,update_index,timestamp) values (?,?,?)'
                bindings3 = ('site',site_id,nowstr())
                try:
                    self.db._db_write(conn,sql3,bindings3,None,None,None)
                except Exception as e:
                    logger.info('error updating master_update_history',
                                exc_info=True)
            else:
                self._send_to_master(('site',site_id,nowstr(),sql,bindings))
            callback(True)

    @dbmethod
    def queue_get_active_tasks(self,gridspec,callback=None):
        """Get a dict of active tasks (queued,processing,reset,resume) on this site and plugin,
           returning {status:{tasks}} where each task = join of search and task tables"""
        if not callback:
            return
        sql = 'select task_id from search '
        sql += 'where gridspec like ? '
        sql += ' and task_status in ("queued","processing","reset","resume")'
        bindings = ('%'+gridspec+'%',)
        cb = partial(self._queue_get_active_tasks_callback,callback=callback)
        self.db.sql_read_task(sql,bindings,callback=cb)
    def _queue_get_active_tasks_callback(self,ret,callback=None):
        if isinstance(ret,Exception):
            callback(ret)
        else:
            tasks = set(row[0] for row in ret)
            sql = 'select * from task where task_id in ('
            sql += ','.join('?' for _ in tasks)+')'
            bindings = tuple(tasks)
            cb = partial(self._queue_get_active_tasks_callback2,callback=callback)
            self.db.sql_read_task(sql,bindings,callback=cb)
    def _queue_get_active_tasks_callback2(self,ret,callback=None):
        if isinstance(ret,Exception):
            callback(ret)
        else:
            task_groups = {}
            if ret:
                tasks = self._queue_get_task_from_ret(ret)
                for task_id in tasks:
                    status = tasks[task_id]['status']
                    if status not in task_groups:
                        task_groups[status] = {}
                    task_groups[status][task_id] = tasks[task_id]
            callback(task_groups)

    @dbmethod
    def queue_get_grid_tasks(self,gridspec,callback=None):
        """Get a list of tasks (queued, processing) on this
           site and plugin, returning [tasks] where each task has
           task_id, grid_queue_id, submit_time, and submit_dir"""
        if not callback:
            return
        sql = 'select task_id from search '
        sql += 'where gridspec like ? '
        sql += ' and task_status in ("queued","processing")'
        bindings = ('%'+gridspec+'%',)
        cb = partial(self._queue_get_grid_tasks_callback,callback=callback)
        self.db.sql_read_task(sql,bindings,callback=cb)
    def _queue_get_grid_tasks_callback(self,ret,callback=None):
        if isinstance(ret,Exception):
            callback(ret)
        else:
            tasks = set(row[0] for row in ret)
            sql = 'select * from task where task_id in ('
            sql += ','.join('?' for _ in tasks)+')'
            bindings = tuple(tasks)
            cb = partial(self._queue_get_grid_tasks_callback2,callback=callback)
            self.db.sql_read_task(sql,bindings,callback=cb)
    def _queue_get_grid_tasks_callback2(self,ret,callback=None):
        if isinstance(ret,Exception):
            callback(ret)
        else:
            task_ret = []
            if ret:
                tasks = self._queue_get_task_from_ret(ret)
                for task_id in tasks:
                    task_ret.append({
                        'task_id': task_id,
                        'grid_queue_id': tasks[task_id]['grid_queue_id'],
                        'submit_time': str2datetime(tasks[task_id]['status_changed']),
                        'submit_dir': tasks[task_id]['submit_dir'],
                    })
            callback(task_ret)

    @dbmethod
    def queue_set_task_status(self,task,status,callback=None):
        """Set the status of a task"""
        if not isinstance(task,Iterable):
            raise Exception('task is not a str or iterable')
        cb = partial(self._queue_set_task_status_blocking,task,status,
                     callback=callback)
        self.db.non_blocking_task(cb)
    def _queue_set_task_status_blocking(self,task,status,callback=None):
        conn,archive_conn = self.db._dbsetup()
        now = nowstr()
        if isinstance(task,String):
            task = [task]
        elif not isinstance(task,Iterable):
            callback(Exception('unknown type for task'))
            return
        
        msql = 'update search set task_status = ? '
        msql += ' where task_id = ?'
        msql2 = 'update task set prev_status = status, '
        msql2 += ' status = ?, status_changed = ? where task_id = ?'

        # process in batches of 900
        try:
            while task:
                t = task[:900]
                task = task[900:]<commit_msg>fix more 1000 binding limit cases<commit_after>"""
Queue database methods
"""

import logging
from datetime import datetime
from functools import partial,reduce
import operator
from collections import OrderedDict, Iterable
import math

from iceprod.core.util import Node_Resources
from iceprod.core.dataclasses import Number,String
from iceprod.core import serialization
from iceprod.core.jsonUtil import json_encode,json_decode,json_compressor

from iceprod.server import GlobalID
from iceprod.server.dbmethods import dbmethod,_Methods_Base,datetime2str,str2datetime,nowstr

logger = logging.getLogger('dbmethods.queue')

class queue(_Methods_Base):
    """
    The Queue DB methods.

    Takes a handle to a subclass of iceprod.server.modules.db.DBAPI
    as an argument.
    """

    def _queue_get_task_from_ret(self,ret):
        tasks = OrderedDict()
        try:
            for row in ret:
                dict_row = self._list_to_dict('task',row)
                dict_row['status_changed'] = str2datetime(dict_row['status_changed'])
                tasks[row[0]] = dict_row
        except:
            return {}
        else:
            return tasks

    @dbmethod
    def queue_get_site_id(self,callback=None):
        """Get the current site_id"""
        sql = 'select site_id from setting'
        bindings = tuple()
        cb = partial(self._queue_get_site_id_callback,callback=callback)
        self.db.sql_read_task(sql,bindings,callback=cb)
    def _queue_get_site_id_callback(self,ret,callback=None):
        if callback:
            if isinstance(ret,Exception):
                callback(ret)
            elif ret is None or len(ret) < 1 or len(ret[0]) < 1:
                callback(Exception('no site id'))
            else:
                callback(ret[0][0])

    @dbmethod
    def queue_set_site_queues(self,site_id,queues,callback=None):
        """Set the site queues"""
        cb = partial(self._queue_set_site_queues_blocking,site_id,queues,
                     callback=callback)
        self.db.blocking_task('queue',cb)
    def _queue_set_site_queues_blocking(self,site_id,queues,callback=None):
        conn,archive_conn = self.db._dbsetup()
        sql = 'select * from site where site_id = ?'
        bindings = (site_id,)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        elif len(ret) > 0 and len(ret[0]) > 0:
            # already a site entry, so just update
            try:
                old_site = self._list_to_dict('site',ret[0])
                old_queues = json_decode(old_site['queues'])
                for k in set(queues) & set(old_queues):
                    try:
                        for kk in set(old_queues[k]['resources']) - set(queues[k]['resources']):
                            queues[k]['resources'][kk] = old_queues[k]['resources'][kk]
                    except:
                        try:
                            queues[k]['resources'] = old_queues[k]['resources']
                        except:
                            queues[k]['resources'] = {}
                queues = json_encode(queues)
            except Exception as e:
                logger.warn('set_site_queues(): cannot encode queues to json')
                callback(e)
                return
            sql = 'update site set queues = ? where site_id = ?'
            bindings = (queues,site_id)
        else:
            # add a new site entry
            try:
                queues = json_encode(queues)
            except Exception as e:
                logger.warn('set_site_queues(): cannot encode queues to json')
                callback(e)
                return
            sql = 'insert into site (site_id,queues) values (?,?)'
            bindings = (site_id,queues)
        try:
            ret = self.db._db_write(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
        else:
            if self._is_master():
                sql3 = 'replace into master_update_history (table_name,update_index,timestamp) values (?,?,?)'
                bindings3 = ('site',site_id,nowstr())
                try:
                    self.db._db_write(conn,sql3,bindings3,None,None,None)
                except Exception as e:
                    logger.info('error updating master_update_history',
                                exc_info=True)
            else:
                self._send_to_master(('site',site_id,nowstr(),sql,bindings))
            callback(True)

    @dbmethod
    def queue_get_active_tasks(self, gridspec, callback=None):
        """Get a dict of active tasks (queued,processing,reset,resume) on this site and plugin,
           returning {status:{tasks}} where each task = join of search and task tables"""
        if not callback:
            return
        cb = partial(self._queue_get_active_tasks_blocking,gridspec,callback=callback)
        self.db.blocking_task('queue',cb)
    def _queue_get_active_tasks_blocking(self, gridspec, callback=None):
        conn,archive_conn = self.db._dbsetup()
        try:
            sql = 'select task_id from search '
            sql += 'where gridspec like ? '
            sql += ' and task_status in ("queued","processing","reset","resume")'
            bindings = ('%'+gridspec+'%',)
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
            tasks = set(row[0] for row in ret)

            sql = 'select * from task where task_id in (%s)'
            task_groups = {}
            tasks = self._queue_get_task_from_ret(self._bulk_select(conn,sql,tasks))
            for task_id in tasks:
                status = tasks[task_id]['status']
                if status not in task_groups:
                    task_groups[status] = {}
                task_groups[status][task_id] = tasks[task_id]
            callback(task_groups)
        except Exception as e:
            logger.info('error getting active tasks', exc_info=True)
            callback(e)

    @dbmethod
    def queue_get_grid_tasks(self,gridspec,callback=None):
        """Get a list of tasks (queued, processing) on this
           site and plugin, returning [tasks] where each task has
           task_id, grid_queue_id, submit_time, and submit_dir"""
        if not callback:
            return
        cb = partial(self._queue_get_grid_tasks_blocking,gridspec,callback=callback)
        self.db.blocking_task('queue',cb)
    def _queue_get_grid_tasks_blocking(self, gridspec, callback=None):
        conn,archive_conn = self.db._dbsetup()
        try:
            sql = 'select task_id from search '
            sql += 'where gridspec like ? '
            sql += ' and task_status in ("queued","processing")'
            bindings = ('%'+gridspec+'%',)
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
            tasks = set(row[0] for row in ret)

            sql = 'select * from task where task_id in (%s)'
            task_ret = []
            tasks = self._queue_get_task_from_ret(self._bulk_select(conn,sql,tasks))
            for task_id in tasks:
                task_ret.append({
                    'task_id': task_id,
                    'grid_queue_id': tasks[task_id]['grid_queue_id'],
                    'submit_time': tasks[task_id]['status_changed'],
                    'submit_dir': tasks[task_id]['submit_dir'],
                })
            callback(task_ret)
        except Exception as e:
            logger.info('error getting grid tasks', exc_info=True)
            callback(e)

    @dbmethod
    def queue_set_task_status(self,task,status,callback=None):
        """Set the status of a task"""
        if not isinstance(task,Iterable):
            raise Exception('task is not a str or iterable')
        cb = partial(self._queue_set_task_status_blocking,task,status,
                     callback=callback)
        self.db.non_blocking_task(cb)
    def _queue_set_task_status_blocking(self,task,status,callback=None):
        conn,archive_conn = self.db._dbsetup()
        now = nowstr()
        if isinstance(task,String):
            task = [task]
        elif not isinstance(task,Iterable):
            callback(Exception('unknown type for task'))
            return
        
        msql = 'update search set task_status = ? '
        msql += ' where task_id = ?'
        msql2 = 'update task set prev_status = status, '
        msql2 += ' status = ?, status_changed = ? where task_id = ?'

        # process in batches of 900
        try:
            while task:
                t = task[:900]
                task = task[900:]
{'fraction_non_alphanumeric': 0.10089385474860335, 'fraction_numerical': 0.0038547486033519555, 'mean_word_length': 2.0166835187057632, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3497411', 'n_tokens_mistral': 5144, 'n_tokens_neox': 4993, 'n_words': 1134}