server/image-hash: improve exception handling

5 files changed, 74 insertions, 30 deletions

diff --git a/server/szurubooru/api/post_api.py b/server/szurubooru/api/post_api.py
index 5b38b18..cbf8f27 100644
--- a/server/szurubooru/api/post_api.py
+++ b/server/szurubooru/api/post_api.py
@@ -1,5 +1,5 @@
 import datetime
-from szurubooru import search, db
+from szurubooru import search, db, errors
 from szurubooru.rest import routes
 from szurubooru.func import (
     auth, tags, posts, snapshots, favorites, scores, util, versions)
@@ -211,6 +211,12 @@ def get_posts_around(ctx, params):
 def get_posts_by_image(ctx, _params=None):
     auth.verify_privilege(ctx.user, 'posts:reverse_search')
     content = ctx.get_file('content', required=True)
+
+    try:
+        lookalikes = posts.search_by_image(content)
+    except (errors.ThirdPartyError, errors.ProcessingError):
+        lookalikes = []
+
     return {
         'exactPost':
             _serialize_post(ctx, posts.search_by_image_exact(content)),
@@ -220,6 +226,6 @@ def get_posts_by_image(ctx, _params=None):
                     'distance': lookalike.distance,
                     'post': _serialize_post(ctx, lookalike.post),
                 }
-                for lookalike in posts.search_by_image(content)
+                for lookalike in lookalikes
             ],
     }
diff --git a/server/szurubooru/errors.py b/server/szurubooru/errors.py
index f7edf85..4fbb67b 100644
--- a/server/szurubooru/errors.py
+++ b/server/szurubooru/errors.py
@@ -46,3 +46,7 @@ class MissingRequiredParameterError(ValidationError):
 
 class InvalidParameterError(ValidationError):
     pass
+
+
+class ThirdPartyError(BaseError):
+    pass
diff --git a/server/szurubooru/facade.py b/server/szurubooru/facade.py
index 7bec33a..30e5836 100644
--- a/server/szurubooru/facade.py
+++ b/server/szurubooru/facade.py
@@ -44,6 +44,13 @@ def _on_processing_error(ex):
     raise _map_error(ex, rest.errors.HttpBadRequest, 'Processing error')
 
 
+def _on_third_party_error(ex):
+    raise _map_error(
+        ex,
+        rest.errors.HttpInternalServerError,
+        'Server configuration error')
+
+
 def _on_stale_data_error(_ex):
     raise rest.errors.HttpConflict(
         name='IntegrityError',
@@ -110,6 +117,7 @@ def create_app():
     rest.errors.handle(errors.IntegrityError, _on_integrity_error)
     rest.errors.handle(errors.NotFoundError, _on_not_found_error)
     rest.errors.handle(errors.ProcessingError, _on_processing_error)
+    rest.errors.handle(errors.ThirdPartyError, _on_third_party_error)
     rest.errors.handle(sqlalchemy.orm.exc.StaleDataError, _on_stale_data_error)
 
     return rest.application
diff --git a/server/szurubooru/func/image_hash.py b/server/szurubooru/func/image_hash.py
index ebd96cd..8f8c4d7 100644
--- a/server/szurubooru/func/image_hash.py
+++ b/server/szurubooru/func/image_hash.py
@@ -1,10 +1,13 @@
+import logging
 import elasticsearch
 import elasticsearch_dsl
+import xml.etree
 from image_match.elasticsearch_driver import SignatureES
-from szurubooru import config
+from szurubooru import config, errors
 
 
 # pylint: disable=invalid-name
+logger = logging.getLogger(__name__)
 es = elasticsearch.Elasticsearch([{
     'host': config.config['elasticsearch']['host'],
     'port': config.config['elasticsearch']['port'],
@@ -12,6 +15,28 @@ es = elasticsearch.Elasticsearch([{
 session = SignatureES(es, index='szurubooru')
 
 
+def _safe_blanket(default_param_factory):
+    def wrapper_outer(target_function):
+        def wrapper_inner(*args, **kwargs):
+            try:
+                return target_function(*args, **kwargs)
+            except elasticsearch.exceptions.NotFoundError:
+                # index not yet created, will be created dynamically by
+                # add_image()
+                return default_param_factory()
+            except elasticsearch.exceptions.ElasticsearchException as ex:
+                logger.warning('Problem with elastic search: %s' % ex)
+                raise errors.ThirdPartyError(
+                    'Error connecting to elastic search.')
+            except xml.etree.ElementTree.ParseError as ex:
+                # image-match issue #60
+                raise errors.ProcessingError('Not an image.')
+            except Exception as ex:
+                raise errors.ThirdPartyError('Unknown error (%s).' % ex)
+        return wrapper_inner
+    return wrapper_outer
+
+
 class Lookalike:
     def __init__(self, score, distance, path):
         self.score = score
@@ -19,39 +44,37 @@ class Lookalike:
         self.path = path
 
 
+@_safe_blanket(lambda: None)
 def add_image(path, image_content):
     if not path or not image_content:
         return
     session.add_image(path=path, img=image_content, bytestream=True)
 
 
+@_safe_blanket(lambda: None)
 def delete_image(path):
     if not path:
         return
-    try:
-        es.delete_by_query(
-            index=session.index,
-            doc_type=session.doc_type,
-            body={'query': {'term': {'path': path}}})
-    except elasticsearch.exceptions.NotFoundError:
-        pass
+    es.delete_by_query(
+        index=session.index,
+        doc_type=session.doc_type,
+        body={'query': {'term': {'path': path}}})
 
 
+@_safe_blanket(lambda: [])
 def search_by_image(image_content):
-    try:
-        for result in session.search_image(
-                path=image_content,  # sic
-                bytestream=True):
-            yield Lookalike(
-                score=result['score'],
-                distance=result['dist'],
-                path=result['path'])
-    except elasticsearch.exceptions.ElasticsearchException:
-        raise
-    except Exception:
-        yield from []
+    ret = []
+    for result in session.search_image(
+            path=image_content,  # sic
+            bytestream=True):
+        ret.append(Lookalike(
+            score=result['score'],
+            distance=result['dist'],
+            path=result['path']))
+    return ret
 
 
+@_safe_blanket(lambda: None)
 def purge():
     es.delete_by_query(
         index=session.index,
@@ -59,12 +82,10 @@ def purge():
         body={'query': {'match_all': {}}})
 
 
+@_safe_blanket(lambda: set())
 def get_all_paths():
-    try:
-        search = (
-            elasticsearch_dsl.Search(
-                using=es, index=session.index, doc_type=session.doc_type)
-            .source(['path']))
-        return set(h.path for h in search.scan())
-    except elasticsearch.exceptions.NotFoundError:
-        return set()
+    search = (
+        elasticsearch_dsl.Search(
+            using=es, index=session.index, doc_type=session.doc_type)
+        .source(['path']))
+    return set(h.path for h in search.scan())
diff --git a/server/szurubooru/rest/errors.py b/server/szurubooru/rest/errors.py
index 3c40b4c..b0f5b88 100644
--- a/server/szurubooru/rest/errors.py
+++ b/server/szurubooru/rest/errors.py
@@ -47,5 +47,10 @@ class HttpMethodNotAllowed(BaseHttpError):
     reason = 'Method Not Allowed'
 
 
+class HttpInternalServerError(BaseHttpError):
+    code = 500
+    reason = 'Internal Server Error'
+
+
 def handle(exception_type, handler):
     error_handlers[exception_type] = handler