diff --git a/trustgraph-cli/trustgraph/cli/load_structured_data.py b/trustgraph-cli/trustgraph/cli/load_structured_data.py index c208a874..eb0b139f 100644 --- a/trustgraph-cli/trustgraph/cli/load_structured_data.py +++ b/trustgraph-cli/trustgraph/cli/load_structured_data.py @@ -564,59 +564,69 @@ def _process_data_pipeline(input_file, descriptor_file, user, collection, sample def _send_to_trustgraph(objects, api_url, flow, batch_size=1000): - """Send ExtractedObject records to TrustGraph""" - from trustgraph.api import Api + """Send ExtractedObject records to TrustGraph using WebSocket""" + import json + import asyncio + from websockets.asyncio.client import connect try: - # Initialize TrustGraph API - api = Api(api_url) - flow_api = api.flow().id(flow) + # Construct objects import URL similar to load_knowledge pattern + if not api_url.endswith("/"): + api_url += "/" - # Process records in batches - total_records = len(objects) - logger.info(f"Importing {total_records} records to TrustGraph in batches of {batch_size}") + # Convert HTTP URL to WebSocket URL if needed + ws_url = api_url.replace("http://", "ws://").replace("https://", "wss://") + objects_url = ws_url + f"api/v1/flow/{flow}/import/objects" - imported_count = 0 - error_count = 0 + logger.info(f"Connecting to objects import endpoint: {objects_url}") - for i in range(0, total_records, batch_size): - batch = objects[i:i + batch_size] - batch_num = (i // batch_size) + 1 - - logger.info(f"Processing batch {batch_num}: records {i+1}-{min(i+batch_size, total_records)}") - - try: - # Send batch to objects import dispatcher - response = flow_api.call( - service="objects-import", - input_data=batch - ) + async def import_objects(): + async with connect(objects_url) as ws: + imported_count = 0 - imported_count += len(batch) - logger.info(f"Batch {batch_num} imported successfully ({len(batch)} records)") - print(f"āœ… Batch {batch_num}/{(total_records + batch_size - 1) // batch_size} imported ({len(batch)} records)") - - except Exception as e: - error_count += len(batch) - logger.error(f"Failed to import batch {batch_num}: {e}") - print(f"āŒ Batch {batch_num} failed: {e}") + for record in objects: + try: + # Send individual ExtractedObject records + await ws.send(json.dumps(record)) + imported_count += 1 + + if imported_count % 100 == 0: + logger.info(f"Imported {imported_count}/{len(objects)} records...") + print(f"āœ… Imported {imported_count}/{len(objects)} records...") + + except Exception as e: + logger.error(f"Failed to send record {imported_count + 1}: {e}") + print(f"āŒ Failed to send record {imported_count + 1}: {e}") + logger.info(f"Successfully imported {imported_count} records to TrustGraph") + return imported_count + + # Run the async import + imported_count = asyncio.run(import_objects()) + # Summary - logger.info(f"Import complete: {imported_count} imported, {error_count} failed") + total_records = len(objects) + failed_count = total_records - imported_count + print(f"\nšŸ“Š Import Summary:") print(f"- Total records: {total_records}") print(f"- Successfully imported: {imported_count}") - print(f"- Failed: {error_count}") + print(f"- Failed: {failed_count}") - if error_count > 0: - print(f"āš ļø {error_count} records failed to import. Check logs for details.") + if failed_count > 0: + print(f"āš ļø {failed_count} records failed to import. Check logs for details.") else: print("āœ… All records imported successfully!") return imported_count + except ImportError as e: + logger.error(f"Failed to import required modules: {e}") + print(f"Error: Required modules not available - {e}") + raise except Exception as e: - logger.error(f"TrustGraph import failed: {e}") + logger.error(f"Failed to import data to TrustGraph: {e}") + print(f"Import failed: {e}") raise