|
1 | 1 | """Database ingestion script."""
|
2 | 2 | import json
|
3 | 3 | import os
|
4 |
| -import sys |
5 |
| - |
6 | 4 | import click
|
7 | 5 | import requests
|
8 | 6 |
|
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 |
13 | 8 | DATA_DIR = os.path.join(os.path.dirname(__file__), "setup_data/")
|
14 | 9 |
|
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 |
| - |
25 | 10 | def load_data(filename):
|
26 |
| - """Load json data.""" |
| 11 | + """Load json data from a file.""" |
27 | 12 | with open(os.path.join(DATA_DIR, filename)) as file:
|
28 | 13 | return json.load(file)
|
29 | 14 |
|
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.""" |
33 | 17 | collection = load_data("collection.json")
|
34 | 18 | collection["id"] = collection_id
|
35 | 19 | try:
|
36 |
| - resp = requests.post(f"{STAC_API_BASE_URL}/collections", json=collection) |
| 20 | + resp = requests.post(f"{base_url}/collections", json=collection) |
37 | 21 | 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']}") |
40 | 24 | 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") |
43 | 27 | except requests.ConnectionError:
|
44 |
| - click.secho("failed to connect") |
45 |
| - |
| 28 | + click.secho("Failed to connect", fg="red") |
46 | 29 |
|
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.""" |
49 | 32 | 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) |
52 | 35 |
|
53 | 36 | for feature in feature_collection["features"]:
|
54 | 37 | try:
|
55 | 38 | 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) |
59 | 40 | 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']}") |
62 | 43 | 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") |
65 | 46 | 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() |
0 commit comments