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..19a19e44 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 +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) @@ -719,6 +728,15 @@ 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() + initialize_redis_page_data(doc_id, page_count) # Update the model with the page count