fix bug with workflow scheduling not getting the db time (#36)

Author: aclowes

yawn/utilities/database.py

@@ -26,3 +26,11 @@ def wrapper(*args, **kwargs):
                 'was used inside an transaction.atomic() block.'
 
     return wrapper
+
+
+def current_time():
+    """Return the database time"""
+    with db.connection.cursor() as cursor:
+        cursor.execute("SELECT STATEMENT_TIMESTAMP()")
+        row = cursor.fetchone()
+    return row[0]

yawn/utilities/tests/test_database.py

@@ -1,6 +1,9 @@
+import datetime
+
 import pytest
 from django.db import connection
 
+from yawn.utilities import database
 from yawn.utilities.database import close_on_exception
 
 
@@ -24,3 +27,7 @@ def example_disconnect():
     # but the exception is caught, and on retry the database reconnects:
     with connection.cursor() as cursor:
         cursor.execute('select 1')
+
+
+def test_current_time():
+    assert isinstance(database.current_time(), datetime.datetime)

yawn/worker/tests/test_main.py

@@ -1,5 +1,11 @@
+import datetime
+from unittest import mock
+
+from django.utils import timezone
+
 from yawn.worker.models import Queue
-from yawn.worker.main import Main
+from yawn.worker.main import Main, State
+from yawn.workflow.models import WorkflowName, Workflow
 
 
 def test_init_default_queue():
@@ -21,3 +27,27 @@ def test_init_custom_queues():
     queue_names = Queue.objects.filter(
         id__in=worker1.queue_ids).values_list('name', flat=True)
     assert sorted(queue_names) == queues
+
+
+@mock.patch('yawn.worker.main.time.sleep')
+def test_run(mock_time):
+    # so the worker exits immediately
+    worker = Main(1, 'test name', ['default'])
+
+    def set_shutdown(_):
+        # stop the worker after one run
+        worker.state = State.shutdown
+
+    mock_time.side_effect = set_shutdown
+    worker.run()
+
+
+def test_schedule_workflows():
+    name = WorkflowName.objects.create(name='workflow1')
+    next_run = timezone.now() - datetime.timedelta(hours=1)
+    workflow = Workflow.objects.create(
+        name=name, version=1, schedule_active=True, schedule='0 0 *', next_run=next_run)
+    worker = Main(1, 'test name', ['default'])
+    worker.schedule_workflows()
+    workflow.refresh_from_db()
+    assert workflow.next_run > next_run

yawn/workflow/models.py

@@ -1,9 +1,8 @@
 from django.db import models
 from django.contrib.postgres import fields
 from django.db.models import functions
-from django.utils import timezone
 
-from yawn.utilities import cron
+from yawn.utilities import cron, database
 from yawn.utilities.cron import Crontab
 
 
@@ -41,8 +40,7 @@ class Meta:
     def save(self, **kwargs):
         if self.schedule_active:
             if not self.next_run:
-                # this call uses the server time instead of the db time...
-                self.next_run = Crontab(self.schedule).next_run(timezone.now())
+                self.next_run = Crontab(self.schedule).next_run(database.current_time())
         else:
             self.next_run = None
         super().save(**kwargs)
@@ -78,6 +76,9 @@ def submit_run(self, parameters=None, scheduled_time=None):
             )
             if not template.upstream.exists():
                 task.enqueue()
+
+        # refresh to get the actual DB submitted time
+        run.refresh_from_db()
         return run
 
     def __str__(self):