1
- """Database ingestion script ."""
1
+ """Data Loader CLI tool ."""
2
2
import json
3
3
import os
4
4
@@ -31,17 +31,23 @@ def load_collection(base_url, collection_id):
31
31
click .secho ("Failed to connect" , fg = "red" )
32
32
33
33
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 ."""
36
36
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
+
39
43
44
+ def load_items_one_by_one (base_url , collection_id , feature_collection ):
45
+ """Load STAC items into the database one by one."""
40
46
for feature in feature_collection ["features" ]:
41
47
try :
42
- feature ["collection" ] = collection
48
+ feature ["collection" ] = collection_id
43
49
resp = requests .post (
44
- f"{ base_url } /collections/{ collection } /items" , json = feature
50
+ f"{ base_url } /collections/{ collection_id } /items" , json = feature
45
51
)
46
52
if resp .status_code == 200 :
47
53
click .echo (f"Status code: { resp .status_code } " )
@@ -53,16 +59,38 @@ def load_items(base_url, collection_id):
53
59
click .secho ("Failed to connect" , fg = "red" )
54
60
55
61
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
+
56
83
@click .command ()
57
84
@click .option ("--base-url" , required = True , help = "Base URL of the STAC API" )
58
85
@click .option (
59
86
"--collection-id" ,
60
87
default = "test-collection" ,
61
88
help = "ID of the collection to which items are added" ,
62
89
)
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 ):
64
92
"""Load STAC items into the database."""
65
- load_items (base_url , collection_id )
93
+ load_items (base_url , collection_id , use_bulk )
66
94
67
95
68
96
if __name__ == "__main__" :
0 commit comments