Skip to content

Commit

Permalink
feat: added load features
Browse files Browse the repository at this point in the history
  • Loading branch information
jankovicgd committed Nov 22, 2024
1 parent a941571 commit 2f85e8c
Show file tree
Hide file tree
Showing 8 changed files with 232 additions and 50 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,4 @@ jobs:
- name: Check with mypy
run: poetry run mypy .
- name: Run tests
run: poetry run pytest -v
run: poetry run pytest -v --cov=src
16 changes: 15 additions & 1 deletion poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 1 addition & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ mkdocs-autorefs = "^1.2.0"
mkdocs-material-extensions = "^1.3.1"
pytest-cov = "^6.0.0"
pytest-mock = "^3.14.0"
respx = "^0.21.1"


[tool.commitizen]
Expand Down Expand Up @@ -76,8 +77,6 @@ exclude = [
module = ["pkg_resources"]
ignore_missing_imports = true

[tool.pytest.ini_options]
addopts = "--cov=src"

[tool.coverage.report]
exclude_also = [
Expand Down
73 changes: 33 additions & 40 deletions src/eodm/cli/load/apps/stac_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,12 @@
import typer

from eodm.cli._globals import DEFAULT_EXTENT
from eodm.cli._serialization import serialize
from eodm.cli._types import Output, OutputType
from eodm.load import load_stac_api_collections, load_stac_api_items

app = typer.Typer(no_args_is_help=True)

HEADERS = {"Content-Type": "application/json"}


@app.callback()
def main():
Expand Down Expand Up @@ -40,13 +41,13 @@ def collection(
)

response = httpx.post(
collections_endpoint, json=collection.to_dict(), headers=HEADERS, verify=verify
collections_endpoint, json=collection.to_dict(), headers={}, verify=verify
)

if update and response.status_code == 409:
collection_endpoint = f"{collections_endpoint}/{collection.id}"
response = httpx.put(
collection_endpoint, json=collection.to_dict(), headers=HEADERS, verify=verify
collection_endpoint, json=collection.to_dict(), headers={}, verify=verify
)

response.raise_for_status()
Expand All @@ -59,30 +60,26 @@ def collections(
verify: bool = True,
update: bool = False,
skip_existing: bool = False,
output: OutputType = Output.default,
) -> None:
"""
Load multiple collections to a stac API
Load multiple collections to a stac API. Collections can be piped from STDIN or a file
with Collection jsons on each line
"""

collections_endpoint = f"{url}/collections"
for line in collections.readlines():
collection_dict = json.loads(line)
response = httpx.post(
collections_endpoint, json=collection_dict, headers=HEADERS, verify=verify
)
if response.status_code == 409:
if update:
collection_endpoint = f"{collections_endpoint}/{collection_dict['id']}"
response = httpx.put(
collection_endpoint,
json=collection_dict,
headers=HEADERS,
verify=verify,
)
if skip_existing:
continue

response.raise_for_status()
_collections = [
pystac.Collection.from_dict(json.loads(line)) for line in collections.readlines()
]
serialize(
load_stac_api_collections(
url=url,
collections=_collections,
verify=verify,
update=update,
skip_existing=skip_existing,
),
output_type=output,
)


@app.command(no_args_is_help=True)
Expand All @@ -92,24 +89,20 @@ def items(
verify: bool = True,
update: bool = False,
skip_existing: bool = False,
output: OutputType = Output.default,
) -> None:
"""
Load multiple items into a STAC API
"""

for line in items.readlines():
item_dict = json.loads(line)
collection_id = item_dict["collection"]
items_endpoint = f"{url}/collections/{collection_id}/items"
response = httpx.post(
items_endpoint, json=item_dict, headers=HEADERS, verify=verify
)
if response.status_code == 409:
if update:
item_endpoint = f"{items_endpoint}/{item_dict['id']}"
response = httpx.put(
item_endpoint, json=item_dict, headers=HEADERS, verify=verify
)
if skip_existing:
continue
response.raise_for_status()
_items = [pystac.Item.from_dict(json.loads(line)) for line in items.readlines()]
serialize(
load_stac_api_items(
url=url,
items=_items,
verify=verify,
update=update,
skip_existing=skip_existing,
),
output_type=output,
)
97 changes: 97 additions & 0 deletions src/eodm/load.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
from typing import Iterable

import httpx
from pystac import Collection, Item

DEFAULT_HEADERS = {"Content-Type": "application/json"}


def load_stac_api_items(
url: str,
items: Iterable[Item],
headers: dict[str, str] | None = None,
verify: bool = True,
update: bool = False,
skip_existing: bool = False,
) -> Iterable[Item]:
"""Load multiple items into a STAC API
Args:
url (str): STAC API url
items (Iterable[Item]): A collection of STAC Items
headers (dict[str, str] | None, optional): Headers to add to the request. Defaults to None.
verify (bool, optional): Verify SSL request. Defaults to True.
update (bool, optional): Update STAC Item with new content. Defaults to False.
skip_existing (bool, optional): Skip Item if exists. Defaults to False.
"""
if not headers:
headers = DEFAULT_HEADERS

for item in items:
collection_id = item.collection_id
items_endpoint = f"{url}/collections/{collection_id}/items"
response = httpx.post(
items_endpoint,
json=item.to_dict(),
headers=headers,
verify=verify,
)
if response.status_code == 409:
if update:
item_endpoint = f"{items_endpoint}/{item.id}"
response = httpx.put(
item_endpoint, json=item.to_dict(), headers=headers, verify=verify
)
if skip_existing:
continue
response.raise_for_status()
yield item


def load_stac_api_collections(
url: str,
collections: Iterable[Collection],
headers: dict[str, str] | None = None,
verify: bool = True,
update: bool = False,
skip_existing: bool = False,
) -> Iterable[Collection]:
"""Load multiple collections to a stac API
Args:
url (str): STAC API URL
collections (Iterable[Collection]): A collection of STAC Collections
headers (dict[str, str] | None, optional): Additional headers to send. Defaults to None.
verify (bool, optional): Verify TLS request. Defaults to True.
update (bool, optional): Update the destination Collections. Defaults to False.
skip_existing (bool, optional): Skip existing Collections. Defaults to False.
Returns:
Iterable[Collection]:
"""

if not headers:
headers = DEFAULT_HEADERS

collections_endpoint = f"{url}/collections"
for collection in collections:
response = httpx.post(
collections_endpoint,
json=collection.to_dict(),
headers=headers,
verify=verify,
)
if response.status_code == 409:
if update:
collection_endpoint = f"{collections_endpoint}/{collection.id}"
response = httpx.put(
collection_endpoint,
json=collection.to_dict(),
headers=headers,
verify=verify,
)
if skip_existing:
continue

response.raise_for_status()
yield collection
49 changes: 47 additions & 2 deletions tests/conftest.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,60 @@
import json
from pathlib import Path

import pytest
import respx
from httpx import Response
from pytest_mock import MockerFixture
from test_data import STAC_COLLECTIONS, STAC_ITEMS


@pytest.fixture()
def mock_stac_collections(mocker: MockerFixture):
def mock_extract_stac_collections(mocker: MockerFixture):
client_mock = mocker.patch("eodm.extract.pystac_client.Client")
client_mock.open().get_collections().__iter__.return_value = STAC_COLLECTIONS


@pytest.fixture()
def mock_stac_items(mocker: MockerFixture):
def mock_extract_stac_items(mocker: MockerFixture):
client_mock = mocker.patch("eodm.extract.pystac_client.Client")
client_mock.open().search().item_collection().__iter__.return_value = STAC_ITEMS


@pytest.fixture()
def stac_collections(tmp_path: Path):
collection_json_path = tmp_path / "collections"

with collection_json_path.open("a", encoding="utf-8", newline="\n") as f:
for collection in STAC_COLLECTIONS:
f.write(json.dumps(collection.to_dict()))
f.write("\n")

return str(collection_json_path)


@pytest.fixture()
def stac_items(tmp_path):
item_json_path = tmp_path / "items"

with item_json_path.open("a", encoding="utf-8", newline="\n") as f:
for item in STAC_ITEMS:
f.write(json.dumps(item.to_dict()))
f.write("\n")

return str(item_json_path)


@pytest.fixture()
def mock_post_stac_api_items_endpoint(respx_mock: respx.MockRouter):
post_mock = respx_mock.post(
"https://example.com/stac-api/collections/sentinel-2-l2a/items"
).mock(return_value=Response(204))
return post_mock


@pytest.fixture()
def mock_post_stac_api_collections_endpoint(respx_mock: respx.MockRouter):
post_mock = respx_mock.post("https://example.com/stac-api/collections").mock(
return_value=Response(204)
)
return post_mock
38 changes: 36 additions & 2 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,50 @@
runner = CliRunner()


def test_extract_stacapi_items(mock_stac_items):
def test_extract_stacapi_items(mock_extract_stac_items):
result = runner.invoke(app, ["extract", "stac-api", "items", "sample", "sample"])
assert result.exit_code == 0
assert "Feature" in result.stdout


def test_extract_stacapi_collections(mock_stac_collections):
def test_extract_stacapi_collections(mock_extract_stac_collections):
result = runner.invoke(
app,
["extract", "stac-api", "collections", "sample"],
)
assert result.exit_code == 0
assert "Collection" in result.stdout


def test_load_stacapi_items(stac_items, mock_post_stac_api_items_endpoint):
result = runner.invoke(
app,
[
"load",
"stac-api",
"items",
"https://example.com/stac-api",
stac_items,
],
)
assert mock_post_stac_api_items_endpoint.called
assert result.exit_code == 0
assert "Feature" in result.stdout


def test_load_stacapi_collections(
stac_collections, mock_post_stac_api_collections_endpoint
):
result = runner.invoke(
app,
[
"load",
"stac-api",
"collections",
"https://example.com/stac-api",
stac_collections,
],
)
assert mock_post_stac_api_collections_endpoint.called
assert result.exit_code == 0
assert "Collection" in result.stdout
4 changes: 2 additions & 2 deletions tests/test_data.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from pystac import Collection, Item

STAC_ITEMS = (
STAC_ITEMS = list(
Item.from_dict(item)
for item in [
{
Expand Down Expand Up @@ -2901,7 +2901,7 @@
]
)

STAC_COLLECTIONS = (
STAC_COLLECTIONS = list(
Collection.from_dict(collection)
for collection in [
{
Expand Down

0 comments on commit 2f85e8c

Please sign in to comment.