pdf_util.py

importtime
importuuid
fromdatetimeimportdatetime
fromenumimportEnum
fromurllib.parseimportquote

fromsqlalchemyimporttext

fromappimports3_conn, logger, db

fromconstimportPDF_ARCHIVE_BUCKET, GROBID_XML_BUCKET, PDF_ARCHIVE_BUCKET_NEW
froms3_utilimportcheck_exists, get_object, harvest_pdf_table, s3


classPDFVersion(Enum):
PUBLISHED='published'
ACCEPTED='accepted'
SUBMITTED='submitted'

defs3_key(self, doi):
returnf"{self.s3_prefix}{quote(doi, safe='')}.pdf"

defgrobid_s3_key(self, doi):
returnf'{self.s3_prefix}{quote(doi, safe="")}.xml'

@property
defs3_prefix(self):
ifnotself==PDFVersion.PUBLISHED:
returnf'{self.value}_'
return''

defs3_url(self, doi):
returnf's3://{PDF_ARCHIVE_BUCKET}/{self.s3_key(doi)}'

@classmethod
deffrom_version_str(cls, version_str: str):
ifnotversion_str:
returnNone
forversionincls:
ifversion.valueinversion_str.lower():
returnversion
returnNone

defvalid_in_s3(self, doi) ->bool:
returncheck_valid_pdf(PDF_ARCHIVE_BUCKET, self.s3_key(doi))

defin_s3(self, doi) ->bool:
returncheck_exists(PDF_ARCHIVE_BUCKET, self.s3_key(doi))

defgrobid_in_s3(self, doi):
returncheck_exists(GROBID_XML_BUCKET, self.grobid_s3_key(doi))

defget_grobid_xml_obj(self, doi):
returnget_object(GROBID_XML_BUCKET, self.grobid_s3_key(doi))

defget_pdf_obj(self, doi):
returnget_object(PDF_ARCHIVE_BUCKET, self.s3_key(doi))


defsave_pdf(doi, content, version=PDFVersion.PUBLISHED):
ifnotcontent:
returnFalse
logger.info(
f'saving {len(content)} characters to {version.s3_url(doi)}')
try:
s3_conn.put_object(
Body=content,
Bucket=PDF_ARCHIVE_BUCKET,
Key=version.s3_key(doi)
 )
returnTrue
exceptExceptionase:
logger.error(f'failed to save pdf: {e}')
returnFalse


defsave_pdf_new(content, native_id, native_id_ns, version=PDFVersion.PUBLISHED, url='', resolved_url=''):
ifnotcontent:
return
new_key=str(uuid.uuid4()) +'.pdf'
encoded_url=quote(str(urlor''))
encoded_resolved_url=quote(str(resolved_urlor''))
s3.put_object(
Bucket=PDF_ARCHIVE_BUCKET_NEW,
Key=new_key,
Body=content,
Metadata={
'url': encoded_url,
'resolved_url': encoded_resolved_url,
'created_date': datetime.utcnow().isoformat(),
'created_timestamp': str(int(time.time())),
'id': new_key.replace('.pdf', ''),
'native_id': native_id.lower().strip(),
'native_id_namespace': native_id_ns
 })
item= {
'id': new_key.replace('.pdf', ''),
'native_id': native_id.lower().strip(),
'native_id_namespace': native_id_ns,
'url': url,
'type': version.value,
's3_key': new_key,
's3_path': f's3://{PDF_ARCHIVE_BUCKET_NEW}/{new_key}',
'created_timestamp': int(time.time()),
'created_date': datetime.utcnow().isoformat()
 }
ifnative_id_ns=='doi':
item['normalized_doi'] =native_id
harvest_pdf_table.put_item(Item=item)


defenqueue_pdf_parsing(doi, version: PDFVersion=PDFVersion.PUBLISHED,
commit=True):
stmnt=text(
"INSERT INTO recordthresher.pdf_update_ingest (doi, pdf_version) VALUES (:doi, :version) ON CONFLICT(doi, pdf_version) DO UPDATE SET finished = NULL;").bindparams(
doi=doi, version=version.value)
db.session.execute(stmnt)
ifcommit:
db.session.commit()


defcheck_valid_pdf(bucket, key, s3=None, _raise=False):
obj=get_object(bucket, key, s3=s3, _raise=_raise)
ifnotobj:
returnFalse
contents=obj['Body'].read()
ifcontentsisnotNone:
returnis_pdf(contents)
returnFalse


defis_pdf(contents: bytes) ->bool:
returncontents.startswith(b"%PDF-")