Updated data ingestion to use spanner API

packages/datacommons-api/datacommons_api/api_cli.py

@@ -16,9 +16,10 @@
 import uvicorn
 
 from datacommons_api.app import app
-from datacommons_api.core.config import initialize_config
+from datacommons_api.core.config import get_config, initialize_config
 from datacommons_api.core.logging import get_logger, setup_logging
-from datacommons_db.session import initialize_db
+from datacommons_db.session import get_session, initialize_db
+from datacommons_api.services.graph_service import GraphService
 
 setup_logging()
 logger = get_logger(__name__)
@@ -72,3 +73,42 @@ def start(
         port=port,
         reload=reload,
     )
+
+
+@api.command()
+@click.option("--gcp-project-id", help="GCP project id.", required=True)
+@click.option(
+    "--gcp-spanner-instance-id", help="GCP Spanner instance id.", required=True
+)
+@click.option(
+    "--gcp-spanner-database-name", help="GCP Spanner database name.", required=True
+)
+@click.option("--yes", is_flag=True, help="Skip confirmation prompt.")
+def drop_tables(
+    gcp_project_id: str,
+    gcp_spanner_instance_id: str,
+    gcp_spanner_database_name: str,
+    yes: bool,
+):
+    """Drop Node and Edge tables from the graph database."""
+    # TODO: Refactor this method to only drop the data from the tables, not the tables themselves.
+    if not yes:
+        click.confirm(
+            "Are you sure you want to drop the Node and Edge tables?", abort=True
+        )
+
+    logger.info("Dropping Node and Edge tables from the graph database")
+    initialize_config(
+        gcp_project_id=gcp_project_id,
+        gcp_spanner_instance_id=gcp_spanner_instance_id,
+        gcp_spanner_database_name=gcp_spanner_database_name,
+    )
+    config = get_config()
+    db = get_session(
+        config.GCP_PROJECT_ID,
+        config.GCP_SPANNER_INSTANCE_ID,
+        config.GCP_SPANNER_DATABASE_NAME,
+    )
+    graph_service = GraphService(db)
+    graph_service.drop_tables()
+    logger.info("Successfully dropped Node and Edge tables")

packages/datacommons-api/datacommons_api/core/config.py

@@ -65,7 +65,7 @@ def validate_config_or_exit(config: Config) -> None:
     # Ensure GCP Spanner is configured
     for var in REQUIRED_ENV_VARS:
         if not getattr(config, var):
-            logger.error("Environment variable %s must be set", var)
+            logger.error("Config variable %s must be set", var)
             sys.exit(1)
 
 

packages/datacommons-api/datacommons_api/endpoints/routers/node_router.py

@@ -29,7 +29,13 @@
 
 
 # JSON-LD endpoint
-@router.get("/nodes/", response_model=JSONLDDocument, response_model_exclude_none=True)
+@router.get("/nodes", response_model=JSONLDDocument, response_model_exclude_none=True)
+@router.get(
+    "/nodes/",
+    response_model=JSONLDDocument,
+    response_model_exclude_none=True,
+    include_in_schema=False,
+)
 def get_nodes(
     limit: int = DEFAULT_NODE_FETCH_LIMIT,
     type_filter: Annotated[
@@ -44,7 +50,13 @@ def get_nodes(
     return graph_service.get_graph_nodes(limit=limit, type_filter=type_filter)
 
 
-@router.post("/nodes/", response_model=UpdateResponse, response_model_exclude_none=True)
+@router.post("/nodes", response_model=UpdateResponse, response_model_exclude_none=True)
+@router.post(
+    "/nodes/",
+    response_model=UpdateResponse,
+    response_model_exclude_none=True,
+    include_in_schema=False,
+)
 def insert_nodes(
     jsonld: JSONLDDocument,
     graph_service: Annotated[GraphService, Depends(with_graph_service)] = None,