From 5ce01ab8ca992023c320294f20356735d550d4d9 Mon Sep 17 00:00:00 2001 From: Bryan Newbold Date: Wed, 10 Nov 2021 15:16:58 -0800 Subject: simple persist worker/tool to backfill grobid_refs --- python/sandcrawler/persist.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) (limited to 'python/sandcrawler') diff --git a/python/sandcrawler/persist.py b/python/sandcrawler/persist.py index 6847e2e..d753380 100644 --- a/python/sandcrawler/persist.py +++ b/python/sandcrawler/persist.py @@ -678,6 +678,12 @@ class PersistHtmlTeiXmlWorker(GenericPersistDocWorker): class PersistCrossrefWorker(SandcrawlerWorker): + """ + Pushes Crossref API JSON records into postgresql. Can also talk to GROBID, + parsed 'unstructured' references, and push the results in to postgresql at + the same time. + """ + def __init__( self, db_url: str, @@ -739,3 +745,37 @@ class PersistCrossrefWorker(SandcrawlerWorker): self.db.commit() return [] + + +class PersistGrobidRefsWorker(SandcrawlerWorker): + """ + Simple persist worker to backfill GROBID references in to postgresql + locally. Consumes the JSON output from GROBID CrossrefRefsWorker. + """ + + def __init__(self, db_url: str, **kwargs): + super().__init__(**kwargs) + self.db = SandcrawlerPostgresClient(db_url) + self.cur = self.db.conn.cursor() + + def process(self, record: Any, key: Optional[str] = None) -> Any: + """Only do batches (as transactions)""" + raise NotImplementedError + + def push_batch(self, batch: list) -> list: + self.counts["total"] += len(batch) + + refs_batch = [] + for record in batch: + assert record["source"] + assert record["source_id"] + refs_batch.append(record) + + resp = self.db.insert_grobid_refs(self.cur, refs_batch) + if len(refs_batch) < len(batch): + self.counts["skip"] += len(batch) - len(refs_batch) + self.counts["insert-grobid_refs"] += resp[0] + self.counts["update-grobid_refs"] += resp[1] + + self.db.commit() + return [] -- cgit v1.2.3