clean_pdfs.py

importgzip
importlogging
importos
importre
importtime
fromargparseimportArgumentParser
fromdatetimeimportdatetime
fromioimportBytesIO
fromqueueimportQueue, Empty
fromthreadingimportThread, Lock
fromurllib.parseimportunquote

importboto3
fromrequestsimportHTTPError
fromsqlalchemyimportcreate_engine, text
fromtenacityimportretry_if_exception_type, stop_after_attempt, retry

fromappimportapp, logger
fromhttp_cacheimporthttp_get

S3_PDF_BUCKET_NAME=os.getenv('AWS_S3_PDF_BUCKET')

OADOI_DB_ENGINE=create_engine(app.config['SQLALCHEMY_DATABASE_URI'])

DECOMPRESSED_LOCK=Lock()
REDOWNLOADED_LOCK=Lock()
FAILED_AND_DELETED_LOCK=Lock()
TOTAL_ATTEMPTED_LOCK=Lock()
OK_LOCK=Lock()
EMPTY_PDF_URL_LOCK=Lock()

DECOMPRESSED=0
REDOWNLOADED=0
FAILED_AND_DELETED=0
TOTAL_ATTEMPTED=0
OK_COUNT=0
EMPTY_PDF_URL=0

libs_to_mum= [
'boto',
'boto3',
'botocore',
's3transfer'
]

forlibinlibs_to_mum:
logging.getLogger(lib).setLevel(logging.CRITICAL)



definc_decompressed():
globalDECOMPRESSED
withDECOMPRESSED_LOCK:
DECOMPRESSED+=1


definc_redownloaded():
globalREDOWNLOADED
withREDOWNLOADED_LOCK:
REDOWNLOADED+=1


definc_failed():
globalFAILED_AND_DELETED
withFAILED_AND_DELETED_LOCK:
FAILED_AND_DELETED+=1


definc_total():
globalTOTAL_ATTEMPTED
withTOTAL_ATTEMPTED_LOCK:
TOTAL_ATTEMPTED+=1


definc_ok():
globalOK_COUNT
withOK_LOCK:
OK_COUNT+=1


definc_empty_pdf_url():
globalEMPTY_PDF_URL
withEMPTY_PDF_URL_LOCK:
EMPTY_PDF_URL+=1


defmake_s3():
session=boto3.session.Session()
returnsession.client('s3',
aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'),
aws_secret_access_key=os.getenv(
'AWS_SECRET_ACCESS_KEY'))


classInvalidPDFException(Exception):
pass


# @retry(retry=retry_if_exception_type(
# InvalidPDFException) | retry_if_exception_type(HTTPError),
# stop=stop_after_attempt(3), reraise=True)
deffetch_pdf(url):
r=http_get(url, ask_slowly=True)
# r = requests.get(url, headers=HEADERS, proxies=CRAWLERA_PROXIES, verify=False)
r.raise_for_status()
# content = r.content
content=r.content_big()
ifnotisinstance(content, bytes):
content=content.encode()
ifnotcontent.startswith(b'%PDF'):
raiseInvalidPDFException(f'Not a valid PDF document: {url}')
returncontent


defdownload_pdf(url, key, s3):
content=fetch_pdf(url)
s3.upload_fileobj(BytesIO(content), S3_PDF_BUCKET_NAME,
key)


defupload(key, body, s3):
s3.upload_fileobj(BytesIO(body), S3_PDF_BUCKET_NAME,
key)


defkey_to_doi(key):
returnre.sub('.pdf$', '', unquote(key))


defprocess_pdf_keys(q: Queue):
s3=make_s3()
withOADOI_DB_ENGINE.connect() asconn:
whileTrue:
try:
key=q.get(timeout=15)
obj_details=s3.get_object(Bucket=S3_PDF_BUCKET_NAME, Key=key)
body=obj_details['Body'].read()
zipped=False
ifbody[:3] ==b'\x1f\x8b\x08':
body=gzip.decompress(body)
zipped=True
ifbody.startswith(b'%PDF') andzipped:
upload(key, body, s3)
inc_decompressed()
elifnotbody.startswith(b'%PDF'):
# object is not PDF, need to attempt to re-download
row=conn.execute(text(
'SELECT scrape_pdf_url FROM pub WHERE id = :doi').bindparams(
doi=key_to_doi(key))).one()
ifrow['scrape_pdf_url'] isNone:
inc_empty_pdf_url()
continue
download_pdf(row['scrape_pdf_url'], key, s3)
inc_redownloaded()
else:
inc_ok()
exceptEmpty:
break
except (HTTPError, InvalidPDFException) ase:
s3.delete_object(Bucket=S3_PDF_BUCKET_NAME, Key=key)
inc_failed()
logger.error(
f'Unable to re-download PDF: {row["scrape_pdf_url"]} - {e}\nDeleted key {key}')
exceptExceptionase:
logger.exception(f'Error with key: {key}', exc_info=True)
finally:
inc_total()


defprint_stats():
start=datetime.now()
whileTrue:
now=datetime.now()
hrs_running= (now-start).total_seconds() / (60*60)
rate_per_hr=round(TOTAL_ATTEMPTED/hrs_running, 2)
redownloaded_pct=round(REDOWNLOADED/TOTAL_ATTEMPTED,
4) *100ifTOTAL_ATTEMPTEDelse0
logger.info(
f'Attempted count: {TOTAL_ATTEMPTED} | '
f'Ok count : {OK_COUNT} | '
f'Redownloaded count: {REDOWNLOADED} | '
f'Redownloaded %: {redownloaded_pct}% | '
f'Deleted count: {FAILED_AND_DELETED} | '
f'Decompressed count: {DECOMPRESSED} | '
f'Empty PDF URL count: {EMPTY_PDF_URL} | '
f'Rate: {rate_per_hr}/hr | '
f'Hrs running: {hrs_running}hrs')
time.sleep(5)


defenqueue_s3_keys(q: Queue):
s3=make_s3()
has_more=True
count=0
cont_token=None
whilehas_more:
kwargs= {'Bucket': S3_PDF_BUCKET_NAME}
ifcont_token:
kwargs['ContinuationToken'] =cont_token
page=s3.list_objects_v2(**kwargs)
count+=page['MaxKeys']
has_more=bool(page.get('NextContinuationToken'))
cont_token=page.get('NextContinuationToken')
forobjinpage['Contents']:
q.put(obj['Key'])


defparse_args():
parser=ArgumentParser()
parser.add_argument('--threads', '-t',
default=10,
type=int,
help='Number of threads to process PDF keys')
args=parser.parse_args()
env_dt=int(os.getenv('PDF_CLEAN_THREADS', 0))
ifenv_dt:
args.download_threads=env_dt
returnargs


defmain():
args=parse_args()

q=Queue(maxsize=args.threads+1)
Thread(target=print_stats, daemon=True).start()
Thread(target=enqueue_s3_keys, args=(q,)).start()

threads= []
for_inrange(args.threads):
t=Thread(target=process_pdf_keys, args=(q,))
t.start()
threads.append(t)
fortinthreads:
t.join()


if__name__=='__main__':
main()