Merge branch '19744-acr-crunchstat' refs #19744

[arvados.git] / tools / crunchstat-summary / crunchstat_summary / reader.py

diff --git a/tools/crunchstat-summary/crunchstat_summary/reader.py b/tools/crunchstat-summary/crunchstat_summary/reader.py
index 8ccdbc2fcf04e45ca3ab3ec6e2270933d050ea1c..0198d765c3533df4cdeb42096fedc0cd57d20051 100644 (file)
--- a/tools/crunchstat-summary/crunchstat_summary/reader.py
+++ b/tools/crunchstat-summary/crunchstat_summary/reader.py
@@ -4,6 +4,7 @@
 
 import arvados
 import itertools
+import json
 import queue
 import threading
 
@@ -11,24 +12,26 @@ from crunchstat_summary import logger
 
 
 class CollectionReader(object):
-    def __init__(self, collection_id):
+    def __init__(self, collection_id, api_client=None, collection_object=None):
         self._collection_id = collection_id
         self._label = collection_id
         self._readers = []
+        self._api_client = api_client
+        self._collection = collection_object or arvados.collection.CollectionReader(self._collection_id, api_client=self._api_client)
 
     def __str__(self):
         return self._label
 
     def __iter__(self):
         logger.debug('load collection %s', self._collection_id)
-        collection = arvados.collection.CollectionReader(self._collection_id)
-        filenames = [filename for filename in collection]
+
+        filenames = [filename for filename in self._collection]
         # Crunch2 has multiple stats files
         if len(filenames) > 1:
             filenames = ['crunchstat.txt', 'arv-mount.txt']
         for filename in filenames:
             try:
-                self._readers.append(collection.open(filename))
+                self._readers.append(self._collection.open(filename, "rt"))
             except IOError:
                 logger.warn('Unable to open %s', filename)
         self._label = "{}/{}".format(self._collection_id, filenames[0])
@@ -43,6 +46,14 @@ class CollectionReader(object):
                 reader.close()
             self._readers = []
 
+    def node_info(self):
+        try:
+            with self._collection.open("node.json", "rt") as f:
+                return json.load(f)
+        except IOError:
+            logger.warn('Unable to open node.json')
+        return {}
+
 
 class LiveLogReader(object):
     EOF = None
@@ -63,7 +74,7 @@ class LiveLogReader(object):
             ['event_type', 'in', self.event_types]]
         try:
             while True:
-                page = arvados.api().logs().index(
+                page = arvados.api().logs().list(
                     limit=1000,
                     order=['id asc'],
                     filters=filters + [['id','>',str(last_id)]],
@@ -105,3 +116,25 @@ class LiveLogReader(object):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         pass
+
+    def node_info(self):
+        return {}
+
+class StubReader(object):
+    def __init__(self, fh):
+        self.fh = fh
+
+    def __str__(self):
+        return ""
+
+    def __iter__(self):
+        return iter(self.fh)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def node_info(self):
+        return {}