Skip to content

Commit

Permalink
Add CmrQuery storage class as a convenience around HttpRequest
Browse files Browse the repository at this point in the history
  • Loading branch information
reweeden committed Dec 10, 2024
1 parent be9c0f9 commit 896d6bf
Show file tree
Hide file tree
Showing 5 changed files with 112 additions and 12 deletions.
6 changes: 6 additions & 0 deletions mandible/metadata_mapper/storage/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,19 @@
StorageError,
)

try:
from .cmr_query import CmrQuery
except ImportError:
from .storage import CmrQuery

try:
from .http_request import HttpRequest
except ImportError:
from .storage import HttpRequest


__all__ = (
"CmrQuery",
"Dummy",
"FilteredStorage",
"HttpRequest",
Expand Down
46 changes: 46 additions & 0 deletions mandible/metadata_mapper/storage/cmr_query.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
import urllib.parse
from dataclasses import InitVar, dataclass
from typing import Optional

from mandible.metadata_mapper.context import Context

from .http_request import HttpRequest


@dataclass
class CmrQuery(HttpRequest):
"""A convenience class for setting neccessary CMR parameters"""

url: InitVar[None] = None

base_url: str = ""
path: str = ""
format: str = ""
token: Optional[str] = None

def __post_init__(self, url: str):
if url:
raise ValueError(
"do not set 'url' directly, use 'base_url' and 'path' instead",
)

def _get_override_request_args(self, context: Context) -> dict:
return {
"headers": self._get_headers(),
"url": self._get_url(),
}

def _get_headers(self) -> Optional[dict]:
if self.token is None:
return self.headers

return {
**(self.headers or {}),
"Authorization": self.token,
}

def _get_url(self) -> str:
path = self.path
if self.format:
path = f"{self.path}.{self.format.lower()}"
return urllib.parse.urljoin(self.base_url, path)
32 changes: 20 additions & 12 deletions mandible/metadata_mapper/storage/http_request.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
class HttpRequest(Storage):
"""A storage which returns the body of an HTTP response"""

# TODO(reweeden): Python 3.10 added support for KW_ONLY arguments which can
# be used to clean up the inheritance here a bit.
url: str
method: str = "GET"
params: Optional[dict] = None
Expand All @@ -26,18 +28,21 @@ class HttpRequest(Storage):
allow_redirects: bool = True

def open_file(self, context: Context) -> IO[bytes]:
response = requests.request(
self.method,
self.url,
params=self.params,
data=self.data,
json=self.json,
headers=self.headers,
cookies=self.cookies,
timeout=self.timeout,
allow_redirects=self.allow_redirects,
stream=True,
)
kwargs = {
"allow_redirects": self.allow_redirects,
"cookies": self.cookies,
"data": self.data,
"headers": self.headers,
"json": self.json,
"method": self.method,
"params": self.params,
"stream": True,
"timeout": self.timeout,
"url": self.url,
# Allow subclasses to override these
**self._get_override_request_args(context),
}
response = requests.request(**kwargs)

# TODO(reweeden): Using response.content causes the entire response
# payload to be loaded into memory immediately. Ideally, we would
Expand All @@ -46,3 +51,6 @@ def open_file(self, context: Context) -> IO[bytes]:
# object, however, this doesn't preform the content decoding that you
# get when using response.content.
return io.BytesIO(response.content)

def _get_override_request_args(self, context: Context) -> dict:
return {}
6 changes: 6 additions & 0 deletions mandible/metadata_mapper/storage/storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,12 @@ def __init__(self):
super().__init__("requests")


@dataclass
class CmrQuery(_PlaceholderBase):
def __init__(self):
super().__init__("requests")


# Define storages that don't require extra dependencies

@dataclass
Expand Down
34 changes: 34 additions & 0 deletions tests/test_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from mandible.metadata_mapper.context import Context
from mandible.metadata_mapper.storage import (
STORAGE_REGISTRY,
CmrQuery,
Dummy,
HttpRequest,
LocalFile,
Expand All @@ -17,6 +18,7 @@

def test_registry():
assert STORAGE_REGISTRY == {
"CmrQuery": CmrQuery,
"Dummy": Dummy,
"HttpRequest": HttpRequest,
"LocalFile": LocalFile,
Expand Down Expand Up @@ -180,3 +182,35 @@ def create_file(bucket, name, contents=None, type="data"):
})
with storage.open_file(context) as f:
assert f.read() == b"Content from file2.txt\n"


def test_cmr_query_params():
with pytest.raises(ValueError):
CmrQuery(url="foobar")

assert CmrQuery(
base_url="http://foo.bar",
path="/search/granules",
)._get_url() == "http://foo.bar/search/granules"
assert CmrQuery(
base_url="http://foo.bar",
path="search/granules",
)._get_url() == "http://foo.bar/search/granules"
assert CmrQuery(
base_url="http://foo.bar/",
path="/search/granules",
)._get_url() == "http://foo.bar/search/granules"
assert CmrQuery(
base_url="http://foo.bar/",
path="search/granules",
)._get_url() == "http://foo.bar/search/granules"

assert CmrQuery(
base_url="http://foo.bar",
path="/search/granules",
format="umm_json",
)._get_url() == "http://foo.bar/search/granules.umm_json"

assert CmrQuery(token="foobar")._get_headers() == {
"Authorization": "foobar",
}

0 comments on commit 896d6bf

Please sign in to comment.