Repository: mlrun/mlrun

Issue:
Persistancy for cron jobs
We'd like to be able to reload schedules after restart

Discussion:

Patch:
diff --git a/mlrun/db/filedb.py b/mlrun/db/filedb.py
index 68df353b9..b5b4d199f 100644
--- a/mlrun/db/filedb.py
+++ b/mlrun/db/filedb.py
@@ -15,7 +15,7 @@
 import json
 import pathlib
 from datetime import datetime, timedelta
-from os import makedirs, path, remove
+from os import makedirs, path, remove, scandir
 
 import yaml
 
@@ -31,6 +31,7 @@ from .base import RunDBError, RunDBInterface
 run_logs = 'runs'
 artifacts_dir = 'artifacts'
 functions_dir = 'functions'
+schedules_dir = 'schedules'
 
 
 class FileRunDB(RunDBInterface):
@@ -41,6 +42,7 @@ class FileRunDB(RunDBInterface):
         self.dirpath = dirpath
         self._datastore = None
         self._subpath = None
+        makedirs(self.schedules_dir, exist_ok=True)
 
     def connect(self, secrets=None):
         sm = StoreManager(secrets)
@@ -257,6 +259,25 @@ class FileRunDB(RunDBInterface):
         return path.join(self.dirpath, '{}/{}/{}{}'.format(
             table, project, tag, key))
 
+    @property
+    def schedules_dir(self):
+        return path.join(self.dirpath, schedules_dir)
+
+    def store_schedule(self, data):
+        sched_id = 1 + sum(1 for _ in scandir(self.schedules_dir))
+        fname = path.join(
+            self.schedules_dir,
+            '{}{}'.format(sched_id, self.format),
+        )
+        with open(fname, 'w') as out:
+            out.write(self._dumps(data))
+
+    def list_schedules(self):
+        pattern = '*{}'.format(self.format)
+        for p in pathlib.Path(self.schedules_dir).glob(pattern):
+            with p.open() as fp:
+                yield self._loads(fp.read())
+
     _encodings = {
         '.yaml': ('to_yaml', dict_to_yaml),
         '.json': ('to_json', dict_to_json),
diff --git a/mlrun/db/httpd.py b/mlrun/db/httpd.py
index aae3bbf58..43c0cfb26 100644
--- a/mlrun/db/httpd.py
+++ b/mlrun/db/httpd.py
@@ -125,6 +125,10 @@ def submit_job():
         return json_error(HTTPStatus.BAD_REQUEST, reason='bad JSON body')
 
     logger.info('submit_job: {}'.format(data))
+    _submit(data)
+
+
+def _submit(data):
     task = data.get('task')
     function = data.get('function')
     url = data.get('functionUrl')
@@ -162,6 +166,7 @@ def submit_job():
         if schedule:
             args = (task, )
             job_id = _scheduler.add(schedule, fn, args)
+            _db.save_schedule(data)
             resp = {'schedule': schedule, 'id': job_id}
         else:
             resp = fn.run(task)
@@ -704,6 +709,11 @@ def init_app():
     periodic.schedule(task, 60)
 
     _scheduler = Scheduler()
+    for data in _db.list_schedules():
+        if 'schedule' not in data:
+            logger.warning('bad scheduler data - %s', data)
+            continue
+        _submit(data)
 
 
 # Don't remove this function, it's an entry point in setup.py
diff --git a/mlrun/db/sqldb.py b/mlrun/db/sqldb.py
index b87f24cf5..f364742ab 100644
--- a/mlrun/db/sqldb.py
+++ b/mlrun/db/sqldb.py
@@ -116,6 +116,12 @@ with warnings.catch_warnings():
         start_time = Column(TIMESTAMP)
         labels = relationship(Label)
 
+    class Schedule(Base, HasStruct):
+        __tablename__ = 'schedules'
+
+        id = Column(Integer, primary_key=True)
+        body = Column(BLOB)
+
 
 class SQLDB(RunDBInterface):
     def __init__(self, dsn):
@@ -328,6 +334,14 @@ class SQLDB(RunDBInterface):
             Artifact.project == project).distinct()
         return [row[0] for row in query]
 
+    def store_schedule(self, data):
+        sched = Schedule()
+        sched.struct = data
+        self._upsert(sched)
+
+    def list_schedules(self):
+        return [s.struct for s in self.session.query(Schedule)]
+
     def _query(self, cls, **kw):
         kw = {k: v for k, v in kw.items() if v}
         return self.session.query(cls).filter_by(**kw)
diff --git a/mlrun/funcdoc.py b/mlrun/funcdoc.py
index 513426a9c..3d54e0c22 100644
--- a/mlrun/funcdoc.py
+++ b/mlrun/funcdoc.py
@@ -107,14 +107,16 @@ def rst_read_doc(lines):
 
 
 def rst_read_section(lines, i):
-    # Skip empty lines
-    while not lines[i].strip() and i < len(lines):
-        i += 1
-
-    # :param path: The path of the file to wrap
-    match = re.match(r':(\w+)(\s+\w+)?:', lines[i])
-    if not match:
-        raise ValueError(f'{i}: bad line - {lines[i]!r}')
+    # Skip empty lines/other lines
+    for i, line in enumerate(lines[i:], i):
+        if not line.strip():
+            continue
+        # :param path: The path of the file to wrap
+        match = re.match(r':\s*(\w+)(\s+\w+)?\s*:', lines[i])
+        if match:
+            break
+    else:
+        return None
 
     tag = match.group(1)
     value = match.group(2).strip() if match.group(2) else ''
@@ -134,7 +136,10 @@ def parse_rst(docstring: str):
     ret = param_dict()
 
     while i != -1:
-        tag, value, text, i = rst_read_section(lines, i)
+        out = rst_read_section(lines, i)
+        if not out:
+            break
+        tag, value, text, i = out
         if tag == 'param':
             params[value] = param_dict(name=value, doc=text)
             names.append(value)


Test patch:
diff --git a/tests/rst_cases.yml b/tests/rst_cases.yml
index 3bbd601ba..6953f4934 100644
--- a/tests/rst_cases.yml
+++ b/tests/rst_cases.yml
@@ -78,3 +78,47 @@
         type: BufferedFileStorage
         default: ""
     lineno: -1
+- name: space in colon
+  text: |
+    Open a file
+
+    :param path : The path of the file to wrap
+  expected:
+    doc: Open a file
+    params:
+      - name: path
+        doc: The path of the file to wrap
+        type: ""
+        default: ""
+    ret:
+        name: ""
+        doc: ""
+        type: ""
+        default: ""
+    lineno: -1
+- name: Data in middle
+  text: |
+    Open a file
+
+    :param path : The path of the file to wrap
+
+    Then we can optionally
+
+    :param mode : File mode
+  expected:
+    doc: Open a file
+    params:
+      - name: path
+        doc: The path of the file to wrap
+        type: ""
+        default: ""
+      - name: mode
+        doc: File mode
+        type: ""
+        default: ""
+    ret:
+        name: ""
+        doc: ""
+        type: ""
+        default: ""
+    lineno: -1
diff --git a/tests/test_filedb.py b/tests/test_filedb.py
index 2cd3feb5e..6f3fb2b58 100644
--- a/tests/test_filedb.py
+++ b/tests/test_filedb.py
@@ -45,3 +45,14 @@ def test_list_fuctions(db: FileRunDB):
 
     out = db.list_functions('', proj)
     assert len(out) == count, 'bad list'
+
+
+def test_schedules(db: FileRunDB):
+    count = 7
+    for i in range(count):
+        data = {'i': i}
+        db.store_schedule(data)
+
+    scheds = list(db.list_schedules())
+    assert count == len(scheds), 'wrong number of schedules'
+    assert set(range(count)) == set(s['i'] for s in scheds), 'bad scheds'
diff --git a/tests/test_sqldb.py b/tests/test_sqldb.py
index ff161f6bb..6639a92c4 100644
--- a/tests/test_sqldb.py
+++ b/tests/test_sqldb.py
@@ -189,3 +189,14 @@ def test_list_runs(db: sqldb.SQLDB):
 
     runs = list(db.list_runs(uid=uid, iter=True))
     assert 5 == len(runs), 'iter=True'
+
+
+def test_schedules(db: sqldb.SQLDB):
+    count = 7
+    for i in range(count):
+        data = {'i': i}
+        db.store_schedule(data)
+
+    scheds = list(db.list_schedules())
+    assert count == len(scheds), 'wrong number of schedules'
+    assert set(range(count)) == set(s['i'] for s in scheds), 'bad scheds'