Repository: jcrist/skein (python)

Issue:
`Client.get_applications` should support more filters
Currently `Client.get_applications` supports filtering on application state. We should also support filtering on:

- name
- user
- queue
- start_time (optionally as a range)
- finish_time (optionally as a range)

Pull request:
Add filter methods to `Client.get_applications`
Allows filtering applications on:

- state
- name
- user
- queue
- start and finish times

These filters are also exposed through the CLI.

Fixes #133.

Patch:
diff --git a/java/src/main/java/com/anaconda/skein/Driver.java b/java/src/main/java/com/anaconda/skein/Driver.java
index 566e5a9..8195a0e 100644
--- a/java/src/main/java/com/anaconda/skein/Driver.java
+++ b/java/src/main/java/com/anaconda/skein/Driver.java
@@ -806,7 +806,48 @@ public class Driver {
         return;
       }
 
-      resp.onNext(MsgUtils.writeApplicationsResponse(reports));
+      // XXX: Much like the Queue API, YARN actually provides a method to have
+      // the server do this filtering in the query, but then doesn't expose it
+      // in `YarnClient`. Newer versions of YARN support querying all but the
+      // ranges in the YarnClient API, we should use these when available.
+      String name = Strings.emptyToNull(req.getName());
+      String user = Strings.emptyToNull(req.getUser());
+      String queue = Strings.emptyToNull(req.getQueue());
+      long startedBegin = req.getStartedBegin();
+      long startedEnd = req.getStartedEnd();
+      if (startedEnd == 0) {
+        startedEnd = Long.MAX_VALUE;
+      }
+      long finishedBegin = req.getFinishedBegin();
+      long finishedEnd = req.getFinishedEnd();
+      if (finishedEnd == 0) {
+        finishedEnd = Long.MAX_VALUE;
+      }
+
+      Msg.ApplicationsResponse.Builder builder = Msg.ApplicationsResponse.newBuilder();
+      for (ApplicationReport report : reports) {
+        if (name != null && !report.getName().equals(name)) {
+          continue;
+        }
+        if (user != null && !report.getUser().equals(user)) {
+          continue;
+        }
+        if (queue != null && !report.getQueue().equals(queue)) {
+          continue;
+        }
+        // Ranges are inclusive on both ends
+        long startTime = report.getStartTime();
+        if (startedBegin > startTime || startedEnd < startTime) {
+          continue;
+        }
+        long finishTime = report.getFinishTime();
+        if (finishedBegin > finishTime || finishedEnd < finishTime) {
+          continue;
+        }
+        builder.addReports(MsgUtils.writeApplicationReport(report));
+      }
+
+      resp.onNext(builder.build());
       resp.onCompleted();
     }
 
diff --git a/java/src/main/java/com/anaconda/skein/MsgUtils.java b/java/src/main/java/com/anaconda/skein/MsgUtils.java
index 641de43..2aaee2c 100644
--- a/java/src/main/java/com/anaconda/skein/MsgUtils.java
+++ b/java/src/main/java/com/anaconda/skein/MsgUtils.java
@@ -236,15 +236,6 @@ public class MsgUtils {
       .build();
   }
 
-  public static Msg.ApplicationsResponse writeApplicationsResponse(
-      List<ApplicationReport> reports) {
-    Msg.ApplicationsResponse.Builder builder = Msg.ApplicationsResponse.newBuilder();
-    for (ApplicationReport report : reports) {
-      builder.addReports(writeApplicationReport(report));
-    }
-    return builder.build();
-  }
-
   public static Msg.NodeReport writeNodeReport(
       NodeReport r) {
     return Msg.NodeReport.newBuilder()
diff --git a/java/src/main/proto/skein.proto b/java/src/main/proto/skein.proto
index d06cf69..3896b08 100644
--- a/java/src/main/proto/skein.proto
+++ b/java/src/main/proto/skein.proto
@@ -276,6 +276,13 @@ message Application {
 
 message ApplicationsRequest {
   repeated ApplicationState.Type states = 1;
+  string name = 2;
+  string user = 3;
+  string queue = 4;
+  int64 started_begin = 5;
+  int64 started_end = 6;
+  int64 finished_begin = 7;
+  int64 finished_end = 8;
 }
 
 
diff --git a/skein/cli.py b/skein/cli.py
index 56b159d..722dc48 100644
--- a/skein/cli.py
+++ b/skein/cli.py
@@ -283,15 +283,39 @@ def application_submit(spec):
 @subcommand(application.subs,
             'ls', 'List applications',
             arg('--all', '-a', action='store_true',
-                help=('Show all applications (default is only active '
-                      'applications)')),
-            arg("--state", "-s", action='append',
-                help=('Filter by application states. May be repeated '
-                      'to select multiple states.')))
-def application_ls(all=False, state=None):
+                help=('Show applications in all states (default is only '
+                      'active applications)')),
+            arg('--state', '-s', action='append',
+                help=('Select applications with this state. May be '
+                      'repeated to filter on multiple states.')),
+            arg('--name', help='Select applications with this name.'),
+            arg('--queue', help='Select applications with this queue.'),
+            arg('--user', help='Select applications with this user.'),
+            arg('--started-begin',
+                help=('Select applications that started after this datetime '
+                      '(inclusive). Accepts several date and time formats '
+                      '(e.g. ``YYYY-M-D H:M:S`` or ``H:M``). See the '
+                      'documentation for more information.')),
+            arg('--started-end',
+                help=('Select applications that started before this datetime '
+                      '(inclusive)')),
+            arg('--finished-begin',
+                help=('Select applications that finished after this datetime '
+                      '(inclusive)')),
+            arg('--finished-end',
+                help=('Select applications that finished before this datetime '
+                      '(inclusive)')))
+def application_ls(all=False, state=None, name=None, queue=None, user=None,
+                   started_begin=None, started_end=None, finished_begin=None,
+                   finished_end=None):
     if all and state is None:
         state = tuple(ApplicationState)
-    apps = get_driver().get_applications(states=state)
+
+    apps = get_driver().get_applications(
+        states=state, name=name, queue=queue, user=user,
+        started_begin=started_begin, started_end=started_end,
+        finished_begin=finished_begin, finished_end=finished_end
+    )
     _print_application_status(apps)
 
 
diff --git a/skein/core.py b/skein/core.py
index e9b1382..d365f45 100644
--- a/skein/core.py
+++ b/skein/core.py
@@ -1,5 +1,6 @@
 from __future__ import print_function, division, absolute_import
 
+import datetime
 import errno
 import json
 import os
@@ -25,7 +26,8 @@ from .model import (Security, ApplicationSpec, ApplicationReport,
                     ApplicationState, ContainerState, Container,
                     FinalStatus, Resources, container_instance_from_string,
                     LogLevel, NodeState, NodeReport, Queue)
-from .utils import cached_property, grpc_fork_support_disabled, pid_exists
+from .utils import (cached_property, grpc_fork_support_disabled, pid_exists,
+                    datetime_to_millis)
 
 
 __all__ = ('Client', 'ApplicationClient', 'properties')
@@ -574,7 +576,34 @@ class Client(_ClientBase):
                                  app_id,
                                  security=security)
 
-    def get_applications(self, states=None):
+    @staticmethod
+    def _parse_datetime(x, parameter_name):
+        if x is None:
+            return None
+        if isinstance(x, datetime.datetime):
+            return x
+        if not isinstance(x, str):
+            raise context.TypeError("%s must be a datetime or str, got %r"
+                                    % (parameter_name, type(x).__name__))
+        for format in ['%Y-%m-%d', '%Y-%m-%d %H:%M', '%Y-%m-%d %H:%M:%S']:
+            try:
+                return datetime.datetime.strptime(x, format)
+            except ValueError:
+                continue
+        # Times without dates refer to today
+        for format in ['%H:%M', '%H:%M:%S']:
+            try:
+                t = datetime.datetime.strptime(x, format).time()
+                return datetime.datetime.combine(datetime.date.today(), t)
+            except ValueError:
+                continue
+
+        raise context.ValueError("Unable to parse %r from %r as datetime"
+                                 % (x, parameter_name))
+
+    def get_applications(self, states=None, name=None, user=None, queue=None,
+                         started_begin=None, started_end=None,
+                         finished_begin=None, finished_end=None):
         """Get the status of current skein applications.
 
         Parameters
@@ -582,6 +611,32 @@ class Client(_ClientBase):
         states : sequence of ApplicationState, optional
             If provided, applications will be filtered to these application
             states. Default is ``['SUBMITTED', 'ACCEPTED', 'RUNNING']``.
+        name : str, optional
+            Only select applications with this name.
+        user : str, optional
+            Only select applications with this user.
+        queue : str, optional
+            Only select applications in this queue.
+        started_begin : datetime or str, optional
+            Only select applications that started after this time (inclusive).
+            Can be either a datetime or a string representation of one. String
+            representations can use any of the following formats:
+
+              - ``YYYY-M-D H:M:S`` (e.g. 2019-4-10 14:50:20)
+              - ``YYYY-M-D H:M`` (e.g. 2019-4-10 14:50)
+              - ``YYYY-M-D`` (e.g. 2019-4-10)
+              - ``H:M:S`` (e.g. 14:50:20, today is used for date)
+              - ``H:M`` (e.g. 14:50, today is used for date)
+
+        started_end : datetime or str, optional
+            Only select applications that started before this time (inclusive).
+            Can be either a datetime or a string representation of one.
+        finished_begin : datetime or str, optional
+            Only select applications that finished after this time (inclusive).
+            Can be either a datetime or a string representation of one.
+        finished_end : datetime or str, optional
+            Only select applications that finished before this time (inclusive).
+            Can be either a datetime or a string representation of one.
 
         Returns
         -------
@@ -595,6 +650,12 @@ class Client(_ClientBase):
         [ApplicationReport<name='demo'>,
          ApplicationReport<name='dask'>,
          ApplicationReport<name='demo'>]
+
+        Get all applications named 'demo' started after 2019-4-10:
+
+        >>> client.get_applications(name='demo', started_begin='2019-4-10')
+        [ApplicationReport<name='demo'>,
+         ApplicationReport<name='demo'>]
         """
         if states is not None:
             states = tuple(ApplicationState(s) for s in states)
@@ -603,7 +664,21 @@ class Client(_ClientBase):
                       ApplicationState.ACCEPTED,
                       ApplicationState.RUNNING)
 
-        req = proto.ApplicationsRequest(states=[str(s) for s in states])
+        started_begin = self._parse_datetime(started_begin, 'started_begin')
+        started_end = self._parse_datetime(started_end, 'started_end')
+        finished_begin = self._parse_datetime(finished_begin, 'finished_begin')
+        finished_end = self._parse_datetime(finished_end, 'finished_end')
+
+        req = proto.ApplicationsRequest(
+            states=[str(s) for s in states],
+            name=name,
+            user=user,
+            queue=queue,
+            started_begin=datetime_to_millis(started_begin),
+            started_end=datetime_to_millis(started_end),
+            finished_begin=datetime_to_millis(finished_begin),
+            finished_end=datetime_to_millis(finished_end)
+        )
         resp = self._call('getApplications', req)
         return sorted((ApplicationReport.from_protobuf(r) for r in resp.reports),
                       key=lambda x: x.id)
diff --git a/skein/utils.py b/skein/utils.py
index 47f5850..aada957 100644
--- a/skein/utils.py
+++ b/skein/utils.py
@@ -4,6 +4,7 @@ import errno
 import fcntl
 import os
 import threading
+import time
 import weakref
 from contextlib import contextmanager
 from datetime import datetime, timedelta
@@ -132,20 +133,35 @@ _EPOCH = datetime(1970, 1, 1, tzinfo=UTC)
 
 
 def datetime_from_millis(x):
+    """Convert milliseconds since the epoch to a naive `datetime.datetime`"""
     if x is None or x == 0:
         return None
-    return _EPOCH + timedelta(milliseconds=x)
+    return datetime.fromtimestamp(x / 1e3)
 
 
 def datetime_to_millis(x):
-    return int((x - _EPOCH).total_seconds() * 1000)
+    """Convert a `datetime.datetime` to milliseconds since the epoch"""
+    if x is None:
+        return None
+    if hasattr(x, 'timestamp'):
+        # Python >= 3.3
+        secs = x.timestamp()
+    elif x.tzinfo is None:
+        # Timezone naive
+        secs = (time.mktime((x.year, x.month, x.day,
+                             x.hour, x.minute, x.second,
+                             -1, -1, -1)) + x.microsecond / 1e6)
+    else:
+        # Timezone aware
+        secs = (x - _EPOCH).total_seconds()
+    return int(secs * 1000)
 
 
 def runtime(start_time, finish_time):
     if start_time is None:
         return timedelta(0)
     if finish_time is None:
-        return datetime.now(UTC) - start_time
+        return datetime.now() - start_time
     return finish_time - start_time
 
 


Test patch:
diff --git a/skein/test/test_core.py b/skein/test/test_core.py
index 2c73dec..c1734cb 100644
--- a/skein/test/test_core.py
+++ b/skein/test/test_core.py
@@ -1,5 +1,6 @@
 from __future__ import print_function, division, absolute_import
 
+import datetime
 import os
 import subprocess
 import time
@@ -307,6 +308,95 @@ def test_get_queue_methods(client):
         client.get_child_queues("missing")
 
 
+@pytest.fixture
+def at_least_3_apps_in_history(client):
+    all_apps = client.get_applications(states=skein.model.ApplicationState.values())
+    all_expected = [a for a in all_apps if a.user == 'testuser' and a.queue == 'default']
+    for _ in range(3 - len(all_expected)):
+        with run_application(client) as app:
+            app.shutdown()
+
+
+def test_get_applications(client, at_least_3_apps_in_history):
+    all_states = skein.model.ApplicationState.values()
+    all_apps = client.get_applications(states=all_states)
+    # This test is a bit tricky, since the many apps may have been run
+    # previously, some not in our test suite. Here we filter to just ones we'd
+    # expect.
+    all_expected = [a for a in all_apps if a.user == 'testuser' and a.queue == 'default']
+    assert all_expected
+
+    # Filter on user and queue to show identical
+    common = {'states': all_states, 'user': 'testuser', 'queue': 'default'}
+    res = client.get_applications(**common)
+    assert res == all_expected
+
+    # Filter on name as well
+    app_name = all_expected[0].name
+    sol = [a for a in all_expected if a.name == app_name]
+    res = client.get_applications(name=app_name, **common)
+    assert res == sol
+
+    # Filter out all apps
+    assert not client.get_applications(states=all_states, user='not-a-real-value')
+    assert not client.get_applications(states=all_states, queue='not-a-real-value')
+    assert not client.get_applications(states=all_states, name='not-a-real-value')
+
+    mid_app = all_expected[len(all_expected) // 2]
+    start_time = mid_app.start_time.replace(microsecond=0)
+    finish_time = mid_app.finish_time.replace(microsecond=0)
+
+    # All apps started after start_time
+    sol = [a for a in all_expected if a.start_time >= start_time]
+    res = client.get_applications(started_begin=start_time, **common)
+    assert res == sol
+
+    # All apps started before start_time
+    sol = [a for a in all_expected if a.start_time <= start_time]
+    res = client.get_applications(started_end=start_time, **common)
+    assert res == sol
+
+    # All apps finished after finish_time
+    sol = [a for a in all_expected if a.finish_time >= finish_time]
+    res = client.get_applications(finished_begin=finish_time, **common)
+    assert res == sol
+
+    # All apps finished before finish_time
+    sol = [a for a in all_expected if a.finish_time <= finish_time]
+    res = client.get_applications(finished_end=finish_time, **common)
+    assert res == sol
+
+    # Check time parsing
+    t = datetime.datetime.now()
+    for format in ['%Y-%m-%d', '%Y-%m-%d %H:%M', '%Y-%m-%d %H:%M:%S']:
+        t_str = t.strftime(format)
+        sol = datetime.datetime.strptime(t_str, format)
+        res = client._parse_datetime(t_str, '')
+        assert sol == res
+
+    for format in ['%H:%M', '%H:%M:%S']:
+        d = datetime.datetime(year=t.year, month=t.month, day=t.day)
+        t_str = t.strftime(format)
+        tm = datetime.datetime.strptime(t_str, format).time()
+        sol = datetime.datetime.combine(d, tm)
+        res = client._parse_datetime(t_str, '')
+        assert sol == res
+
+    sol = [a for a in all_expected if a.finish_time <= finish_time]
+    res = client.get_applications(
+        finished_end=finish_time.strftime('%Y-%m-%d %H:%M:%S'), **common
+    )
+    assert res == sol
+
+    with pytest.raises(TypeError) as exc:
+        res = client.get_applications(finished_end=123)
+    assert 'finished_end' in str(exc.value)
+
+    with pytest.raises(ValueError) as exc:
+        res = client.get_applications(finished_end='4-1-2019')
+    assert 'finished_end' in str(exc.value)
+
+
 def test_appclient_and_security_in_container(monkeypatch, tmpdir, security):
     # Not running in a container
     with pytest.raises(ValueError) as exc:
diff --git a/skein/test/test_model.py b/skein/test/test_model.py
index 8bcf5ab..a999218 100644
--- a/skein/test/test_model.py
+++ b/skein/test/test_model.py
@@ -7,7 +7,7 @@ import pickle
 
 import pytest
 
-from skein.compatibility import UTC, math_ceil
+from skein.compatibility import math_ceil
 from skein.model import (ApplicationSpec, Service, Resources, File,
                          ApplicationState, FinalStatus, FileType, ACLs, Master,
                          Container, ApplicationReport, ResourceUsageReport,
@@ -566,10 +566,8 @@ def test_enums():
 
 
 def test_container():
-    start = datetime.datetime(2018, 6, 7, 23, 24, 25, 26 * 1000,
-                              tzinfo=UTC)
-    finish = datetime.datetime(2018, 6, 7, 23, 21, 25, 26 * 1000,
-                               tzinfo=UTC)
+    start = datetime.datetime(2018, 6, 7, 23, 24, 25, 26 * 1000)
+    finish = datetime.datetime(2018, 6, 7, 23, 21, 25, 26 * 1000)
 
     kwargs = dict(service_name="foo",
                   instance=0,
@@ -596,9 +594,9 @@ def test_container():
 
     assert c2.runtime == c2.finish_time - c2.start_time
 
-    before = datetime.datetime.now(UTC)
+    before = datetime.datetime.now()
     runtime = c.runtime
-    after = datetime.datetime.now(UTC)
+    after = datetime.datetime.now()
     assert (before - c.start_time) <= runtime <= (after - c.start_time)
 
     assert c3.runtime == datetime.timedelta(0)
@@ -621,8 +619,8 @@ def test_application_report():
                                 Resources(memory=256, vcores=2),
                                 Resources(memory=384, vcores=3))
 
-    start = datetime.datetime(2018, 6, 7, 23, 24, 25, 26 * 1000, tzinfo=UTC)
-    finish = datetime.datetime(2018, 6, 7, 23, 21, 25, 26 * 1000, tzinfo=UTC)
+    start = datetime.datetime(2018, 6, 7, 23, 24, 25, 26 * 1000)
+    finish = datetime.datetime(2018, 6, 7, 23, 21, 25, 26 * 1000)
 
     kwargs = dict(id='application_1528138529205_0001',
                   name='test',
@@ -652,9 +650,9 @@ def test_application_report():
     check_base_methods(a, b)
 
     assert b.runtime == b.finish_time - b.start_time
-    before = datetime.datetime.now(UTC)
+    before = datetime.datetime.now()
     runtime = a.runtime
-    after = datetime.datetime.now(UTC)
+    after = datetime.datetime.now()
     assert (before - a.start_time) <= runtime <= (after - a.start_time)
 
 
diff --git a/skein/test/test_utils.py b/skein/test/test_utils.py
index 0bcf86a..ae08d48 100644
--- a/skein/test/test_utils.py
+++ b/skein/test/test_utils.py
@@ -1,10 +1,12 @@
 from __future__ import absolute_import, print_function, division
 
 import multiprocessing
-from datetime import timedelta
+from datetime import timedelta, datetime
 
+from skein.compatibility import UTC
 from skein.utils import (humanize_timedelta, format_table,
-                         format_comma_separated_list, lock_file)
+                         format_comma_separated_list, lock_file,
+                         datetime_from_millis, datetime_to_millis)
 
 
 def test_humanize_timedelta():
@@ -13,6 +15,24 @@ def test_humanize_timedelta():
     assert humanize_timedelta(timedelta(0, 6001)) == '1h 40m'
 
 
+def test_datetime_conversion():
+    # Timezone naive
+    now = datetime.now()
+    now2 = datetime_from_millis(datetime_to_millis(now))
+    now3 = datetime_from_millis(datetime_to_millis(now2))
+    # Rounded to nearest millisecond
+    assert abs(now2 - now).total_seconds() < 0.001
+    assert now2 == now3
+
+    # Timezone aware
+    now = datetime.now(UTC)
+    now2 = datetime_from_millis(datetime_to_millis(now)).replace(tzinfo=UTC)
+    now3 = datetime_from_millis(datetime_to_millis(now2)).replace(tzinfo=UTC)
+    # Rounded to nearest millisecond
+    assert abs(now2 - now).total_seconds() < 0.001
+    assert now2 == now3
+
+
 def test_format_table():
     res = format_table(['fruit', 'color'],
                        [('apple', 'red'),