queue_pdf_url_check.py

importargparse
importlogging
importos
importrandom
fromdatetimeimportdatetime
frommultiprocessingimportPool, current_process
fromtimeimportsleep
fromtimeimporttime

fromsqlalchemyimportorm
fromsqlalchemyimporttext
fromsqlalchemy.ormimportmake_transient

fromappimportdb
fromappimportlogger
fromhttp_cacheimporthttp_get, get_session_id
frompdf_urlimportPdfUrl
fromqueue_mainimportDbQueue
fromutilimportelapsed
fromutilimportrun_sql
fromutilimportsafe_commit
fromwebpageimportis_a_pdf_page

importendpoint# magic
importpmh_record# magic


defcheck_pdf_urls(pdf_urls):
forurlinpdf_urls:
make_transient(url)

# free up the connection while doing net IO
safe_commit(db)
db.engine.dispose()

req_pool=get_request_pool()

checked_pdf_urls=req_pool.map(get_pdf_url_status, pdf_urls, chunksize=1)
req_pool.close()
req_pool.join()

row_dicts= [x.__dict__forxinchecked_pdf_urls]
forrow_dictinrow_dicts:
row_dict.pop('_sa_instance_state')

db.session.bulk_update_mappings(PdfUrl, row_dicts)

start_time=time()
commit_success=safe_commit(db)
ifnotcommit_success:
logger.info("COMMIT fail")
logger.info("commit took {} seconds".format(elapsed(start_time, 2)))


defget_pdf_url_status(pdf_url):
worker=current_process()
logger.info('{} checking pdf url: {}'.format(worker, pdf_url))
domains_to_skip= ['iop.org', 'sciencedirect.com', 'wiley.com', 'acs.org', 'sagepub.com']

is_pdf=False
http_status=None

try:
ifany(domaininpdf_url.urlfordomainindomains_to_skip):
logger.info(f"Skipping {pdf_url.url} because it's in domains to skip")
pdf_url.last_checked=datetime.utcnow()
returnpdf_url

response=http_get(
url=pdf_url.url, ask_slowly=True, stream=True,
publisher=pdf_url.publisher, session_id=get_session_id(),
verify=True
 )
exceptExceptionase:
logger.error("{} failed to get response: {}".format(worker, e))
else:
withresponse:
try:
is_pdf=is_a_pdf_page(response, pdf_url.publisher)
http_status=response.status_code
exceptExceptionase:
logger.error("{} failed reading response: {}".format(worker, e))

pdf_url.is_pdf=is_pdf
pdf_url.http_status=http_status
pdf_url.last_checked=datetime.utcnow()

logger.info('{} updated pdf url: {}'.format(worker, pdf_url))

returnpdf_url


defget_request_pool():
num_request_workers=int(os.getenv('PDF_REQUEST_PROCS_PER_WORKER', 10))
returnPool(processes=num_request_workers, maxtasksperchild=10)


classDbQueuePdfUrlCheck(DbQueue):
deftable_name(self, job_type):
return'pdf_url_check_queue'

defprocess_name(self, job_type):
return'run_pdf_url_check'

defworker_run(self, **kwargs):
run_class=PdfUrl

single_url=kwargs.get("id", None)
chunk_size=kwargs.get("chunk", 100)
limit=kwargs.get("limit", None)

iflimitisNone:
limit=float("inf")

ifsingle_url:
objects= [run_class.query.filter(run_class.url==single_url).first()]
check_pdf_urls(objects)
else:
index=0
num_updated=0
start_time=time()

whilenum_updated<limit:
new_loop_start_time=time()

objects=self.fetch_queue_chunk(chunk_size)

ifnotobjects:
sleep(5)
continue

check_pdf_urls(objects)

object_ids= [obj.urlforobjinobjects]
object_ids_str=",".join(["'{}'".format(oid.replace("'", "''")) foroidinobject_ids])
object_ids_str=object_ids_str.replace("%", "%%") # sql escaping

base_retry_interval=r"(case when pdf_url.url ~* 'academic\.oup\.com/' then interval '5 minutes' else interval '2 hours' end)"
retry_backoff=r"(case when pdf_url.url ~* 'academic\.oup\.com/' then 2 else 4 end)"

sql_command="""
 update {queue_table} q
 set
 finished = now(),
 started = null,
 retry_interval = least(
 case when is_pdf then {base_retry_interval} else {retry_backoff} * coalesce(retry_interval, {base_retry_interval}) end,
 '12 months'
 ),
 retry_at = now() + case when is_pdf then '2 months' else coalesce(retry_interval, {base_retry_interval}) end
 from
 pdf_url
 where
 pdf_url.url = q.url
 and q.url in ({ids})
 """.format(
base_retry_interval=base_retry_interval,
retry_backoff=retry_backoff,
queue_table=self.table_name(None), ids=object_ids_str
 )
run_sql(db, sql_command)

index+=1
num_updated+=len(objects)
self.print_update(new_loop_start_time, chunk_size, limit, start_time, index)

deffetch_queue_chunk(self, chunk_size):
logger.info("looking for new jobs")

text_query_pattern="""
 with update_chunk as (
 select url
 from {queue_table}
 where started is null
 order by retry_at nulls first, finished nulls first, started, rand
 limit {chunk_size}
 for update skip locked
 )
 update {queue_table} queue_rows_to_update
 set started=now()
 from update_chunk
 where update_chunk.url = queue_rows_to_update.url
 returning update_chunk.url;
 """
text_query=text_query_pattern.format(
chunk_size=chunk_size,
queue_table=self.table_name(None)
 )

job_time=time()
row_list=db.engine.execute(text(text_query).execution_options(autocommit=True)).fetchall()
object_ids= [row[0] forrowinrow_list]
logger.info("got ids, took {} seconds".format(elapsed(job_time)))

job_time=time()
q=db.session.query(PdfUrl).options(orm.undefer('*')).filter(PdfUrl.url.in_(object_ids))
objects=q.all()
logger.info("got pdf_url objects in {} seconds".format(elapsed(job_time)))

random.shuffle(objects)
returnobjects


if__name__=="__main__":
ifos.getenv('OADOI_LOG_SQL'):
logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
db.session.configure()

parser=argparse.ArgumentParser(description="Run stuff.")
parser.add_argument('--id', nargs="?", type=str, help="id of the one thing you want to update (case sensitive)")
parser.add_argument('--doi', nargs="?", type=str, help="id of the one thing you want to update (case insensitive)")

parser.add_argument('--reset', default=False, action='store_true', help="do you want to just reset?")
parser.add_argument('--run', default=False, action='store_true', help="to run the queue")
parser.add_argument('--status', default=False, action='store_true', help="to logger.info(the status")
parser.add_argument('--dynos', default=None, type=int, help="scale to this many dynos")
parser.add_argument('--logs', default=False, action='store_true', help="logger.info(out logs")
parser.add_argument('--monitor', default=False, action='store_true', help="monitor till done, then turn off dynos")
parser.add_argument('--kick', default=False, action='store_true', help="put started but unfinished dois back to unstarted so they are retried")
parser.add_argument('--limit', "-l", nargs="?", type=int, help="how many jobs to do")
parser.add_argument('--chunk', "-ch", nargs="?", default=100, type=int, help="how many to take off db at once")

parsed_args=parser.parse_args()

job_type="normal"# should be an object attribute
my_queue=DbQueuePdfUrlCheck()
my_queue.parsed_vars=vars(parsed_args)
my_queue.run_right_thing(parsed_args, job_type)