Skip to content

Commit 8b4720a

Browse files
committed
add bulk insert option
1 parent 491496b commit 8b4720a

File tree

2 files changed

+50
-10
lines changed

2 files changed

+50
-10
lines changed

README.md

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,19 @@ 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
96+
## Ingesting Sample Data CLI Tool
97+
98+
```shell
99+
Usage: data_loader.py [OPTIONS]
100+
101+
Load STAC items into the database.
102+
103+
Options:
104+
--base-url TEXT Base URL of the STAC API [required]
105+
--collection-id TEXT ID of the collection to which items are added
106+
--use-bulk Use bulk insert method for items
107+
--help Show this message and exit.
108+
```
97109

98110
```shell
99111
cd data_loader

data_loader/data_loader.py

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
"""Database ingestion script."""
1+
"""Data Loader CLI tool."""
22
import json
33
import os
44

@@ -31,17 +31,23 @@ def load_collection(base_url, collection_id):
3131
click.secho("Failed to connect", fg="red")
3232

3333

34-
def load_items(base_url, collection_id):
35-
"""Load STAC items into the database."""
34+
def load_items(base_url, collection_id, use_bulk):
35+
"""Load STAC items into the database based on the method selected."""
3636
feature_collection = load_data("sentinel-s2-l2a-cogs_0_100.json")
37-
collection = collection_id
38-
load_collection(base_url, collection)
37+
load_collection(base_url, collection_id)
38+
if use_bulk:
39+
load_items_bulk_insert(base_url, collection_id, feature_collection)
40+
else:
41+
load_items_one_by_one(base_url, collection_id, feature_collection)
42+
3943

44+
def load_items_one_by_one(base_url, collection_id, feature_collection):
45+
"""Load STAC items into the database one by one."""
4046
for feature in feature_collection["features"]:
4147
try:
42-
feature["collection"] = collection
48+
feature["collection"] = collection_id
4349
resp = requests.post(
44-
f"{base_url}/collections/{collection}/items", json=feature
50+
f"{base_url}/collections/{collection_id}/items", json=feature
4551
)
4652
if resp.status_code == 200:
4753
click.echo(f"Status code: {resp.status_code}")
@@ -53,16 +59,38 @@ def load_items(base_url, collection_id):
5359
click.secho("Failed to connect", fg="red")
5460

5561

62+
def load_items_bulk_insert(base_url, collection_id, feature_collection):
63+
"""Load STAC items into the database via bulk insert."""
64+
try:
65+
for i, _ in enumerate(feature_collection["features"]):
66+
feature_collection["features"][i]["collection"] = collection_id
67+
resp = requests.post(
68+
f"{base_url}/collections/{collection_id}/items", json=feature_collection
69+
) # Adjust this endpoint as necessary
70+
if resp.status_code == 200:
71+
click.echo(f"Status code: {resp.status_code}")
72+
click.echo("Bulk inserted items successfully.")
73+
elif resp.status_code == 204:
74+
click.echo(f"Status code: {resp.status_code}")
75+
click.echo("Bulk update successful, no content returned.")
76+
elif resp.status_code == 409:
77+
click.echo(f"Status code: {resp.status_code}")
78+
click.echo("Conflict detected, some items might already exist.")
79+
except requests.ConnectionError:
80+
click.secho("Failed to connect", fg="red")
81+
82+
5683
@click.command()
5784
@click.option("--base-url", required=True, help="Base URL of the STAC API")
5885
@click.option(
5986
"--collection-id",
6087
default="test-collection",
6188
help="ID of the collection to which items are added",
6289
)
63-
def main(base_url, collection_id):
90+
@click.option("--use-bulk", is_flag=True, help="Use bulk insert method for items")
91+
def main(base_url, collection_id, use_bulk):
6492
"""Load STAC items into the database."""
65-
load_items(base_url, collection_id)
93+
load_items(base_url, collection_id, use_bulk)
6694

6795

6896
if __name__ == "__main__":

0 commit comments

Comments
 (0)