Skip to content

Commit 50c3d26

Browse files
committed
update data loader to accept base_url
1 parent 866d45c commit 50c3d26

File tree

3 files changed

+40
-43
lines changed

3 files changed

+40
-43
lines changed

README.md

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,13 @@ get the next page of results.
9393
curl -X "GET" "http://localhost:8080/collections?limit=1&token=example_token"
9494
```
9595

96+
## Ingesting Sample Data
97+
98+
```shell
99+
cd data_loader
100+
python3 data_loader.py --base-url http://localhost:8080
101+
```
102+
96103
## Testing
97104

98105
```shell

data_loader/data_loader.py

Lines changed: 32 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,69 +1,59 @@
11
"""Database ingestion script."""
22
import json
33
import os
4-
import sys
5-
64
import click
75
import requests
86

9-
if len(sys.argv) != 2:
10-
print("Usage: python data_loader.py <opensearch|elasticsearch>")
11-
sys.exit(1)
12-
7+
# Define the directory where your data files are located
138
DATA_DIR = os.path.join(os.path.dirname(__file__), "setup_data/")
149

15-
backend = sys.argv[1].lower()
16-
17-
if backend == "opensearch":
18-
STAC_API_BASE_URL = "http://localhost:8082"
19-
elif backend == "elasticsearch":
20-
STAC_API_BASE_URL = "http://localhost:8080"
21-
else:
22-
print("Invalid backend tag. Enter either 'opensearch' or 'elasticsearch'.")
23-
24-
2510
def load_data(filename):
26-
"""Load json data."""
11+
"""Load json data from a file."""
2712
with open(os.path.join(DATA_DIR, filename)) as file:
2813
return json.load(file)
2914

30-
31-
def load_collection(collection_id):
32-
"""Load stac collection into the database."""
15+
def load_collection(base_url, collection_id):
16+
"""Load a STAC collection into the database."""
3317
collection = load_data("collection.json")
3418
collection["id"] = collection_id
3519
try:
36-
resp = requests.post(f"{STAC_API_BASE_URL}/collections", json=collection)
20+
resp = requests.post(f"{base_url}/collections", json=collection)
3721
if resp.status_code == 200:
38-
print(f"Status code: {resp.status_code}")
39-
print(f"Added collection: {collection['id']}")
22+
click.echo(f"Status code: {resp.status_code}")
23+
click.echo(f"Added collection: {collection['id']}")
4024
elif resp.status_code == 409:
41-
print(f"Status code: {resp.status_code}")
42-
print(f"Collection: {collection['id']} already exists")
25+
click.echo(f"Status code: {resp.status_code}")
26+
click.echo(f"Collection: {collection['id']} already exists")
4327
except requests.ConnectionError:
44-
click.secho("failed to connect")
45-
28+
click.secho("Failed to connect", fg="red")
4629

47-
def load_items():
48-
"""Load stac items into the database."""
30+
def load_items(base_url, collection_id):
31+
"""Load STAC items into the database."""
4932
feature_collection = load_data("sentinel-s2-l2a-cogs_0_100.json")
50-
collection = "test-collection"
51-
load_collection(collection)
33+
collection = collection_id
34+
load_collection(base_url, collection)
5235

5336
for feature in feature_collection["features"]:
5437
try:
5538
feature["collection"] = collection
56-
resp = requests.post(
57-
f"{STAC_API_BASE_URL}/collections/{collection}/items", json=feature
58-
)
39+
resp = requests.post(f"{base_url}/collections/{collection}/items", json=feature)
5940
if resp.status_code == 200:
60-
print(f"Status code: {resp.status_code}")
61-
print(f"Added item: {feature['id']}")
41+
click.echo(f"Status code: {resp.status_code}")
42+
click.echo(f"Added item: {feature['id']}")
6243
elif resp.status_code == 409:
63-
print(f"Status code: {resp.status_code}")
64-
print(f"Item: {feature['id']} already exists")
44+
click.echo(f"Status code: {resp.status_code}")
45+
click.echo(f"Item: {feature['id']} already exists")
6546
except requests.ConnectionError:
66-
click.secho("failed to connect")
67-
68-
69-
load_items()
47+
click.secho("Failed to connect", fg="red")
48+
49+
@click.command()
50+
@click.option('--base-url', required=True, help='Base URL of the STAC API')
51+
@click.option('--collection-id', default='test-collection', help='ID of the collection to which items are added')
52+
def main(base_url, collection_id):
53+
"""
54+
Load STAC items into the database.
55+
"""
56+
load_items(base_url, collection_id)
57+
58+
if __name__ == '__main__':
59+
main()

data_loader/setup_data/collection.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
2-
"id":"sentinel-s2-l2a-cogs",
2+
"id":"sentinel-s2-l2a-cogs-test",
33
"stac_version":"1.0.0",
44
"description":"Sentinel-2a and Sentinel-2b imagery, processed to Level 2A (Surface Reflectance) and converted to Cloud-Optimized GeoTIFFs",
55
"links":[

0 commit comments

Comments
 (0)