Merge pull request #559 from christianlouis/copilot/refactor-multi-tenant-polling-engine

feat(tasks): Multi-tenant polling engine for IMAP and watch folder ingestion
This commit is contained in:
Christian Krakau-Louis
2026-03-08 20:23:04 +01:00
committed by GitHub
7 changed files with 840 additions and 47 deletions
+15 -29
View File
@@ -63,15 +63,13 @@ def test_task():
check_credentials.apply_async(countdown=10) # Run 10 seconds after worker starts
celery.conf.beat_schedule = {
"poll-inboxes-every-minute": (
{
"task": "app.tasks.imap_tasks.pull_all_inboxes",
"schedule": crontab(minute="*/1"), # every 1 minute
"options": {"expires": 55}, # Ensure tasks don't pile up
}
if (settings.imap1_host or settings.imap2_host)
else None
),
# IMAP polling — always enabled because per-user IMAP integrations may
# exist in the database even when no system-level IMAP hosts are configured.
"poll-inboxes-every-minute": {
"task": "app.tasks.imap_tasks.pull_all_inboxes",
"schedule": crontab(minute="*/1"), # every 1 minute
"options": {"expires": 55}, # Ensure tasks don't pile up
},
# Add Uptime Kuma ping task if configured
"ping-uptime-kuma": (
{
@@ -100,27 +98,15 @@ celery.conf.beat_schedule = {
"schedule": crontab(minute="*/1"), # Every minute
"options": {"expires": 55}, # Must complete within 55 seconds
},
# Watch folder scanning — polls local paths, FTP, SFTP, and cloud ingest folders.
# Watch folder scanning — always enabled because per-user WATCH_FOLDER
# integrations may exist in the database even when no system-level watch
# folder settings are configured.
# Schedule is controlled by WATCH_FOLDER_POLL_INTERVAL (default: 1 minute).
"scan-watch-folders": (
{
"task": "app.tasks.watch_folder_tasks.scan_all_watch_folders",
"schedule": crontab(minute=f"*/{max(1, settings.watch_folder_poll_interval)}"),
"options": {"expires": 55},
}
if (
settings.watch_folders
or settings.ftp_ingest_enabled
or settings.sftp_ingest_enabled
or settings.dropbox_ingest_enabled
or settings.google_drive_ingest_enabled
or settings.onedrive_ingest_enabled
or settings.nextcloud_ingest_enabled
or settings.s3_ingest_enabled
or settings.webdav_ingest_enabled
)
else None
),
"scan-watch-folders": {
"task": "app.tasks.watch_folder_tasks.scan_all_watch_folders",
"schedule": crontab(minute=f"*/{max(1, settings.watch_folder_poll_interval)}"),
"options": {"expires": 55},
},
# Backfill embeddings for files that were processed before the
# embedding pipeline was enabled, or where the embedding task failed.
"backfill-missing-embeddings": {
+103 -5
View File
@@ -149,6 +149,9 @@ def pull_all_inboxes():
# Per-user IMAP accounts from the database
_pull_user_imap_accounts()
# Per-user IMAP integrations from the UserIntegration model
_pull_user_integration_imap()
logger.info("Finished pull_all_inboxes")
finally:
@@ -176,6 +179,7 @@ def _pull_user_imap_accounts() -> None:
password=_decrypt_imap_password(acct.password),
use_ssl=acct.use_ssl,
delete_after_process=acct.delete_after_process,
owner_id=acct.owner_id,
)
# Record successful poll
acct.last_checked_at = datetime.now(timezone.utc)
@@ -202,6 +206,91 @@ def _pull_user_imap_accounts() -> None:
logger.error("Failed to process per-user IMAP accounts: %s", exc)
def _pull_user_integration_imap() -> None:
"""Iterate over all active IMAP UserIntegrations and pull their inboxes.
This polls the ``user_integrations`` table for records with
``integration_type='IMAP'``, ``direction='SOURCE'``, and ``is_active=True``.
Each integration's config/credentials are decoded and passed to
:func:`pull_inbox` with the owning user's ``owner_id`` so that ingested
documents are correctly attributed.
Individual connection failures are caught and recorded on the integration
without crashing the polling loop.
"""
try:
import json as _json
from app.models import IntegrationDirection, IntegrationType, UserIntegration
from app.utils.encryption import decrypt_value
db = _get_db_session()
try:
integrations = (
db.query(UserIntegration)
.filter(
UserIntegration.integration_type == IntegrationType.IMAP,
UserIntegration.direction == IntegrationDirection.SOURCE,
UserIntegration.is_active.is_(True),
)
.all()
)
logger.info("Processing %d IMAP UserIntegration(s)", len(integrations))
for integ in integrations:
account_identifier = f"integration_{integ.owner_id}_{integ.id}"
try:
cfg = _json.loads(integ.config) if integ.config else {}
raw_creds = decrypt_value(integ.credentials) if integ.credentials else None
creds = _json.loads(raw_creds) if raw_creds else {}
host = cfg.get("host")
port = int(cfg.get("port", 993))
username = cfg.get("username")
password = creds.get("password")
use_ssl = cfg.get("use_ssl", True)
delete_after = cfg.get("delete_after_process", False)
if not (host and username and password):
logger.warning(
"IMAP integration %d (owner %s) has incomplete config — skipping.",
integ.id,
integ.owner_id,
)
continue
pull_inbox(
mailbox_key=account_identifier,
host=host,
port=port,
username=username,
password=password,
use_ssl=use_ssl,
delete_after_process=delete_after,
owner_id=integ.owner_id,
)
integ.last_used_at = datetime.now(timezone.utc)
integ.last_error = None
db.commit()
except Exception as exc: # noqa: BLE001
error_msg = str(exc)[:_MAX_ERROR_LENGTH]
logger.error(
"Error pulling IMAP integration %d (owner %s): %s",
integ.id,
integ.owner_id,
error_msg,
)
try:
integ.last_used_at = datetime.now(timezone.utc)
integ.last_error = error_msg
db.commit()
except Exception: # noqa: BLE001
db.rollback()
finally:
db.close()
except Exception as exc: # noqa: BLE001
logger.error("Failed to process IMAP UserIntegrations: %s", exc)
def check_and_pull_mailbox(
mailbox_key: str,
host: str | None,
@@ -228,7 +317,7 @@ def check_and_pull_mailbox(
)
def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_after_process):
def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_after_process, owner_id=None):
"""
Connects to the IMAP inbox, fetches new unread emails from the last 3 days,
and processes attachments while preserving the original unread status.
@@ -238,6 +327,10 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
- Runs an X-GM-RAW query: "in:anywhere in:unread newer_than:3d has:attachment".
For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter.
Args:
owner_id: Optional user identifier. When provided, ingested documents are
attributed to this user via ``process_document`` / ``convert_to_pdf``.
"""
logger.info("Connecting to %s at %s:%s (SSL=%s)", mailbox_key, host, port, use_ssl)
processed_emails = load_processed_emails()
@@ -300,7 +393,7 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
# Process attachments (and convert non-PDF files).
# We call the function without assigning its return value since it is not used.
fetch_attachments_and_enqueue(email_message)
fetch_attachments_and_enqueue(email_message, owner_id=owner_id)
if settings.imap_readonly_mode:
logger.info("Readonly mode: skipping mailbox modifications for %s in %s", msg_id, mailbox_key)
@@ -329,7 +422,7 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
logger.exception("Error pulling mailbox %s: %s", mailbox_key, e)
def fetch_attachments_and_enqueue(email_message):
def fetch_attachments_and_enqueue(email_message, owner_id: str | None = None):
"""
Extracts attachments from the email and processes only allowed file types.
@@ -354,6 +447,11 @@ def fetch_attachments_and_enqueue(email_message):
If the attachment is a PDF (by extension or MIME type), it is enqueued for upload;
any other allowed file is enqueued for conversion to PDF.
Args:
email_message: The parsed email message to extract attachments from.
owner_id: Optional user identifier forwarded to ``process_document`` /
``convert_to_pdf`` for multi-tenant attribution.
Returns True if at least one allowed attachment was processed.
"""
has_attachment = False
@@ -381,11 +479,11 @@ def fetch_attachments_and_enqueue(email_message):
# If it's a PDF by MIME type or extension, process it directly
if mime_type == "application/pdf" or is_pdf_by_extension:
process_document.delay(file_path)
process_document.delay(file_path, owner_id=owner_id)
logger.info("Enqueued PDF for upload: %s (MIME: %s)", filename, mime_type)
elif mime_type in ALLOWED_MIME_TYPES:
# Other allowed files are sent for conversion
convert_to_pdf.delay(file_path)
convert_to_pdf.delay(file_path, owner_id=owner_id)
logger.info("Enqueued file for conversion to PDF: %s", filename)
has_attachment = True
+242 -4
View File
@@ -39,8 +39,27 @@ FTP_INGEST_CACHE_FILE = os.path.join(settings.workdir, "ftp_ingest_processed.jso
# Cache file for tracking already-ingested files (SFTP watch folder)
SFTP_INGEST_CACHE_FILE = os.path.join(settings.workdir, "sftp_ingest_processed.json")
# Per-integration watch-folder cache file prefix
_USER_WF_CACHE_PREFIX = os.path.join(settings.workdir, "user_wf_")
_CACHE_RETENTION_DAYS = 30
# Maximum length to store as last_error on UserIntegration to prevent DB bloat
_MAX_ERROR_LENGTH = 500
# Database session factory (imported lazily to avoid circular imports)
_db_session_factory = None
def _get_db_session():
"""Return a new SQLAlchemy session (lazy import to avoid startup issues)."""
global _db_session_factory # noqa: PLW0603
if _db_session_factory is None:
from app.database import SessionLocal
_db_session_factory = SessionLocal
return _db_session_factory()
# ---------------------------------------------------------------------------
# Locking helpers
@@ -122,17 +141,24 @@ def _is_allowed_file(filename: str) -> bool:
return ext in ALLOWED_EXTENSIONS or filename.lower().endswith(".pdf")
def _enqueue_file(file_path: str, *, filename: str | None = None) -> None:
"""Enqueue a local file path for document processing."""
def _enqueue_file(file_path: str, *, filename: str | None = None, owner_id: str | None = None) -> None:
"""Enqueue a local file path for document processing.
Args:
file_path: Absolute path to the file on disk.
filename: Optional display filename (defaults to basename of *file_path*).
owner_id: Optional user identifier forwarded to ``process_document`` /
``convert_to_pdf`` for multi-tenant attribution.
"""
fname = filename or os.path.basename(file_path)
_, ext = os.path.splitext(fname)
mime_check = ext.lower() in {".pdf"}
if mime_check or fname.lower().endswith(".pdf"):
process_document.delay(file_path)
process_document.delay(file_path, owner_id=owner_id)
logger.info("Enqueued for processing: %s", fname)
else:
convert_to_pdf.delay(file_path)
convert_to_pdf.delay(file_path, owner_id=owner_id)
logger.info("Enqueued for PDF conversion: %s", fname)
@@ -1288,6 +1314,216 @@ def scan_webdav_watch_folder() -> dict:
return {"status": "ok", "files_enqueued": n, "folder": ingest_folder}
# ---------------------------------------------------------------------------
# Per-user watch folder integration scanning
# ---------------------------------------------------------------------------
def _is_safe_watch_path(folder_path: str) -> bool:
"""Validate that a user-configured watch folder path is safe.
Rejects paths that attempt directory traversal (``..``), use relative
references, or are not absolute. This prevents a malicious user from
configuring a watch folder that could escape its intended directory.
Args:
folder_path: The path to validate.
Returns:
``True`` if the path is considered safe, ``False`` otherwise.
"""
if not folder_path:
return False
# Must be absolute
if not os.path.isabs(folder_path):
logger.warning("Rejecting non-absolute watch folder path: %s", folder_path)
return False
# Resolve to canonical path and ensure no traversal components exist
resolved = os.path.realpath(folder_path)
if ".." in folder_path.split(os.sep):
logger.warning("Rejecting path with traversal components: %s", folder_path)
return False
# Ensure resolved path matches the original intent (no symlink escapes)
if resolved != os.path.normpath(folder_path):
logger.warning(
"Watch folder path resolves differently (possible symlink escape): %s -> %s",
folder_path,
resolved,
)
return False
return True
def _scan_user_watch_folder(
folder_path: str,
cache: dict[str, str],
delete_after: bool,
owner_id: str,
) -> int:
"""Scan a user-configured local watch folder, attributing files to *owner_id*.
Delegates to the same file-scanning logic as :func:`_scan_local_folder` but
passes ``owner_id`` to :func:`_enqueue_file` so that ingested documents are
correctly attributed to the user.
Args:
folder_path: Absolute directory path to scan.
cache: In-memory dict of already-processed file keys.
delete_after: Whether to remove the source file after ingestion.
owner_id: The user to attribute ingested documents to.
Returns:
Number of files newly enqueued.
"""
if not os.path.isdir(folder_path):
logger.warning("User watch folder does not exist or is not a directory: %s", folder_path)
return 0
count = 0
try:
entries = os.scandir(folder_path)
except PermissionError as exc:
logger.error("Cannot scan user watch folder %s: %s", folder_path, exc)
return 0
for entry in entries:
if not entry.is_file(follow_symlinks=True):
continue
if not _is_allowed_file(entry.name):
continue
abs_path = entry.path
if abs_path in cache:
continue
dest_filename = f"uwf_{owner_id}_{entry.name}"
dest_path = os.path.join(settings.workdir, dest_filename)
if os.path.exists(dest_path):
base, ext2 = os.path.splitext(dest_filename)
dest_path = os.path.join(settings.workdir, f"{base}_{int(datetime.now().timestamp())}{ext2}")
try:
import shutil
shutil.copy2(abs_path, dest_path)
except OSError as exc:
logger.error("Failed to copy %s to workdir: %s", abs_path, exc)
continue
_enqueue_file(dest_path, owner_id=owner_id)
_mark_processed(cache, abs_path)
count += 1
if delete_after:
try:
os.remove(abs_path)
logger.info("Deleted source file after ingestion: %s", abs_path)
except OSError as exc:
logger.warning("Could not delete source file %s: %s", abs_path, exc)
return count
def _pull_user_integration_watch_folders() -> dict:
"""Iterate over all active WATCH_FOLDER UserIntegrations and scan their paths.
Polls the ``user_integrations`` table for records with
``integration_type='WATCH_FOLDER'``, ``direction='SOURCE'``, and
``is_active=True``. Each integration's config is decoded and the
configured ``folder_path`` is scanned for new files, which are enqueued
with the owning user's ``owner_id``.
Path traversal protection is enforced on the configured path.
Individual integration failures are caught and recorded without crashing
the polling loop.
Returns:
Summary dict with ``status`` and ``integrations_processed`` count.
"""
try:
import json as _json
from app.models import IntegrationDirection, IntegrationType, UserIntegration
db = _get_db_session()
try:
integrations = (
db.query(UserIntegration)
.filter(
UserIntegration.integration_type == IntegrationType.WATCH_FOLDER,
UserIntegration.direction == IntegrationDirection.SOURCE,
UserIntegration.is_active.is_(True),
)
.all()
)
logger.info("Processing %d WATCH_FOLDER UserIntegration(s)", len(integrations))
total_files = 0
for integ in integrations:
try:
cfg = _json.loads(integ.config) if integ.config else {}
folder_path = cfg.get("folder_path", "")
delete_after = cfg.get("delete_after_process", False)
if not folder_path:
logger.warning(
"Watch folder integration %d (owner %s) has no folder_path — skipping.",
integ.id,
integ.owner_id,
)
continue
if not _is_safe_watch_path(folder_path):
error_msg = f"Unsafe watch folder path rejected: {folder_path}"
logger.error(
"Watch folder integration %d (owner %s): %s",
integ.id,
integ.owner_id,
error_msg,
)
integ.last_error = error_msg[:_MAX_ERROR_LENGTH]
db.commit()
continue
cache_file = f"{_USER_WF_CACHE_PREFIX}{integ.id}.json"
cache = _load_cache(cache_file)
n = _scan_user_watch_folder(folder_path, cache, delete_after, integ.owner_id)
_save_cache(cache_file, cache)
total_files += n
integ.last_used_at = datetime.now(timezone.utc)
integ.last_error = None
db.commit()
logger.info(
"Watch folder integration %d (owner %s): %d file(s) enqueued from %s",
integ.id,
integ.owner_id,
n,
folder_path,
)
except Exception as exc: # noqa: BLE001
error_msg = str(exc)[:_MAX_ERROR_LENGTH]
logger.error(
"Error scanning watch folder integration %d (owner %s): %s",
integ.id,
integ.owner_id,
error_msg,
)
try:
integ.last_used_at = datetime.now(timezone.utc)
integ.last_error = error_msg
db.commit()
except Exception: # noqa: BLE001
db.rollback()
return {"status": "ok", "integrations_processed": len(integrations), "files_enqueued": total_files}
finally:
db.close()
except Exception as exc: # noqa: BLE001
logger.error("Failed to process WATCH_FOLDER UserIntegrations: %s", exc)
return {"status": "error", "reason": str(exc)[:_MAX_ERROR_LENGTH]}
@shared_task
def scan_all_watch_folders() -> dict:
"""
@@ -1303,6 +1539,7 @@ def scan_all_watch_folders() -> dict:
7. Nextcloud ingest folder (if enabled)
8. Amazon S3 ingest prefix (if enabled)
9. WebDAV ingest folder (if enabled)
10. Per-user WATCH_FOLDER integrations from the database
"""
if not _acquire_lock(WATCH_FOLDER_LOCK_KEY):
logger.info("Watch folder scan already running — skipping this cycle.")
@@ -1319,6 +1556,7 @@ def scan_all_watch_folders() -> dict:
results["nextcloud"] = scan_nextcloud_watch_folder()
results["s3"] = scan_s3_watch_folder()
results["webdav"] = scan_webdav_watch_folder()
results["user_watch_folders"] = _pull_user_integration_watch_folders()
finally:
_release_lock(WATCH_FOLDER_LOCK_KEY)
+30 -3
View File
@@ -267,12 +267,27 @@ DocuElevate can poll a WebDAV folder for new files. It reuses the existing WebDA
| `WEBDAV_INGEST_FOLDER` | WebDAV folder path to poll. Uses the existing WebDAV URL and credentials. | *(empty)* |
| `WEBDAV_INGEST_DELETE_AFTER_PROCESS` | Delete files from WebDAV after they are downloaded and enqueued. | `false` |
#### Per-User Watch Folder Integrations
In addition to system-level watch folders, each user can configure personal watch folder sources through the **Integrations** dashboard (`/integrations`). Documents ingested from per-user watch folder integrations are automatically attributed to the owning user's `owner_id`.
Per-user watch folder integrations are stored in the `user_integrations` table with `integration_type='WATCH_FOLDER'` and `direction='SOURCE'`. The `config` JSON field stores:
- `folder_path` — absolute path to the directory to scan
- `delete_after_process` — whether to remove source files after ingestion (default: `false`)
> **Security**: Path traversal protection is enforced on user-configured watch folder paths. Relative paths, `..` components, and symlink escapes are rejected to prevent access to files outside the intended directory.
- Individual scan failures are handled gracefully and recorded on the integration's `last_error` field without interrupting the scanning of other integrations.
- The scan runs alongside the system-level watch folder polling cycle.
### IMAP Email Ingestion
DocuElevate can automatically pull document attachments from IMAP mailboxes — no need to forward emails manually. Configure one or two mailboxes and DocuElevate polls them on the schedule you set.
DocuElevate can automatically pull document attachments from IMAP mailboxes — no need to forward emails manually. Configure one or two system-wide mailboxes using environment variables, and/or let each user configure their own IMAP sources via the **Integrations** dashboard.
> **For HP Scanners (Scan to Email)**: If your scanner is set up to email scanned documents to a dedicated mailbox, configure that mailbox in DocuElevate using the settings below. DocuElevate will automatically retrieve the scanned PDFs from the inbox and process them. You do **not** need to configure DocuElevate as an email server — it acts as an email *client* that reads from your existing mailbox.
#### System-Level IMAP Configuration
| **Variable** | **Description** | **Example** |
|-------------------------------|--------------------------------------------------------------|-------------------|
| `IMAP1_HOST` | Hostname for first IMAP server. | `mail.example.com`|
@@ -283,6 +298,15 @@ DocuElevate can automatically pull document attachments from IMAP mailboxes —
| `IMAP1_POLL_INTERVAL_MINUTES` | Frequency in minutes to poll for new mail. | `5` |
| `IMAP_READONLY_MODE` | When `true`, fetches and processes attachments but does **not** modify the mailbox (no starring, labeling, deleting, or flag changes). Use for pre-production instances sharing a mailbox with production. Default: `false`. | `false` |
#### Per-User IMAP Integrations
In addition to system-level mailboxes, each user can configure personal IMAP sources through the **Integrations** dashboard (`/integrations`). Documents ingested from per-user IMAP integrations are automatically attributed to the owning user's `owner_id`.
Per-user IMAP integrations are stored in the `user_integrations` table with `integration_type='IMAP'` and `direction='SOURCE'`. Credentials are encrypted at rest using Fernet encryption.
- Individual connection failures are handled gracefully and recorded on the integration's `last_error` field without interrupting the polling of other integrations.
- The polling loop runs every minute and processes all active IMAP sources (system-level and per-user) in sequence.
### Authentication
| **Variable** | **Description** |
@@ -316,8 +340,11 @@ Requires `AUTH_ENABLED=true`.
#### Unclaimed Documents
Documents ingested without a user session (e.g. via IMAP polling, API calls without authentication,
or legacy imports) have `owner_id = NULL`. These are called **unclaimed** documents.
Documents ingested via **system-level** sources (environment variable IMAP mailboxes, system watch folders)
without a user session have `owner_id = NULL` unless `DEFAULT_OWNER_ID` is set. These are called **unclaimed** documents.
Documents ingested via **per-user integrations** (IMAP or Watch Folder integrations configured through
the Integrations dashboard) are automatically attributed to the owning user's `owner_id` and are never unclaimed.
- When `UNOWNED_DOCS_VISIBLE_TO_ALL=true` (default), every authenticated user sees unclaimed
documents alongside their own files. This allows users to discover and claim them.
+3 -3
View File
@@ -163,10 +163,10 @@ class TestBeatScheduleConfiguration:
# -- Conditional schedule entries ----------------------------------------
def test_imap_schedule_absent_when_not_configured(self):
"""Test IMAP polling absent when neither imap host is set."""
def test_imap_schedule_always_present(self):
"""Test IMAP polling is always scheduled (user integrations may exist)."""
mod = _reload_celery_worker()
assert "poll-inboxes-every-minute" not in mod.celery.conf.beat_schedule
assert "poll-inboxes-every-minute" in mod.celery.conf.beat_schedule
def test_uptime_kuma_schedule_absent_when_not_configured(self):
"""Test Uptime Kuma ping absent when url is not set."""
+194
View File
@@ -1354,3 +1354,197 @@ class TestEmailAlreadyHasLabelExceptions:
# Should return False on error
assert result is False
# ---------------------------------------------------------------------------
# Tests for multi-tenant IMAP user integration polling
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestFetchAttachmentsOwnerIdPassthrough:
"""Test that fetch_attachments_and_enqueue forwards owner_id correctly."""
@patch("app.tasks.imap_tasks.process_document")
@patch("app.tasks.imap_tasks.convert_to_pdf")
def test_pdf_attachment_forwards_owner_id(self, mock_convert, mock_process, tmp_path):
"""PDF attachments should forward owner_id to process_document."""
msg = EmailMessage()
msg["Subject"] = "Test"
msg.add_attachment(b"%PDF-1.4", maintype="application", subtype="pdf", filename="doc.pdf")
with patch("app.tasks.imap_tasks.settings") as mock_settings:
mock_settings.workdir = str(tmp_path)
result = fetch_attachments_and_enqueue(msg, owner_id="user-42")
assert result is True
mock_process.delay.assert_called_once()
call_kwargs = mock_process.delay.call_args
assert call_kwargs.kwargs.get("owner_id") == "user-42"
@patch("app.tasks.imap_tasks.process_document")
@patch("app.tasks.imap_tasks.convert_to_pdf")
def test_non_pdf_attachment_forwards_owner_id(self, mock_convert, mock_process, tmp_path):
"""Non-PDF attachments should forward owner_id to convert_to_pdf."""
msg = EmailMessage()
msg["Subject"] = "Test"
msg.add_attachment(
b"excel content",
maintype="application",
subtype="vnd.ms-excel",
filename="spreadsheet.xls",
)
with patch("app.tasks.imap_tasks.settings") as mock_settings:
mock_settings.workdir = str(tmp_path)
result = fetch_attachments_and_enqueue(msg, owner_id="user-99")
assert result is True
mock_convert.delay.assert_called_once()
call_kwargs = mock_convert.delay.call_args
assert call_kwargs.kwargs.get("owner_id") == "user-99"
@patch("app.tasks.imap_tasks.process_document")
@patch("app.tasks.imap_tasks.convert_to_pdf")
def test_no_owner_id_defaults_to_none(self, mock_convert, mock_process, tmp_path):
"""When owner_id is not provided, it defaults to None."""
msg = EmailMessage()
msg["Subject"] = "Test"
msg.add_attachment(b"%PDF-1.4", maintype="application", subtype="pdf", filename="doc.pdf")
with patch("app.tasks.imap_tasks.settings") as mock_settings:
mock_settings.workdir = str(tmp_path)
fetch_attachments_and_enqueue(msg)
call_kwargs = mock_process.delay.call_args
assert call_kwargs.kwargs.get("owner_id") is None
@pytest.mark.unit
class TestPullUserIntegrationImap:
"""Tests for _pull_user_integration_imap() multi-tenant IMAP polling."""
@patch("app.tasks.imap_tasks._get_db_session")
@patch("app.tasks.imap_tasks.pull_inbox")
def test_polls_active_imap_integrations(self, mock_pull, mock_session_factory):
"""Active IMAP integrations should be polled with correct owner_id."""
from app.tasks.imap_tasks import _pull_user_integration_imap
mock_integ = MagicMock()
mock_integ.id = 10
mock_integ.owner_id = "owner-abc"
mock_integ.config = '{"host": "imap.example.com", "port": 993, "username": "user@test.com", "use_ssl": true}'
mock_integ.credentials = "enc:encrypted"
mock_integ.is_active = True
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
with patch("app.utils.encryption.decrypt_value", return_value='{"password": "secret"}'):
_pull_user_integration_imap()
mock_pull.assert_called_once()
call_kwargs = mock_pull.call_args
assert call_kwargs.kwargs.get("owner_id") == "owner-abc"
assert call_kwargs.kwargs.get("host") == "imap.example.com"
assert call_kwargs.kwargs.get("password") == "secret" # noqa: S105
@patch("app.tasks.imap_tasks._get_db_session")
@patch("app.tasks.imap_tasks.pull_inbox")
def test_skips_incomplete_config(self, mock_pull, mock_session_factory):
"""Integrations missing host/username/password should be skipped."""
from app.tasks.imap_tasks import _pull_user_integration_imap
mock_integ = MagicMock()
mock_integ.id = 11
mock_integ.owner_id = "owner-xyz"
mock_integ.config = '{"host": "", "port": 993, "username": ""}'
mock_integ.credentials = None
mock_integ.is_active = True
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
_pull_user_integration_imap()
mock_pull.assert_not_called()
@patch("app.tasks.imap_tasks._get_db_session")
@patch("app.tasks.imap_tasks.pull_inbox")
def test_handles_connection_failure_gracefully(self, mock_pull, mock_session_factory):
"""Connection failures should be recorded but not crash the loop."""
from app.tasks.imap_tasks import _pull_user_integration_imap
mock_integ = MagicMock()
mock_integ.id = 12
mock_integ.owner_id = "owner-fail"
mock_integ.config = '{"host": "bad.host", "port": 993, "username": "u@x.com", "use_ssl": true}'
mock_integ.credentials = "enc:encrypted"
mock_integ.is_active = True
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
mock_pull.side_effect = Exception("Connection refused")
with patch("app.utils.encryption.decrypt_value", return_value='{"password": "p"}'):
# Should not raise
_pull_user_integration_imap()
# last_error should be recorded
assert mock_integ.last_error is not None
assert "Connection refused" in mock_integ.last_error
mock_db.commit.assert_called()
@patch("app.tasks.imap_tasks._get_db_session")
def test_handles_db_failure_gracefully(self, mock_session_factory):
"""Database failures should be caught without crashing."""
from app.tasks.imap_tasks import _pull_user_integration_imap
mock_session_factory.side_effect = Exception("DB unavailable")
# Should not raise
_pull_user_integration_imap()
@patch("app.tasks.imap_tasks._get_db_session")
@patch("app.tasks.imap_tasks.pull_inbox")
def test_records_last_used_at_on_success(self, mock_pull, mock_session_factory):
"""Successful polling should update last_used_at and clear last_error."""
from app.tasks.imap_tasks import _pull_user_integration_imap
mock_integ = MagicMock()
mock_integ.id = 13
mock_integ.owner_id = "owner-ok"
mock_integ.config = '{"host": "imap.ok.com", "port": 993, "username": "ok@ok.com", "use_ssl": true}'
mock_integ.credentials = "enc:encrypted"
mock_integ.last_error = "previous error"
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
with patch("app.utils.encryption.decrypt_value", return_value='{"password": "ok"}'):
_pull_user_integration_imap()
assert mock_integ.last_error is None
assert mock_integ.last_used_at is not None
@pytest.mark.unit
class TestPullAllInboxesCallsIntegrations:
"""Test that pull_all_inboxes calls both legacy and new integration polling."""
@patch("app.tasks.imap_tasks._pull_user_integration_imap")
@patch("app.tasks.imap_tasks._pull_user_imap_accounts")
@patch("app.tasks.imap_tasks.check_and_pull_mailbox")
@patch("app.tasks.imap_tasks.acquire_lock", return_value=True)
@patch("app.tasks.imap_tasks.release_lock")
def test_calls_both_legacy_and_integration_polling(
self, mock_release, mock_lock, mock_check, mock_legacy, mock_integ
):
"""pull_all_inboxes should call both _pull_user_imap_accounts and _pull_user_integration_imap."""
pull_all_inboxes()
mock_legacy.assert_called_once()
mock_integ.assert_called_once()
+253 -3
View File
@@ -834,7 +834,7 @@ class TestEnqueueFileNonPdf:
patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
):
_enqueue_file("/tmp/doc.docx")
mock_conv.delay.assert_called_once_with("/tmp/doc.docx")
mock_conv.delay.assert_called_once_with("/tmp/doc.docx", owner_id=None)
mock_proc.delay.assert_not_called()
def test_pdf_triggers_process_document(self):
@@ -846,7 +846,7 @@ class TestEnqueueFileNonPdf:
patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
):
_enqueue_file("/tmp/report.pdf")
mock_proc.delay.assert_called_once_with("/tmp/report.pdf")
mock_proc.delay.assert_called_once_with("/tmp/report.pdf", owner_id=None)
mock_conv.delay.assert_not_called()
def test_pdf_detected_via_filename_param(self):
@@ -858,7 +858,7 @@ class TestEnqueueFileNonPdf:
patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
):
_enqueue_file("/tmp/somefile", filename="renamed.pdf")
mock_proc.delay.assert_called_once_with("/tmp/somefile")
mock_proc.delay.assert_called_once_with("/tmp/somefile", owner_id=None)
mock_conv.delay.assert_not_called()
@@ -3622,3 +3622,253 @@ class TestS3DownloadFailureNoPartialFile:
mock_settings.workdir = str(tmp_path)
count = _scan_s3_prefix("inbox/", cache, False)
assert count == 0
# ---------------------------------------------------------------------------
# Tests for multi-tenant watch folder integration polling
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestIsSafeWatchPath:
"""Tests for _is_safe_watch_path path traversal security."""
def test_absolute_path_is_safe(self):
from app.tasks.watch_folder_tasks import _is_safe_watch_path
assert _is_safe_watch_path("/data/watch") is True
def test_empty_path_is_unsafe(self):
from app.tasks.watch_folder_tasks import _is_safe_watch_path
assert _is_safe_watch_path("") is False
def test_relative_path_is_unsafe(self):
from app.tasks.watch_folder_tasks import _is_safe_watch_path
assert _is_safe_watch_path("relative/path") is False
def test_traversal_path_is_unsafe(self):
from app.tasks.watch_folder_tasks import _is_safe_watch_path
assert _is_safe_watch_path("/data/../etc/passwd") is False
def test_double_dot_component_is_unsafe(self):
from app.tasks.watch_folder_tasks import _is_safe_watch_path
assert _is_safe_watch_path("/data/watch/../../secret") is False
def test_none_path_is_unsafe(self):
from app.tasks.watch_folder_tasks import _is_safe_watch_path
assert _is_safe_watch_path(None) is False
@pytest.mark.unit
class TestEnqueueFileOwnerIdPassthrough:
"""Test that _enqueue_file forwards owner_id to downstream tasks."""
def test_pdf_forwards_owner_id(self):
from app.tasks.watch_folder_tasks import _enqueue_file
with (
patch("app.tasks.watch_folder_tasks.convert_to_pdf") as mock_conv,
patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
):
_enqueue_file("/tmp/report.pdf", owner_id="user-123")
mock_proc.delay.assert_called_once_with("/tmp/report.pdf", owner_id="user-123")
mock_conv.delay.assert_not_called()
def test_non_pdf_forwards_owner_id(self):
from app.tasks.watch_folder_tasks import _enqueue_file
with (
patch("app.tasks.watch_folder_tasks.convert_to_pdf") as mock_conv,
patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
):
_enqueue_file("/tmp/doc.docx", owner_id="user-456")
mock_conv.delay.assert_called_once_with("/tmp/doc.docx", owner_id="user-456")
mock_proc.delay.assert_not_called()
@pytest.mark.unit
class TestScanUserWatchFolder:
"""Tests for _scan_user_watch_folder."""
def test_scans_and_attributes_files_to_owner(self, tmp_path):
from app.tasks.watch_folder_tasks import _scan_user_watch_folder
# Create a test file
(tmp_path / "test.pdf").write_bytes(b"%PDF-1.4")
cache: dict[str, str] = {}
with (
patch("app.tasks.watch_folder_tasks._enqueue_file") as mock_enqueue,
patch("app.tasks.watch_folder_tasks.settings") as mock_settings,
):
mock_settings.workdir = str(tmp_path / "workdir")
os.makedirs(mock_settings.workdir, exist_ok=True)
count = _scan_user_watch_folder(str(tmp_path), cache, False, "owner-77")
assert count == 1
mock_enqueue.assert_called_once()
call_kwargs = mock_enqueue.call_args
assert call_kwargs.kwargs.get("owner_id") == "owner-77"
def test_skips_already_processed_files(self, tmp_path):
from app.tasks.watch_folder_tasks import _scan_user_watch_folder
test_file = tmp_path / "test.pdf"
test_file.write_bytes(b"%PDF-1.4")
cache = {str(test_file): datetime.now(timezone.utc).isoformat()}
with (
patch("app.tasks.watch_folder_tasks._enqueue_file") as mock_enqueue,
patch("app.tasks.watch_folder_tasks.settings") as mock_settings,
):
mock_settings.workdir = str(tmp_path / "workdir")
count = _scan_user_watch_folder(str(tmp_path), cache, False, "owner-77")
assert count == 0
mock_enqueue.assert_not_called()
def test_returns_zero_for_nonexistent_dir(self):
from app.tasks.watch_folder_tasks import _scan_user_watch_folder
count = _scan_user_watch_folder("/nonexistent/path", {}, False, "owner-1")
assert count == 0
@pytest.mark.unit
class TestPullUserIntegrationWatchFolders:
"""Tests for _pull_user_integration_watch_folders."""
@patch("app.tasks.watch_folder_tasks._get_db_session")
@patch("app.tasks.watch_folder_tasks._scan_user_watch_folder", return_value=2)
@patch("app.tasks.watch_folder_tasks._load_cache", return_value={})
@patch("app.tasks.watch_folder_tasks._save_cache")
def test_polls_active_watch_folder_integrations(self, mock_save, mock_load, mock_scan, mock_session_factory):
"""Active WATCH_FOLDER integrations should be scanned."""
from app.tasks.watch_folder_tasks import _pull_user_integration_watch_folders
mock_integ = MagicMock()
mock_integ.id = 20
mock_integ.owner_id = "owner-wf"
mock_integ.config = '{"folder_path": "/data/scans", "delete_after_process": false}'
mock_integ.is_active = True
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
with patch("app.tasks.watch_folder_tasks._is_safe_watch_path", return_value=True):
result = _pull_user_integration_watch_folders()
assert result["status"] == "ok"
assert result["integrations_processed"] == 1
assert result["files_enqueued"] == 2
mock_scan.assert_called_once()
@patch("app.tasks.watch_folder_tasks._get_db_session")
def test_rejects_unsafe_paths(self, mock_session_factory):
"""Integrations with unsafe paths should be rejected."""
from app.tasks.watch_folder_tasks import _pull_user_integration_watch_folders
mock_integ = MagicMock()
mock_integ.id = 21
mock_integ.owner_id = "owner-bad"
mock_integ.config = '{"folder_path": "/data/../etc/passwd"}'
mock_integ.is_active = True
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
with patch("app.tasks.watch_folder_tasks._is_safe_watch_path", return_value=False):
result = _pull_user_integration_watch_folders()
assert result["status"] == "ok"
assert mock_integ.last_error is not None
@patch("app.tasks.watch_folder_tasks._get_db_session")
@patch("app.tasks.watch_folder_tasks._scan_user_watch_folder")
@patch("app.tasks.watch_folder_tasks._load_cache", return_value={})
@patch("app.tasks.watch_folder_tasks._save_cache")
def test_handles_scan_failure_gracefully(self, mock_save, mock_load, mock_scan, mock_session_factory):
"""Scan failures should be recorded but not crash the loop."""
from app.tasks.watch_folder_tasks import _pull_user_integration_watch_folders
mock_integ = MagicMock()
mock_integ.id = 22
mock_integ.owner_id = "owner-err"
mock_integ.config = '{"folder_path": "/data/broken"}'
mock_integ.is_active = True
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
mock_scan.side_effect = Exception("Disk error")
with patch("app.tasks.watch_folder_tasks._is_safe_watch_path", return_value=True):
result = _pull_user_integration_watch_folders()
assert mock_integ.last_error is not None
assert "Disk error" in mock_integ.last_error
@patch("app.tasks.watch_folder_tasks._get_db_session")
def test_handles_db_failure_gracefully(self, mock_session_factory):
"""Database failures should return error status without crashing."""
from app.tasks.watch_folder_tasks import _pull_user_integration_watch_folders
mock_session_factory.side_effect = Exception("DB unavailable")
result = _pull_user_integration_watch_folders()
assert result["status"] == "error"
@patch("app.tasks.watch_folder_tasks._get_db_session")
def test_skips_integration_without_folder_path(self, mock_session_factory):
"""Integrations without folder_path should be skipped."""
from app.tasks.watch_folder_tasks import _pull_user_integration_watch_folders
mock_integ = MagicMock()
mock_integ.id = 23
mock_integ.owner_id = "owner-empty"
mock_integ.config = '{"delete_after_process": false}'
mock_integ.is_active = True
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.all.return_value = [mock_integ]
mock_session_factory.return_value = mock_db
result = _pull_user_integration_watch_folders()
assert result["status"] == "ok"
@pytest.mark.unit
class TestScanAllWatchFoldersIncludesUserIntegrations:
"""Test that scan_all_watch_folders calls _pull_user_integration_watch_folders."""
def test_includes_user_watch_folders(self):
from app.tasks.watch_folder_tasks import scan_all_watch_folders
with (
patch("app.tasks.watch_folder_tasks._acquire_lock", return_value=True),
patch("app.tasks.watch_folder_tasks._release_lock"),
patch("app.tasks.watch_folder_tasks.scan_local_watch_folders", return_value={"status": "ok"}),
patch("app.tasks.watch_folder_tasks.scan_ftp_watch_folder", return_value={"status": "skipped"}),
patch("app.tasks.watch_folder_tasks.scan_sftp_watch_folder", return_value={"status": "skipped"}),
patch("app.tasks.watch_folder_tasks.scan_dropbox_watch_folder", return_value={"status": "skipped"}),
patch("app.tasks.watch_folder_tasks.scan_google_drive_watch_folder", return_value={"status": "skipped"}),
patch("app.tasks.watch_folder_tasks.scan_onedrive_watch_folder", return_value={"status": "skipped"}),
patch("app.tasks.watch_folder_tasks.scan_nextcloud_watch_folder", return_value={"status": "skipped"}),
patch("app.tasks.watch_folder_tasks.scan_s3_watch_folder", return_value={"status": "skipped"}),
patch("app.tasks.watch_folder_tasks.scan_webdav_watch_folder", return_value={"status": "skipped"}),
patch(
"app.tasks.watch_folder_tasks._pull_user_integration_watch_folders",
return_value={"status": "ok", "integrations_processed": 1, "files_enqueued": 3},
) as mock_pull,
):
result = scan_all_watch_folders()
mock_pull.assert_called_once()
assert result["results"]["user_watch_folders"]["files_enqueued"] == 3