diff --git a/trustgraph-cli/scripts/tg-load-pdf b/trustgraph-cli/scripts/tg-load-pdf index a0d2b3bc..f27aaf5a 100755 --- a/trustgraph-cli/scripts/tg-load-pdf +++ b/trustgraph-cli/scripts/tg-load-pdf @@ -6,21 +6,19 @@ Loads a PDF document into TrustGraph processing. import pulsar from pulsar.schema import JsonSchema -import base64 import hashlib import argparse import os import time import uuid -from trustgraph.schema import Document, document_ingest_queue -from trustgraph.schema import Metadata, Triple, Value -from trustgraph.log_level import LogLevel -from trustgraph.knowledge import hash, to_uri, Uri +from trustgraph.api import Api +from trustgraph.knowledge import hash, to_uri from trustgraph.knowledge import PREF_PUBEV, PREF_DOC, PREF_ORG from trustgraph.knowledge import Organization, PublicationEvent from trustgraph.knowledge import DigitalDocument +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' @@ -28,24 +26,13 @@ class Loader: def __init__( self, - pulsar_host, - output_queue, + url, user, collection, - log_level, metadata, ): - self.client = pulsar.Client( - pulsar_host, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - - self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(Document), - chunking_enabled=True, - ) + self.api = Api(url) self.user = user self.collection = collection @@ -68,49 +55,18 @@ class Loader: id = to_uri(PREF_DOC, id) - triples = [] - - def emit(t): - triples.append(t) - self.metadata.id = id - self.metadata.emit(emit) - r = Document( - metadata=Metadata( - id=id, - metadata=[ - Triple( - s=Value( - value=t["s"], - is_uri=isinstance(t["s"], Uri) - ), - p=Value( - value=t["p"], - is_uri=isinstance(t["p"], Uri) - ), - o=Value( - value=t["o"], - is_uri=isinstance(t["o"], Uri) - ), - ) - for t in triples - ], - user=self.user, - collection=self.collection, - ), - data=base64.b64encode(data), + self.api.load_document( + document=data, id=id, metadata=self.metadata, +# user=self.user, +# collection=self.collection, ) - self.producer.send(r) - print(f"{file}: Loaded successfully.") except Exception as e: print(f"{file}: Failed: {str(e)}", flush=True) - - def __del__(self): - self.client.close() def main(): @@ -119,29 +75,20 @@ def main(): description=__doc__, ) - default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') - default_output_queue = document_ingest_queue - parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-o', '--output-queue', - default=default_output_queue, - help=f'Output queue (default: {default_output_queue})' - ) - - parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -183,7 +130,7 @@ def main(): ) parser.add_argument( - '--url', help=f'Document URL' + '--document-url', help=f'Document URL' ) parser.add_argument( @@ -194,14 +141,6 @@ def main(): '--identifier', '--id', help=f'Document ID' ) - parser.add_argument( - '-l', '--log-level', - type=LogLevel, - default=LogLevel.ERROR, - choices=list(LogLevel), - help=f'Output queue (default: info)' - ) - parser.add_argument( 'files', nargs='+', help=f'File to load' @@ -221,7 +160,7 @@ def main(): copyright_holder=args.copyright_holder, copyright_year=args.copyright_year, license=args.license, - url=args.url, + url=args.document_url, keywords=args.keyword, ) @@ -239,11 +178,9 @@ def main(): ) p = Loader( - pulsar_host=args.pulsar_host, - output_queue=args.output_queue, + url=args.url, user=args.user, collection=args.collection, - log_level=args.log_level, metadata=document, )