diff --git a/config/aws/lambda/template_params.yaml b/config/aws/lambda/template_params.yaml index 8e03318d..8c6041d3 100644 --- a/config/aws/lambda/template_params.yaml +++ b/config/aws/lambda/template_params.yaml @@ -38,6 +38,7 @@ Globals: SIDEKICK_PREPROCESS_TOPIC: "{{resolve:ssm:/dc/{$ENV$}/lambdas/sidekick_preprocess/topic:latest}}" RETRY_ERROR_TOPIC: "{{resolve:ssm:/dc/{$ENV$}/lambdas/retry_errors/topic:latest}}" PDF_SIZE_LIMIT: "{{resolve:ssm:/dc/{$ENV$}/params/pdf_size_limit:latest}}" + PAGE_COUNT_LIMIT: "{{resolve:ssm:/dc/{$ENV$}/params/page_count_limit:latest}}" BLOCK_SIZE: "{{resolve:ssm:/dc/{$ENV$}/params/block_size:latest}}" IMPORT_DIR: "{{resolve:ssm:/dc/{$ENV$}/lambdas/config/import_dir:latest}}" IMPORT_BUCKET: "{{resolve:ssm:/dc/{$ENV$}/lambdas/config/import_bucket:latest}}" diff --git a/documentcloud/documents/processing/info_and_image/main.py b/documentcloud/documents/processing/info_and_image/main.py index 6290b426..3e9afd1e 100755 --- a/documentcloud/documents/processing/info_and_image/main.py +++ b/documentcloud/documents/processing/info_and_image/main.py @@ -91,7 +91,12 @@ TEXT_POSITION_BATCH = env.int( "TEXT_POSITION_BATCH", 3 ) # Number of pages to pull text positions from with each function + +# PDF upload size limit is 500MB PDF_SIZE_LIMIT = env.int("PDF_SIZE_LIMIT", 501 * 1024 * 1024) +# PDF page count limit, default 2000 +PAGE_COUNT_LIMIT = env.int("PAGE_COUNT_LIMIT", 2000) + BLOCK_SIZE = env.int( "BLOCK_SIZE", 8 * 1024 * 1024 ) # Block size to use for reading chunks of the PDF @@ -170,6 +175,10 @@ class PdfSizeError(Exception): pass +class PageCountError(Exception): + pass + + def initialize_redis_page_data(doc_id, page_count): """Initialize Redis fields to manage page dimensions and processing""" dimensions_field = redis_fields.dimensions(doc_id) @@ -710,7 +719,13 @@ def process_pdf(data, _context=None): if storage.size(doc_path) > PDF_SIZE_LIMIT: # If not, remove the PDF storage.delete(path.path(doc_id)) - raise PdfSizeError() + # Display 500MB instead of 501MB so -1 + raise PdfSizeError( + f"This document is larger than the " + f"{(PDF_SIZE_LIMIT // (1024 * 1024)) - 1}MB limit. " + f"Please either compress or split the PDF to conform" + f" to the limit." + ) # files are always uploaded to S3 as private, set to public on S3 # if uploaded publicly to DocumentCloud @@ -719,6 +734,19 @@ def process_pdf(data, _context=None): # Extract the page count and store it in Redis page_count = extract_pagecount(doc_id, slug) + + # Ensure page count is within the limit + if page_count > PAGE_COUNT_LIMIT: + # If it is beyond the limit + # delete it from storage + # and raise PageCountError + storage.delete(path.path(doc_id)) + raise PageCountError( + f"This document has {page_count} pages, which exceeds the " + f"{PAGE_COUNT_LIMIT}-page limit. " + f"Please split the PDF before upload." + ) + initialize_redis_page_data(doc_id, page_count) # Update the model with the page count