add data directory option, move data loader tool

Author: jonhealy1

README.md

@@ -104,11 +104,12 @@ Options:
   --base-url TEXT       Base URL of the STAC API  [required]
   --collection-id TEXT  ID of the collection to which items are added
   --use-bulk            Use bulk insert method for items
+  --data-dir PATH       Directory containing collection.json and feature
+                        collection file
   --help                Show this message and exit.
 ```
 
 ```shell
-cd data_loader   
 python3 data_loader.py --base-url http://localhost:8080
 ```  
 

data_loader/data_loader.py renamed to data_loader.py

@@ -1,23 +1,20 @@
-"""Data Loader CLI tool."""
 import json
 import os
-
 import click
 import requests
 
-# Define the directory where your data files are located
-DATA_DIR = os.path.join(os.path.dirname(__file__), "setup_data/")
-
-
-def load_data(filename):
-    """Load json data from a file."""
-    with open(os.path.join(DATA_DIR, filename)) as file:
+def load_data(data_dir, filename):
+    """Load json data from a file within the specified data directory."""
+    filepath = os.path.join(data_dir, filename)
+    if not os.path.exists(filepath):
+        click.secho(f"File not found: {filepath}", fg="red", err=True)
+        raise click.Abort()
+    with open(filepath) as file:
         return json.load(file)
 
-
-def load_collection(base_url, collection_id):
+def load_collection(base_url, collection_id, data_dir):
     """Load a STAC collection into the database."""
-    collection = load_data("collection.json")
+    collection = load_data(data_dir, "collection.json")
     collection["id"] = collection_id
     try:
         resp = requests.post(f"{base_url}/collections", json=collection)
@@ -28,45 +25,45 @@ def load_collection(base_url, collection_id):
             click.echo(f"Status code: {resp.status_code}")
             click.echo(f"Collection: {collection['id']} already exists")
     except requests.ConnectionError:
-        click.secho("Failed to connect", fg="red")
-
+        click.secho("Failed to connect", fg="red", err=True)
 
-def load_items(base_url, collection_id, use_bulk):
+def load_items(base_url, collection_id, use_bulk, data_dir):
     """Load STAC items into the database based on the method selected."""
-    feature_collection = load_data("sentinel-s2-l2a-cogs_0_100.json")
-    load_collection(base_url, collection_id)
+    # Attempt to dynamically find a suitable feature collection file
+    feature_files = [file for file in os.listdir(data_dir) if file.endswith('.json') and file != "collection.json"]
+    if not feature_files:
+        click.secho("No feature collection files found in the specified directory.", fg="red", err=True)
+        raise click.Abort()
+    feature_collection_file = feature_files[0]  # Use the first found feature collection file
+    feature_collection = load_data(data_dir, feature_collection_file)
+
+    load_collection(base_url, collection_id, data_dir)
     if use_bulk:
-        load_items_bulk_insert(base_url, collection_id, feature_collection)
+        load_items_bulk_insert(base_url, collection_id, feature_collection, data_dir)
     else:
-        load_items_one_by_one(base_url, collection_id, feature_collection)
+        load_items_one_by_one(base_url, collection_id, feature_collection, data_dir)
 
-
-def load_items_one_by_one(base_url, collection_id, feature_collection):
+def load_items_one_by_one(base_url, collection_id, feature_collection, data_dir):
     """Load STAC items into the database one by one."""
     for feature in feature_collection["features"]:
         try:
             feature["collection"] = collection_id
-            resp = requests.post(
-                f"{base_url}/collections/{collection_id}/items", json=feature
-            )
+            resp = requests.post(f"{base_url}/collections/{collection_id}/items", json=feature)
             if resp.status_code == 200:
                 click.echo(f"Status code: {resp.status_code}")
                 click.echo(f"Added item: {feature['id']}")
             elif resp.status_code == 409:
                 click.echo(f"Status code: {resp.status_code}")
                 click.echo(f"Item: {feature['id']} already exists")
         except requests.ConnectionError:
-            click.secho("Failed to connect", fg="red")
-
+            click.secho("Failed to connect", fg="red", err=True)
 
-def load_items_bulk_insert(base_url, collection_id, feature_collection):
+def load_items_bulk_insert(base_url, collection_id, feature_collection, data_dir):
     """Load STAC items into the database via bulk insert."""
     try:
         for i, _ in enumerate(feature_collection["features"]):
             feature_collection["features"][i]["collection"] = collection_id
-        resp = requests.post(
-            f"{base_url}/collections/{collection_id}/items", json=feature_collection
-        )  # Adjust this endpoint as necessary
+        resp = requests.post(f"{base_url}/collections/{collection_id}/items", json=feature_collection)  # Adjust this endpoint as necessary
         if resp.status_code == 200:
             click.echo(f"Status code: {resp.status_code}")
             click.echo("Bulk inserted items successfully.")
@@ -77,21 +74,16 @@ def load_items_bulk_insert(base_url, collection_id, feature_collection):
             click.echo(f"Status code: {resp.status_code}")
             click.echo("Conflict detected, some items might already exist.")
     except requests.ConnectionError:
-        click.secho("Failed to connect", fg="red")
-
+        click.secho("Failed to connect", fg="red", err=True)
 
 @click.command()
 @click.option("--base-url", required=True, help="Base URL of the STAC API")
-@click.option(
-    "--collection-id",
-    default="test-collection",
-    help="ID of the collection to which items are added",
-)
+@click.option("--collection-id", default="test-collection", help="ID of the collection to which items are added")
 @click.option("--use-bulk", is_flag=True, help="Use bulk insert method for items")
-def main(base_url, collection_id, use_bulk):
+@click.option("--data-dir", type=click.Path(exists=True), default="sample_data/", help="Directory containing collection.json and feature collection file")
+def main(base_url, collection_id, use_bulk, data_dir):
     """Load STAC items into the database."""
-    load_items(base_url, collection_id, use_bulk)
-
+    load_items(base_url, collection_id, use_bulk, data_dir)
 
 if __name__ == "__main__":
-    main()
+    main()