From e0eb007f637fbeca6c5dc130c21eabe474758a63 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Fri, 5 Sep 2025 15:20:21 +0100
Subject: [PATCH] Structured data loader CLI

---
 .../trustgraph/cli/load_structured_data.py    | 48 ++++++++++++-------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/trustgraph-cli/trustgraph/cli/load_structured_data.py b/trustgraph-cli/trustgraph/cli/load_structured_data.py
index 37d46d65..57e77579 100644
--- a/trustgraph-cli/trustgraph/cli/load_structured_data.py
+++ b/trustgraph-cli/trustgraph/cli/load_structured_data.py
@@ -804,32 +804,46 @@ def load_structured_data(
                 print(json.dumps(processed_records[0], indent=2))
             return
         
-        # Import to TrustGraph using objects dispatcher
+        # Import to TrustGraph using objects import endpoint via WebSocket
         logger.info(f"Importing {len(processed_records)} records to TrustGraph...")
         
         try:
-            import trustgraph.publish.writer as writer
+            import asyncio
+            import json
+            from websockets.asyncio.client import connect
             
-            # Initialize the objects writer
-            objects_writer = writer.Writer(api_url, "objects")
+            # Construct objects import URL similar to load_knowledge pattern
+            if not api_url.endswith("/"):
+                api_url += "/"
             
-            # Send records in batches
-            imported_count = 0
-            for i in range(0, len(processed_records), batch_size):
-                batch = processed_records[i:i + batch_size]
-                
-                for record in batch:
-                    objects_writer.write(record)
-                    imported_count += 1
-                
-                logger.info(f"Imported {imported_count}/{len(processed_records)} records...")
+            # Convert HTTP URL to WebSocket URL if needed
+            ws_url = api_url.replace("http://", "ws://").replace("https://", "wss://")
+            objects_url = ws_url + "api/v1/flow/default/import/objects"
             
-            logger.info(f"Successfully imported {imported_count} records to TrustGraph")
+            logger.info(f"Connecting to objects import endpoint: {objects_url}")
+            
+            async def import_objects():
+                async with connect(objects_url) as ws:
+                    imported_count = 0
+                    
+                    for record in processed_records:
+                        # 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(processed_records)} records...")
+                    
+                    logger.info(f"Successfully imported {imported_count} records to TrustGraph")
+                    return imported_count
+            
+            # Run the async import
+            imported_count = asyncio.run(import_objects())
             print(f"Import completed: {imported_count} records imported to schema '{schema_name}'")
             
         except ImportError as e:
-            logger.error(f"Failed to import TrustGraph writer: {e}")
-            print(f"Error: TrustGraph writer not available - {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"Failed to import data to TrustGraph: {e}")