MLRun Errors + v3io access forbidden uses errors (#393)

mlrun/api/api/endpoints/files.py

@@ -4,6 +4,7 @@
 
 from mlrun.api.api.utils import log_and_raise, get_obj_path, get_secrets
 from mlrun.datastore import get_object_stat, store_manager
+from mlrun.errors import MLRunDataStoreError
 
 router = APIRouter()
 
@@ -38,8 +39,11 @@ def get_files(
             }
 
         body = obj.get(size, offset)
-    except FileNotFoundError as e:
-        log_and_raise(status.HTTP_404_NOT_FOUND, path=objpath, err=str(e))
+    except FileNotFoundError as exc:
+        log_and_raise(status.HTTP_404_NOT_FOUND, path=objpath, err=str(exc))
+    except MLRunDataStoreError as exc:
+        log_and_raise(exc.response.status_code, path=objpath, err=str(exc))
+
     if body is None:
         log_and_raise(status.HTTP_404_NOT_FOUND, path=objpath)
 
@@ -65,8 +69,10 @@ def get_filestat(request: Request, schema: str = "", path: str = "", user: str =
     stat = None
     try:
         stat = get_object_stat(path, secrets)
-    except FileNotFoundError as e:
-        log_and_raise(status.HTTP_404_NOT_FOUND, path=path, err=str(e))
+    except FileNotFoundError as exc:
+        log_and_raise(status.HTTP_404_NOT_FOUND, path=path, err=str(exc))
+    except MLRunDataStoreError as exc:
+        log_and_raise(exc.response.status_code, path=path, err=str(exc))
 
     ctype, _ = mimetypes.guess_type(path)
     if not ctype:

mlrun/datastore/base.py

@@ -20,6 +20,7 @@
 import urllib3
 import pandas as pd
 
+import mlrun.errors
 from mlrun.utils import logger
 
 verify_ssl = False
@@ -265,34 +266,35 @@ def basic_auth_header(user, password):
 
 def http_get(url, headers=None, auth=None):
     try:
-        resp = requests.get(url, headers=headers, auth=auth, verify=verify_ssl)
+        response = requests.get(url, headers=headers, auth=auth, verify=verify_ssl)
     except OSError as e:
         raise OSError('error: cannot connect to {}: {}'.format(url, e))
 
-    if not resp.ok:
-        raise OSError('failed to read file in {}'.format(url))
-    return resp.content
+    mlrun.errors.raise_for_status(response)
+
+    return response.content
 
 
 def http_head(url, headers=None, auth=None):
     try:
-        resp = requests.head(url, headers=headers, auth=auth, verify=verify_ssl)
+        response = requests.head(url, headers=headers, auth=auth, verify=verify_ssl)
     except OSError as e:
         raise OSError('error: cannot connect to {}: {}'.format(url, e))
-    if not resp.ok:
-        raise OSError('failed to read file head in {}'.format(url))
-    return resp.headers
+
+    mlrun.errors.raise_for_status(response)
+
+    return response.headers
 
 
 def http_put(url, data, headers=None, auth=None):
     try:
-        resp = requests.put(
+        response = requests.put(
             url, data=data, headers=headers, auth=auth, verify=verify_ssl
         )
     except OSError as e:
         raise OSError('error: cannot connect to {}: {}'.format(url, e))
-    if not resp.ok:
-        raise OSError('failed to upload to {} {}'.format(url, resp.status_code))
+
+    mlrun.errors.raise_for_status(response)
 
 
 def http_upload(url, file_path, headers=None, auth=None):

mlrun/datastore/v3io.py

@@ -18,6 +18,7 @@
 import time
 import v3io.dataplane
 
+import mlrun.errors
 from ..platforms.iguazio import split_path
 from .base import (
     DataStore,
@@ -101,12 +102,19 @@ def listdir(self, key):
         # without the trailing slash
         subpath_length = len(subpath) - 1
 
-        response = v3io_client.get_container_contents(
-            container=container,
-            path=subpath,
-            get_all_attributes=False,
-            directories_only=False,
-        )
+        try:
+            response = v3io_client.get_container_contents(
+                container=container,
+                path=subpath,
+                get_all_attributes=False,
+                directories_only=False,
+            )
+        except RuntimeError as exc:
+            if 'Permission denied' in str(exc):
+                raise mlrun.errors.AccessDeniedError(
+                    f'Access denied to path: {key}'
+                ) from exc
+            raise
 
         # todo: full = key, size, last_modified
         return [obj.key[subpath_length:] for obj in response.output.contents]

mlrun/errors.py

@@ -0,0 +1,74 @@
+import requests
+from fastapi import status
+
+
+class MLRunBaseError(Exception):
+    """
+    A base class from which all other exceptions inherit.
+    If you want to catch all errors that the MLRun SDK might raise,
+    catch this base exception.
+    """
+
+    pass
+
+
+class MLRunHTTPError(MLRunBaseError, requests.HTTPError):
+    def __init__(
+        self, message: str, response: requests.Response = None, status_code: int = None
+    ):
+
+        # because response object is probably with an error, it returns False, so we
+        # should use 'is None' specifically
+        if response is None:
+            response = requests.Response()
+        if status_code:
+            response.status_code = status_code
+
+        requests.HTTPError.__init__(self, message, response=response)
+
+
+class MLRunDataStoreError(MLRunHTTPError):
+    error_status_code = None
+
+    def __init__(self, message: str, response: requests.Response = None):
+        super(MLRunDataStoreError, self).__init__(
+            message, response=response, status_code=self.error_status_code
+        )
+
+
+def raise_for_status(response: requests.Response):
+    """
+    Raise a specific MLRunSDK error depending on the given response status code.
+    If no specific error exists, raises an MLRunHTTPError
+    """
+    try:
+        response.raise_for_status()
+    except requests.HTTPError as exc:
+        try:
+            raise STATUS_ERRORS[response.status_code](
+                str(exc), response=response
+            ) from exc
+        except KeyError:
+            raise MLRunHTTPError(str(exc), response=response) from exc
+
+
+# Specific Errors
+
+
+class UnauthorizedError(MLRunDataStoreError):
+    error_status_code = status.HTTP_401_UNAUTHORIZED
+
+
+class AccessDeniedError(MLRunDataStoreError):
+    error_status_code = status.HTTP_403_FORBIDDEN
+
+
+class NotFoundError(MLRunDataStoreError):
+    error_status_code = status.HTTP_404_NOT_FOUND
+
+
+STATUS_ERRORS = {
+    status.HTTP_401_UNAUTHORIZED: UnauthorizedError,
+    status.HTTP_403_FORBIDDEN: AccessDeniedError,
+    status.HTTP_404_NOT_FOUND: NotFoundError,
+}

tests/test_datastores.py

@@ -14,9 +14,16 @@
 from os import listdir
 from tempfile import TemporaryDirectory
 
-from tests.conftest import rundb_path
-import mlrun
 import pandas as pd
+import requests
+from fastapi import status
+from unittest.mock import Mock
+import pytest
+
+import mlrun
+import mlrun.errors
+import v3io.dataplane
+from tests.conftest import rundb_path
 
 mlrun.mlconf.dbpath = rundb_path
 
@@ -94,3 +101,46 @@ def test_parse_url_preserve_case():
     expected_endpoint = 'Hedi'
     _, endpoint, _ = mlrun.datastore.datastore.parse_url(url)
     assert expected_endpoint, endpoint
+
+
+def test_forbidden_file_access(monkeypatch):
+    class MockV3ioClient:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def get_container_contents(self, *args, **kwargs):
+            raise RuntimeError('Permission denied')
+
+    def mock_get(*args, **kwargs):
+        mock_forbidden_response = Mock()
+        mock_forbidden_response.status_code = status.HTTP_403_FORBIDDEN
+        mock_forbidden_response.raise_for_status = Mock(
+            side_effect=requests.HTTPError('Error', response=mock_forbidden_response)
+        )
+        return mock_forbidden_response
+
+    monkeypatch.setattr(requests, "get", mock_get)
+    monkeypatch.setattr(requests, "head", mock_get)
+    monkeypatch.setattr(v3io.dataplane, "Client", MockV3ioClient)
+
+    store = mlrun.datastore.datastore.StoreManager(
+        secrets={'V3IO_ACCESS_KEY': 'some-access-key'}
+    )
+
+    with pytest.raises(mlrun.errors.AccessDeniedError) as access_denied_exc:
+        obj = store.object('v3io://some-system/some-dir/')
+        obj.listdir()
+
+    assert access_denied_exc.value.response.status_code == status.HTTP_403_FORBIDDEN
+
+    with pytest.raises(mlrun.errors.AccessDeniedError) as access_denied_exc:
+        obj = store.object('v3io://some-system/some-dir/some-file')
+        obj.get()
+
+    assert access_denied_exc.value.response.status_code == status.HTTP_403_FORBIDDEN
+
+    with pytest.raises(mlrun.errors.AccessDeniedError) as access_denied_exc:
+        obj = store.object('v3io://some-system/some-dir/some-file')
+        obj.stat()
+
+    assert access_denied_exc.value.response.status_code == status.HTTP_403_FORBIDDEN