Initial import from Docker volume
This commit is contained in:
root
262
ir/queue_.py
Executable file
262
ir/queue_.py
Executable file
@@ -0,0 +1,262 @@
|
||||
# This file is part of Tryton. The COPYRIGHT file at the top level of
|
||||
# this repository contains the full copyright notices and license terms.
|
||||
import datetime
|
||||
|
||||
from sql import Literal, Null, With
|
||||
from sql.aggregate import Min
|
||||
from sql.functions import CurrentTimestamp, Extract
|
||||
|
||||
from trytond.config import config
|
||||
from trytond.model import Index, ModelSQL, fields
|
||||
from trytond.pool import Pool
|
||||
from trytond.tools import grouped_slice
|
||||
from trytond.transaction import (
|
||||
Transaction, inactive_records, without_check_access)
|
||||
|
||||
has_worker = config.getboolean('queue', 'worker', default=False)
|
||||
clean_days = config.getint('queue', 'clean_days', default=30)
|
||||
batch_size = config.getint('queue', 'batch_size', default=20)
|
||||
|
||||
|
||||
class Queue(ModelSQL):
|
||||
"Queue"
|
||||
__name__ = 'ir.queue'
|
||||
name = fields.Char("Name", required=True)
|
||||
|
||||
data = fields.Dict(None, "Data")
|
||||
|
||||
enqueued_at = fields.Timestamp("Enqueued at", required=True)
|
||||
dequeued_at = fields.Timestamp("Dequeued at")
|
||||
finished_at = fields.Timestamp("Finished at")
|
||||
|
||||
scheduled_at = fields.Timestamp("Scheduled at",
|
||||
help="When the task can start.")
|
||||
expected_at = fields.Timestamp("Expected at",
|
||||
help="When the task should be done.")
|
||||
|
||||
@classmethod
|
||||
def __setup__(cls):
|
||||
super().__setup__()
|
||||
table = cls.__table__()
|
||||
|
||||
cls._sql_indexes.add(
|
||||
Index(
|
||||
table,
|
||||
(table.scheduled_at, Index.Range(nulls_first=True)),
|
||||
(table.expected_at, Index.Range(nulls_first=True)),
|
||||
(table.dequeued_at, Index.Equality()),
|
||||
(table.name, Index.Equality())))
|
||||
|
||||
@classmethod
|
||||
def default_enqueued_at(cls):
|
||||
return datetime.datetime.now()
|
||||
|
||||
@classmethod
|
||||
def copy(cls, records, default=None):
|
||||
if default is None:
|
||||
default = {}
|
||||
else:
|
||||
default = default.copy()
|
||||
default.setdefault('enqueued_at')
|
||||
default.setdefault('dequeued_at')
|
||||
default.setdefault('finished_at')
|
||||
return super(Queue, cls).copy(records, default=default)
|
||||
|
||||
@classmethod
|
||||
def push(cls, name, data, scheduled_at=None, expected_at=None):
|
||||
transaction = Transaction()
|
||||
database = transaction.database
|
||||
cursor = transaction.connection.cursor()
|
||||
with without_check_access():
|
||||
record, = cls.create([{
|
||||
'name': name,
|
||||
'data': data,
|
||||
'scheduled_at': scheduled_at,
|
||||
'expected_at': expected_at,
|
||||
}])
|
||||
if database.has_channel():
|
||||
cursor.execute('NOTIFY "%s"', (cls.__name__,))
|
||||
if not has_worker:
|
||||
transaction.tasks.append(record.id)
|
||||
return record.id
|
||||
|
||||
@classmethod
|
||||
def pull(cls, database, connection, name=None):
|
||||
cursor = connection.cursor()
|
||||
queue = cls.__table__()
|
||||
queue_c = cls.__table__()
|
||||
queue_s = cls.__table__()
|
||||
|
||||
candidates = With('id', 'scheduled_at', 'expected_at',
|
||||
query=queue_c.select(
|
||||
queue_c.id,
|
||||
queue_c.scheduled_at,
|
||||
queue_c.expected_at,
|
||||
where=((queue_c.name == name) if name else Literal(True))
|
||||
& (queue_c.dequeued_at == Null),
|
||||
order_by=[
|
||||
queue_c.scheduled_at.nulls_first,
|
||||
queue_c.expected_at.nulls_first]))
|
||||
selected = queue_s.select(
|
||||
queue_s.id,
|
||||
where=((queue_s.name == name) if name else Literal(True))
|
||||
& (queue_s.dequeued_at == Null)
|
||||
& ((queue_s.scheduled_at <= CurrentTimestamp())
|
||||
| (queue_s.scheduled_at == Null)),
|
||||
order_by=[
|
||||
queue_s.scheduled_at.nulls_first,
|
||||
queue_s.expected_at.nulls_first],
|
||||
limit=1)
|
||||
if database.has_select_for():
|
||||
For = database.get_select_for_skip_locked()
|
||||
selected.for_ = For('UPDATE')
|
||||
|
||||
next_timeout = With('seconds', query=candidates.select(
|
||||
Min(Extract('EPOCH',
|
||||
candidates.scheduled_at - CurrentTimestamp())
|
||||
),
|
||||
where=candidates.scheduled_at >= CurrentTimestamp()))
|
||||
|
||||
task_id, seconds = None, None
|
||||
if database.has_returning():
|
||||
query = queue.update([queue.dequeued_at], [CurrentTimestamp()],
|
||||
where=queue.id.in_(selected),
|
||||
with_=[candidates, next_timeout],
|
||||
returning=[
|
||||
queue.id, next_timeout.select(next_timeout.seconds)])
|
||||
cursor.execute(*query)
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
task_id, seconds = row
|
||||
else:
|
||||
query = queue.select(queue.id,
|
||||
where=queue.id.in_(selected),
|
||||
with_=[candidates])
|
||||
cursor.execute(*query)
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
task_id, = row
|
||||
query = queue.update([queue.dequeued_at], [CurrentTimestamp()],
|
||||
where=queue.id == task_id)
|
||||
cursor.execute(*query)
|
||||
query = next_timeout.select(next_timeout.seconds)
|
||||
cursor.execute(*query)
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
seconds, = row
|
||||
|
||||
if not task_id and database.has_channel():
|
||||
cursor.execute('LISTEN "%s"', (cls.__name__,))
|
||||
return task_id, seconds
|
||||
|
||||
def run(self):
|
||||
transaction = Transaction()
|
||||
Model = Pool().get(self.data['model'])
|
||||
with transaction.set_user(self.data['user']), \
|
||||
transaction.set_context(
|
||||
self.data['context'], _skip_warnings=True):
|
||||
instances = self.data['instances']
|
||||
# Ensure record ids still exist
|
||||
if isinstance(instances, int):
|
||||
with inactive_records():
|
||||
if Model.search([('id', '=', instances)]):
|
||||
instances = Model(instances)
|
||||
else:
|
||||
instances = None
|
||||
else:
|
||||
ids = set()
|
||||
with inactive_records():
|
||||
for sub_ids in grouped_slice(instances):
|
||||
records = Model.search([('id', 'in', list(sub_ids))])
|
||||
ids.update(map(int, records))
|
||||
if ids:
|
||||
instances = Model.browse(
|
||||
[i for i in instances if i in ids])
|
||||
else:
|
||||
instances = None
|
||||
if instances is not None:
|
||||
getattr(Model, self.data['method'])(
|
||||
instances, *self.data['args'], **self.data['kwargs'])
|
||||
if not self.dequeued_at:
|
||||
self.dequeued_at = datetime.datetime.now()
|
||||
self.finished_at = datetime.datetime.now()
|
||||
self.save()
|
||||
|
||||
@classmethod
|
||||
def clean(cls, date=None):
|
||||
if date is None:
|
||||
date = (
|
||||
datetime.datetime.now() - datetime.timedelta(days=clean_days))
|
||||
tasks = cls.search(['OR',
|
||||
('dequeued_at', '<', date),
|
||||
('finished_at', '<', date),
|
||||
])
|
||||
cls.delete(tasks)
|
||||
|
||||
@classmethod
|
||||
def caller(cls, model):
|
||||
return _Model(cls, model)
|
||||
|
||||
|
||||
class _Model(object):
|
||||
def __init__(self, queue, model):
|
||||
self.__queue = queue
|
||||
self.__model = model
|
||||
|
||||
def __getattr__(self, name):
|
||||
return _Method(self.__queue, self.__model, name)
|
||||
|
||||
|
||||
class _Method(object):
|
||||
def __init__(self, queue, model, name):
|
||||
self.__queue = queue
|
||||
self.__model = model
|
||||
self.__name = name
|
||||
|
||||
def __call__(self, instances, *args, **kwargs):
|
||||
transaction = Transaction()
|
||||
context = transaction.context.copy()
|
||||
name = context.pop('queue_name', 'default')
|
||||
now = datetime.datetime.now()
|
||||
scheduled_at = context.pop('queue_scheduled_at', None)
|
||||
if scheduled_at is not None:
|
||||
scheduled_at = now + scheduled_at
|
||||
expected_at = context.pop('queue_expected_at', None)
|
||||
queue_batch = context.pop('queue_batch', None)
|
||||
context.pop('_check_access', None)
|
||||
context.pop('language', None)
|
||||
if expected_at is not None:
|
||||
expected_at = now + expected_at
|
||||
try:
|
||||
instances = list(map(int, instances))
|
||||
except TypeError:
|
||||
instances = int(instances)
|
||||
|
||||
def _push(instances):
|
||||
data = {
|
||||
'model': self.__model.__name__,
|
||||
'method': self.__name,
|
||||
'user': transaction.user,
|
||||
'context': context,
|
||||
'instances': instances,
|
||||
'args': args,
|
||||
'kwargs': kwargs,
|
||||
}
|
||||
return self.__queue.push(
|
||||
name, data,
|
||||
scheduled_at=scheduled_at, expected_at=expected_at)
|
||||
|
||||
if isinstance(instances, list):
|
||||
if has_worker and queue_batch:
|
||||
if isinstance(queue_batch, bool):
|
||||
count = batch_size
|
||||
else:
|
||||
count = int(queue_batch)
|
||||
else:
|
||||
count = len(instances)
|
||||
task_ids = []
|
||||
for sub_instances in grouped_slice(instances, count=count):
|
||||
task_ids.append(_push(list(sub_instances)))
|
||||
return task_ids
|
||||
else:
|
||||
return _push(instances)
|
||||
Reference in New Issue
Block a user