wehub-resource-sync
fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
1048 lines
45 KiB
Python
1048 lines
45 KiB
Python
"""Source document management upload functionality."""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import uuid
|
|
import zipfile
|
|
|
|
from flask import current_app, jsonify, make_response, request
|
|
from flask_restx import fields, Namespace, Resource
|
|
from sqlalchemy import text as sql_text
|
|
|
|
from application.api import api
|
|
from application.api.user.tasks import ingest, ingest_connector_task, ingest_remote
|
|
from application.api.user.team_sharing import effective_write_owner
|
|
from application.core.settings import settings
|
|
from application.storage.db.source_ids import derive_source_id as _derive_source_id
|
|
from application.parser.connectors.connector_creator import ConnectorCreator
|
|
from application.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS
|
|
from application.storage.db.repositories.idempotency import IdempotencyRepository
|
|
from application.storage.db.repositories.sources import SourcesRepository
|
|
from application.storage.db.source_config import SourceConfig
|
|
from application.storage.db.session import db_readonly, db_session
|
|
from application.storage.storage_creator import StorageCreator
|
|
from application.stt.upload_limits import (
|
|
AudioFileTooLargeError,
|
|
build_stt_file_size_limit_message,
|
|
enforce_audio_file_size_limit,
|
|
is_audio_filename,
|
|
)
|
|
from application.utils import check_required_fields, safe_filename
|
|
|
|
|
|
sources_upload_ns = Namespace(
|
|
"sources", description="Source document management operations", path="/api"
|
|
)
|
|
|
|
|
|
_IDEMPOTENCY_KEY_MAX_LEN = 256
|
|
|
|
|
|
def _read_idempotency_key():
|
|
"""Return (key, error_response). Empty header → (None, None); oversized → (None, 400)."""
|
|
key = request.headers.get("Idempotency-Key")
|
|
if not key:
|
|
return None, None
|
|
if len(key) > _IDEMPOTENCY_KEY_MAX_LEN:
|
|
return None, make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": (
|
|
f"Idempotency-Key exceeds maximum length of "
|
|
f"{_IDEMPOTENCY_KEY_MAX_LEN} characters"
|
|
),
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
return key, None
|
|
|
|
|
|
def _scoped_idempotency_key(idempotency_key, scope):
|
|
"""``{scope}:{key}`` so different users can't collide on the same key."""
|
|
if not idempotency_key or not scope:
|
|
return None
|
|
return f"{scope}:{idempotency_key}"
|
|
|
|
|
|
def _parse_source_config(raw):
|
|
"""Strict-validate an optional ``config`` JSON blob from the request.
|
|
|
|
Returns ``(config_dict_or_None, error_response)``. Absent/empty → ``None``
|
|
(worker uses classic defaults). Invalid JSON or a config that fails
|
|
``SourceConfig`` validation → a 400 (strict-on-write, D7).
|
|
|
|
Dedup note: ``config`` does NOT participate in the idempotency key, so a
|
|
same-key retry with a different config is ignored — config is immutable for
|
|
the dedup window (the first request's config wins). Changing chunking later
|
|
requires an explicit re-ingest (D8).
|
|
"""
|
|
if not raw:
|
|
return None, None
|
|
try:
|
|
parsed = json.loads(raw) if isinstance(raw, str) else raw
|
|
if not isinstance(parsed, dict):
|
|
raise ValueError("config must be a JSON object")
|
|
# Validate strictly; normalize to a plain dict for the Celery payload.
|
|
return SourceConfig.model_validate(parsed).model_dump(), None
|
|
except Exception:
|
|
return None, make_response(
|
|
jsonify({"success": False, "message": "Invalid source config"}),
|
|
400,
|
|
)
|
|
|
|
|
|
def _claim_task_or_get_cached(key, task_name):
|
|
"""Claim ``key`` for this request OR return the winner's cached payload.
|
|
|
|
Pre-generates the celery task_id so a losing writer sees the same
|
|
id immediately. Returns ``(task_id, cached_response)``; non-None
|
|
cached means the caller should return without enqueuing. The
|
|
cached payload mirrors the fresh-request response shape (including
|
|
``source_id``) so the frontend can correlate SSE ingest events to
|
|
the cached upload task without an extra round-trip — but only when
|
|
the cached row actually exists; the "deduplicated" sentinel
|
|
deliberately omits ``source_id`` so the frontend doesn't bind to a
|
|
phantom source.
|
|
"""
|
|
predetermined_id = str(uuid.uuid4())
|
|
with db_session() as conn:
|
|
claimed = IdempotencyRepository(conn).claim_task(
|
|
key=key, task_name=task_name, task_id=predetermined_id,
|
|
)
|
|
if claimed is not None:
|
|
return claimed["task_id"], None
|
|
with db_readonly() as conn:
|
|
existing = IdempotencyRepository(conn).get_task(key)
|
|
cached_id = existing.get("task_id") if existing else None
|
|
payload: dict = {
|
|
"success": True,
|
|
"task_id": cached_id or "deduplicated",
|
|
}
|
|
# Only surface ``source_id`` when there's a real winner whose worker
|
|
# is publishing SSE events tagged with that id. The "deduplicated"
|
|
# branch means the lock row vanished — we have nothing to correlate.
|
|
if cached_id is not None:
|
|
payload["source_id"] = str(_derive_source_id(key))
|
|
return None, payload
|
|
|
|
|
|
def _release_claim(key):
|
|
"""Drop a pending claim so a client retry can re-claim it."""
|
|
try:
|
|
with db_session() as conn:
|
|
conn.execute(
|
|
sql_text(
|
|
"DELETE FROM task_dedup WHERE idempotency_key = :k "
|
|
"AND status = 'pending'"
|
|
),
|
|
{"k": key},
|
|
)
|
|
except Exception:
|
|
current_app.logger.exception(
|
|
"Failed to release task_dedup claim for key=%s", key,
|
|
)
|
|
|
|
|
|
|
|
|
|
def _enforce_audio_path_size_limit(file_path: str, filename: str) -> None:
|
|
if not is_audio_filename(filename):
|
|
return
|
|
enforce_audio_file_size_limit(os.path.getsize(file_path))
|
|
|
|
|
|
@sources_upload_ns.route("/upload")
|
|
class UploadFile(Resource):
|
|
@api.expect(
|
|
api.model(
|
|
"UploadModel",
|
|
{
|
|
"user": fields.String(required=True, description="User ID"),
|
|
"name": fields.String(required=True, description="Job name"),
|
|
"file": fields.Raw(required=True, description="File(s) to upload"),
|
|
},
|
|
)
|
|
)
|
|
@api.doc(
|
|
description=(
|
|
"Uploads a file to be vectorized and indexed. Honors an optional "
|
|
"``Idempotency-Key`` header: a repeat request with the same key "
|
|
"within 24h returns the original cached response without re-enqueuing."
|
|
),
|
|
)
|
|
def post(self):
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
idempotency_key, key_error = _read_idempotency_key()
|
|
if key_error is not None:
|
|
return key_error
|
|
source_config, config_error = _parse_source_config(
|
|
request.form.get("config")
|
|
)
|
|
if config_error is not None:
|
|
return config_error
|
|
# User-scoped to avoid cross-user collisions; also feeds
|
|
# ``_derive_source_id`` so uuid5 stays user-disjoint.
|
|
scoped_key = _scoped_idempotency_key(idempotency_key, user)
|
|
# Claim before enqueue; the loser returns the winner's task_id.
|
|
predetermined_task_id = None
|
|
if scoped_key:
|
|
predetermined_task_id, cached = _claim_task_or_get_cached(
|
|
scoped_key, "ingest",
|
|
)
|
|
if cached is not None:
|
|
return make_response(jsonify(cached), 200)
|
|
data = request.form
|
|
files = request.files.getlist("file")
|
|
required_fields = ["user", "name"]
|
|
missing_fields = check_required_fields(data, required_fields)
|
|
if missing_fields or not files or all(file.filename == "" for file in files):
|
|
if scoped_key:
|
|
_release_claim(scoped_key)
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"status": "error",
|
|
"message": "Missing required fields or files",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
job_name = request.form["name"]
|
|
|
|
# Create safe versions for filesystem operations
|
|
|
|
safe_user = safe_filename(user)
|
|
dir_name = safe_filename(job_name)
|
|
base_path = f"{settings.UPLOAD_FOLDER}/{safe_user}/{dir_name}"
|
|
file_name_map = {}
|
|
|
|
try:
|
|
storage = StorageCreator.get_storage()
|
|
|
|
for file in files:
|
|
original_filename = os.path.basename(file.filename)
|
|
safe_file = safe_filename(original_filename)
|
|
if original_filename:
|
|
file_name_map[safe_file] = original_filename
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_file_path = os.path.join(temp_dir, safe_file)
|
|
file.save(temp_file_path)
|
|
_enforce_audio_path_size_limit(temp_file_path, safe_file)
|
|
|
|
# Only extract actual .zip files, not Office formats (.docx, .xlsx, .pptx)
|
|
# which are technically zip archives but should be processed as-is
|
|
is_office_format = safe_file.lower().endswith(
|
|
(".docx", ".xlsx", ".pptx", ".odt", ".ods", ".odp", ".epub")
|
|
)
|
|
if zipfile.is_zipfile(temp_file_path) and not is_office_format:
|
|
try:
|
|
with zipfile.ZipFile(temp_file_path, "r") as zip_ref:
|
|
zip_ref.extractall(path=temp_dir)
|
|
|
|
# Walk through extracted files and upload them
|
|
|
|
for root, _, files in os.walk(temp_dir):
|
|
for extracted_file in files:
|
|
if (
|
|
os.path.join(root, extracted_file)
|
|
== temp_file_path
|
|
):
|
|
continue
|
|
rel_path = os.path.relpath(
|
|
os.path.join(root, extracted_file), temp_dir
|
|
)
|
|
storage_path = f"{base_path}/{rel_path}"
|
|
_enforce_audio_path_size_limit(
|
|
os.path.join(root, extracted_file),
|
|
extracted_file,
|
|
)
|
|
|
|
with open(
|
|
os.path.join(root, extracted_file), "rb"
|
|
) as f:
|
|
storage.save_file(f, storage_path)
|
|
except Exception as e:
|
|
current_app.logger.error(
|
|
f"Error extracting zip: {e}", exc_info=True
|
|
)
|
|
# If zip extraction fails, save the original zip file
|
|
|
|
file_path = f"{base_path}/{safe_file}"
|
|
with open(temp_file_path, "rb") as f:
|
|
storage.save_file(f, file_path)
|
|
else:
|
|
# For non-zip files, save directly
|
|
|
|
file_path = f"{base_path}/{safe_file}"
|
|
with open(temp_file_path, "rb") as f:
|
|
storage.save_file(f, file_path)
|
|
# Mint the source UUID up here so the HTTP response and the
|
|
# worker's SSE envelopes share one id. With an idempotency
|
|
# key we reuse the deterministic uuid5 (retried task lands on
|
|
# the same source row); without a key we fall back to uuid4.
|
|
# The worker is told to use this id verbatim — see
|
|
# ``ingest_worker(source_id=...)``.
|
|
source_uuid = (
|
|
_derive_source_id(scoped_key) if scoped_key else uuid.uuid4()
|
|
)
|
|
ingest_kwargs = dict(
|
|
args=(
|
|
settings.UPLOAD_FOLDER,
|
|
list(SUPPORTED_SOURCE_EXTENSIONS),
|
|
job_name,
|
|
user,
|
|
),
|
|
kwargs={
|
|
"file_path": base_path,
|
|
"filename": dir_name,
|
|
"file_name_map": file_name_map,
|
|
"config": source_config,
|
|
# Scoped so the worker dedup row matches the HTTP claim.
|
|
"idempotency_key": scoped_key or idempotency_key,
|
|
"source_id": str(source_uuid),
|
|
},
|
|
)
|
|
if predetermined_task_id is not None:
|
|
ingest_kwargs["task_id"] = predetermined_task_id
|
|
task = ingest.apply_async(**ingest_kwargs)
|
|
except AudioFileTooLargeError:
|
|
if scoped_key:
|
|
_release_claim(scoped_key)
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": build_stt_file_size_limit_message(),
|
|
}
|
|
),
|
|
413,
|
|
)
|
|
except Exception as err:
|
|
current_app.logger.error(f"Error uploading file: {err}", exc_info=True)
|
|
if scoped_key:
|
|
_release_claim(scoped_key)
|
|
return make_response(jsonify({"success": False}), 400)
|
|
# Predetermined id matches the dedup-claim row; loser GET sees same.
|
|
response_task_id = predetermined_task_id or task.id
|
|
# ``source_uuid`` was minted above and passed to the worker as
|
|
# ``source_id``; the worker uses it verbatim for every SSE event,
|
|
# so the frontend can correlate inbound ``source.ingest.*`` to
|
|
# this upload regardless of whether an idempotency key was set.
|
|
response_payload: dict = {
|
|
"success": True,
|
|
"task_id": response_task_id,
|
|
"source_id": str(source_uuid),
|
|
}
|
|
return make_response(jsonify(response_payload), 200)
|
|
|
|
|
|
@sources_upload_ns.route("/remote")
|
|
class UploadRemote(Resource):
|
|
@api.expect(
|
|
api.model(
|
|
"RemoteUploadModel",
|
|
{
|
|
"user": fields.String(required=True, description="User ID"),
|
|
"source": fields.String(
|
|
required=True, description="Source of the data"
|
|
),
|
|
"name": fields.String(required=True, description="Job name"),
|
|
"data": fields.String(required=True, description="Data to process"),
|
|
"repo_url": fields.String(description="GitHub repository URL"),
|
|
},
|
|
)
|
|
)
|
|
@api.doc(
|
|
description=(
|
|
"Uploads remote source for vectorization. Honors an optional "
|
|
"``Idempotency-Key`` header: a repeat request with the same key "
|
|
"within 24h returns the original cached response without re-enqueuing."
|
|
),
|
|
)
|
|
def post(self):
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
idempotency_key, key_error = _read_idempotency_key()
|
|
if key_error is not None:
|
|
return key_error
|
|
source_config, config_error = _parse_source_config(
|
|
request.form.get("config")
|
|
)
|
|
if config_error is not None:
|
|
return config_error
|
|
scoped_key = _scoped_idempotency_key(idempotency_key, user)
|
|
data = request.form
|
|
required_fields = ["user", "source", "name", "data"]
|
|
missing_fields = check_required_fields(data, required_fields)
|
|
if missing_fields:
|
|
return missing_fields
|
|
task_name_for_dedup = (
|
|
"ingest_connector_task"
|
|
if data.get("source") in ConnectorCreator.get_supported_connectors()
|
|
else "ingest_remote"
|
|
)
|
|
predetermined_task_id = None
|
|
if scoped_key:
|
|
predetermined_task_id, cached = _claim_task_or_get_cached(
|
|
scoped_key, task_name_for_dedup,
|
|
)
|
|
if cached is not None:
|
|
return make_response(jsonify(cached), 200)
|
|
# Mint the source UUID up here so the HTTP response and the
|
|
# worker's SSE envelopes share one id. Same pattern as
|
|
# ``UploadFile.post``: with an idempotency key we reuse the
|
|
# deterministic uuid5 (retried task lands on the same source
|
|
# row); without a key we fall back to uuid4. The worker is told
|
|
# to use this id verbatim — see ``remote_worker`` and
|
|
# ``ingest_connector``. Without this the no-key path would mint
|
|
# a random uuid4 inside the worker that the frontend has no way
|
|
# to correlate SSE events to.
|
|
source_uuid = (
|
|
_derive_source_id(scoped_key) if scoped_key else uuid.uuid4()
|
|
)
|
|
try:
|
|
config = json.loads(data["data"])
|
|
source_data = None
|
|
|
|
if data["source"] == "github":
|
|
source_data = config.get("repo_url")
|
|
elif data["source"] in ["crawler", "url"]:
|
|
source_data = config.get("url")
|
|
elif data["source"] == "reddit":
|
|
source_data = config
|
|
elif data["source"] == "s3":
|
|
source_data = config
|
|
elif data["source"] in ConnectorCreator.get_supported_connectors():
|
|
session_token = config.get("session_token")
|
|
if not session_token:
|
|
if scoped_key:
|
|
_release_claim(scoped_key)
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"error": f"Missing session_token in {data['source']} configuration",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
# Process file_ids
|
|
|
|
file_ids = config.get("file_ids", [])
|
|
if isinstance(file_ids, str):
|
|
file_ids = [id.strip() for id in file_ids.split(",") if id.strip()]
|
|
elif not isinstance(file_ids, list):
|
|
file_ids = []
|
|
# Process folder_ids
|
|
|
|
folder_ids = config.get("folder_ids", [])
|
|
if isinstance(folder_ids, str):
|
|
folder_ids = [
|
|
id.strip() for id in folder_ids.split(",") if id.strip()
|
|
]
|
|
elif not isinstance(folder_ids, list):
|
|
folder_ids = []
|
|
config["file_ids"] = file_ids
|
|
config["folder_ids"] = folder_ids
|
|
|
|
connector_kwargs = {
|
|
"kwargs": {
|
|
"job_name": data["name"],
|
|
"user": user,
|
|
"source_type": data["source"],
|
|
"session_token": session_token,
|
|
"file_ids": file_ids,
|
|
"folder_ids": folder_ids,
|
|
"recursive": config.get("recursive", False),
|
|
"retriever": config.get("retriever", "classic"),
|
|
"config": source_config,
|
|
"idempotency_key": scoped_key or idempotency_key,
|
|
"source_id": str(source_uuid),
|
|
},
|
|
}
|
|
if predetermined_task_id is not None:
|
|
connector_kwargs["task_id"] = predetermined_task_id
|
|
task = ingest_connector_task.apply_async(**connector_kwargs)
|
|
response_task_id = predetermined_task_id or task.id
|
|
# ``source_uuid`` was minted above and passed to the
|
|
# worker as ``source_id``; the worker uses it verbatim
|
|
# for every SSE event, so the frontend can correlate
|
|
# inbound ``source.ingest.*`` regardless of whether an
|
|
# idempotency key was set.
|
|
response_payload = {
|
|
"success": True,
|
|
"task_id": response_task_id,
|
|
"source_id": str(source_uuid),
|
|
}
|
|
return make_response(jsonify(response_payload), 200)
|
|
remote_kwargs = {
|
|
"kwargs": {
|
|
"source_data": source_data,
|
|
"job_name": data["name"],
|
|
"user": user,
|
|
"loader": data["source"],
|
|
"config": source_config,
|
|
"idempotency_key": scoped_key or idempotency_key,
|
|
"source_id": str(source_uuid),
|
|
},
|
|
}
|
|
if predetermined_task_id is not None:
|
|
remote_kwargs["task_id"] = predetermined_task_id
|
|
task = ingest_remote.apply_async(**remote_kwargs)
|
|
except Exception as err:
|
|
current_app.logger.error(
|
|
f"Error uploading remote source: {err}", exc_info=True
|
|
)
|
|
if scoped_key:
|
|
_release_claim(scoped_key)
|
|
return make_response(jsonify({"success": False}), 400)
|
|
response_task_id = predetermined_task_id or task.id
|
|
response_payload = {
|
|
"success": True,
|
|
"task_id": response_task_id,
|
|
"source_id": str(source_uuid),
|
|
}
|
|
return make_response(jsonify(response_payload), 200)
|
|
|
|
|
|
@sources_upload_ns.route("/manage_source_files")
|
|
class ManageSourceFiles(Resource):
|
|
@api.expect(
|
|
api.model(
|
|
"ManageSourceFilesModel",
|
|
{
|
|
"source_id": fields.String(
|
|
required=True, description="Source ID to modify"
|
|
),
|
|
"operation": fields.String(
|
|
required=True,
|
|
description="Operation: 'add', 'remove', or 'remove_directory'",
|
|
),
|
|
"file_paths": fields.List(
|
|
fields.String,
|
|
required=False,
|
|
description="File paths to remove (for remove operation)",
|
|
),
|
|
"directory_path": fields.String(
|
|
required=False,
|
|
description="Directory path to remove (for remove_directory operation)",
|
|
),
|
|
"file": fields.Raw(
|
|
required=False, description="Files to add (for add operation)"
|
|
),
|
|
"parent_dir": fields.String(
|
|
required=False,
|
|
description="Parent directory path relative to source root",
|
|
),
|
|
},
|
|
)
|
|
)
|
|
@api.doc(
|
|
description="Add files, remove files, or remove directories from an existing source",
|
|
)
|
|
def post(self):
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Unauthorized"}), 401
|
|
)
|
|
user = decoded_token.get("sub")
|
|
idempotency_key, key_error = _read_idempotency_key()
|
|
if key_error is not None:
|
|
return key_error
|
|
scoped_key = _scoped_idempotency_key(idempotency_key, user)
|
|
source_id = request.form.get("source_id")
|
|
operation = request.form.get("operation")
|
|
|
|
if not source_id or not operation:
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "source_id and operation are required",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
if operation not in ["add", "remove", "remove_directory"]:
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "operation must be 'add', 'remove', or 'remove_directory'",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
# Resolve the owner to run the file mutation + reingest AS: ``user``
|
|
# when they own the source, the real owner when ``user`` holds a team
|
|
# ``editor`` grant (add/remove files is an editor-allowed write). The
|
|
# storage layout and vector partition are keyed by source_id
|
|
# (owner-agnostic), so running the ops as the owner is correct.
|
|
try:
|
|
with db_readonly() as conn:
|
|
source = SourcesRepository(conn).get_any(source_id, user)
|
|
owner = user
|
|
if source is None:
|
|
owner = effective_write_owner(conn, "source", source_id, user)
|
|
if owner:
|
|
source = SourcesRepository(conn).get_any(source_id, owner)
|
|
if not source:
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "Source not found or access denied",
|
|
}
|
|
),
|
|
404,
|
|
)
|
|
except Exception as err:
|
|
current_app.logger.error(f"Error finding source: {err}", exc_info=True)
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Database error"}), 500
|
|
)
|
|
resolved_source_id = str(source["id"])
|
|
# Flips to True after each branch's ``apply_async`` returns
|
|
# successfully — at that point the worker owns the predetermined
|
|
# task_id. The outer ``except`` only releases the claim while
|
|
# this is False, so a post-``apply_async`` failure (jsonify,
|
|
# make_response, etc.) doesn't double-enqueue on the next retry.
|
|
claim_transferred = False
|
|
try:
|
|
storage = StorageCreator.get_storage()
|
|
source_file_path = source.get("file_path", "")
|
|
parent_dir = request.form.get("parent_dir", "")
|
|
file_name_map = source.get("file_name_map") or {}
|
|
if isinstance(file_name_map, str):
|
|
try:
|
|
file_name_map = json.loads(file_name_map)
|
|
except Exception:
|
|
file_name_map = {}
|
|
if not isinstance(file_name_map, dict):
|
|
file_name_map = {}
|
|
|
|
if parent_dir and (parent_dir.startswith("/") or ".." in parent_dir):
|
|
return make_response(
|
|
jsonify(
|
|
{"success": False, "message": "Invalid parent directory path"}
|
|
),
|
|
400,
|
|
)
|
|
if operation == "add":
|
|
files = request.files.getlist("file")
|
|
if not files or all(file.filename == "" for file in files):
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "No files provided for add operation",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
|
|
# Claim before any storage mutation so a duplicate request
|
|
# short-circuits without touching the filesystem. Mirrors
|
|
# the pattern in ``UploadFile.post`` / ``UploadRemote.post``
|
|
# — without it ``.delay()`` would enqueue twice for two
|
|
# racing same-key POSTs (the worker decorator only
|
|
# deduplicates *after* completion).
|
|
predetermined_task_id = None
|
|
if scoped_key:
|
|
predetermined_task_id, cached = _claim_task_or_get_cached(
|
|
scoped_key, "reingest_source_task",
|
|
)
|
|
if cached is not None:
|
|
# Frontend keys reingest polling on
|
|
# ``reingest_task_id``; the shared cache helper
|
|
# writes ``task_id``. Alias here so a dedup
|
|
# response doesn't silently break FileTree's
|
|
# poller. Override ``source_id`` too — the
|
|
# helper derives it from the scoped key, which
|
|
# is correct for upload but wrong for reingest
|
|
# (the worker publishes events scoped to the
|
|
# actual source row id).
|
|
cached_task_id = cached.pop("task_id", None)
|
|
if cached_task_id is not None:
|
|
cached["reingest_task_id"] = cached_task_id
|
|
cached["source_id"] = resolved_source_id
|
|
return make_response(jsonify(cached), 200)
|
|
|
|
added_files = []
|
|
map_updated = False
|
|
|
|
target_dir = source_file_path
|
|
if parent_dir:
|
|
target_dir = f"{source_file_path}/{parent_dir}"
|
|
for file in files:
|
|
if file.filename:
|
|
original_filename = os.path.basename(file.filename)
|
|
safe_filename_str = safe_filename(original_filename)
|
|
file_path = f"{target_dir}/{safe_filename_str}"
|
|
|
|
# Save file to storage
|
|
|
|
storage.save_file(file, file_path)
|
|
added_files.append(safe_filename_str)
|
|
if original_filename:
|
|
relative_key = (
|
|
f"{parent_dir}/{safe_filename_str}"
|
|
if parent_dir
|
|
else safe_filename_str
|
|
)
|
|
file_name_map[relative_key] = original_filename
|
|
map_updated = True
|
|
|
|
if map_updated:
|
|
with db_session() as conn:
|
|
SourcesRepository(conn).update(
|
|
resolved_source_id, owner,
|
|
{"file_name_map": dict(file_name_map)},
|
|
)
|
|
# Trigger re-ingestion pipeline
|
|
|
|
from application.api.user.tasks import reingest_source_task
|
|
|
|
task = reingest_source_task.apply_async(
|
|
kwargs={
|
|
"source_id": resolved_source_id,
|
|
"user": owner,
|
|
"idempotency_key": scoped_key or idempotency_key,
|
|
},
|
|
task_id=predetermined_task_id,
|
|
)
|
|
claim_transferred = True
|
|
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": True,
|
|
"message": f"Added {len(added_files)} files",
|
|
"added_files": added_files,
|
|
"parent_dir": parent_dir,
|
|
"reingest_task_id": task.id,
|
|
# ``source_id`` lets the frontend correlate
|
|
# inbound ``source.ingest.*`` SSE events
|
|
# (emitted by ``reingest_source_worker``)
|
|
# back to the reingest task — matches the
|
|
# upload route's source-id contract.
|
|
"source_id": resolved_source_id,
|
|
}
|
|
),
|
|
200,
|
|
)
|
|
elif operation == "remove":
|
|
file_paths_str = request.form.get("file_paths")
|
|
if not file_paths_str:
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "file_paths required for remove operation",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
try:
|
|
file_paths = (
|
|
json.loads(file_paths_str)
|
|
if isinstance(file_paths_str, str)
|
|
else file_paths_str
|
|
)
|
|
except Exception:
|
|
return make_response(
|
|
jsonify(
|
|
{"success": False, "message": "Invalid file_paths format"}
|
|
),
|
|
400,
|
|
)
|
|
# Path-traversal guard runs *before* the claim so a 400
|
|
# for an invalid path doesn't leave a pending dedup row.
|
|
for file_path in file_paths:
|
|
if ".." in str(file_path) or str(file_path).startswith("/"):
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "Invalid file path",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
|
|
# Claim before any storage mutation. See ``add`` branch
|
|
# comment for rationale.
|
|
predetermined_task_id = None
|
|
if scoped_key:
|
|
predetermined_task_id, cached = _claim_task_or_get_cached(
|
|
scoped_key, "reingest_source_task",
|
|
)
|
|
if cached is not None:
|
|
cached_task_id = cached.pop("task_id", None)
|
|
if cached_task_id is not None:
|
|
cached["reingest_task_id"] = cached_task_id
|
|
# Override the helper's synthetic source_id (uuid5
|
|
# of the scoped key) with the real source row id
|
|
# — the reingest worker publishes SSE events
|
|
# scoped to ``resolved_source_id`` and FileTree
|
|
# correlates on it.
|
|
cached["source_id"] = resolved_source_id
|
|
return make_response(jsonify(cached), 200)
|
|
|
|
# Remove files from storage and directory structure
|
|
|
|
removed_files = []
|
|
map_updated = False
|
|
for file_path in file_paths:
|
|
full_path = f"{source_file_path}/{file_path}"
|
|
|
|
# Remove from storage
|
|
|
|
if storage.file_exists(full_path):
|
|
storage.delete_file(full_path)
|
|
removed_files.append(file_path)
|
|
if file_path in file_name_map:
|
|
file_name_map.pop(file_path, None)
|
|
map_updated = True
|
|
|
|
if map_updated and isinstance(file_name_map, dict):
|
|
with db_session() as conn:
|
|
SourcesRepository(conn).update(
|
|
resolved_source_id, owner,
|
|
{"file_name_map": dict(file_name_map)},
|
|
)
|
|
# Trigger re-ingestion pipeline
|
|
|
|
from application.api.user.tasks import reingest_source_task
|
|
|
|
task = reingest_source_task.apply_async(
|
|
kwargs={
|
|
"source_id": resolved_source_id,
|
|
"user": owner,
|
|
"idempotency_key": scoped_key or idempotency_key,
|
|
},
|
|
task_id=predetermined_task_id,
|
|
)
|
|
claim_transferred = True
|
|
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": True,
|
|
"message": f"Removed {len(removed_files)} files",
|
|
"removed_files": removed_files,
|
|
"reingest_task_id": task.id,
|
|
"source_id": resolved_source_id,
|
|
}
|
|
),
|
|
200,
|
|
)
|
|
elif operation == "remove_directory":
|
|
directory_path = request.form.get("directory_path")
|
|
if not directory_path:
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "directory_path required for remove_directory operation",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
# Validate directory path (prevent path traversal)
|
|
|
|
if directory_path.startswith("/") or ".." in directory_path:
|
|
current_app.logger.warning(
|
|
f"Invalid directory path attempted for removal. "
|
|
f"User: {user}, Source ID: {source_id}, Directory path: {directory_path}"
|
|
)
|
|
return make_response(
|
|
jsonify(
|
|
{"success": False, "message": "Invalid directory path"}
|
|
),
|
|
400,
|
|
)
|
|
full_directory_path = (
|
|
f"{source_file_path}/{directory_path}"
|
|
if directory_path
|
|
else source_file_path
|
|
)
|
|
|
|
if not storage.is_directory(full_directory_path):
|
|
current_app.logger.warning(
|
|
f"Directory not found or is not a directory for removal. "
|
|
f"User: {user}, Source ID: {source_id}, Directory path: {directory_path}, "
|
|
f"Full path: {full_directory_path}"
|
|
)
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "Directory not found or is not a directory",
|
|
}
|
|
),
|
|
404,
|
|
)
|
|
|
|
# Claim before mutation. See ``add`` branch for rationale.
|
|
predetermined_task_id = None
|
|
if scoped_key:
|
|
predetermined_task_id, cached = _claim_task_or_get_cached(
|
|
scoped_key, "reingest_source_task",
|
|
)
|
|
if cached is not None:
|
|
cached_task_id = cached.pop("task_id", None)
|
|
if cached_task_id is not None:
|
|
cached["reingest_task_id"] = cached_task_id
|
|
# Same source_id override as the ``remove`` /
|
|
# ``add`` cached branches — the helper's synthetic
|
|
# id doesn't match what reingest_source_worker
|
|
# tags its SSE events with.
|
|
cached["source_id"] = resolved_source_id
|
|
return make_response(jsonify(cached), 200)
|
|
|
|
success = storage.remove_directory(full_directory_path)
|
|
|
|
if not success:
|
|
current_app.logger.error(
|
|
f"Failed to remove directory from storage. "
|
|
f"User: {user}, Source ID: {source_id}, Directory path: {directory_path}, "
|
|
f"Full path: {full_directory_path}"
|
|
)
|
|
# Release so a client retry can reclaim — otherwise
|
|
# the next request would silently 200-cache to the
|
|
# task_id that never enqueued.
|
|
if scoped_key:
|
|
_release_claim(scoped_key)
|
|
return make_response(
|
|
jsonify(
|
|
{"success": False, "message": "Failed to remove directory"}
|
|
),
|
|
500,
|
|
)
|
|
current_app.logger.info(
|
|
f"Successfully removed directory. "
|
|
f"User: {user}, Source ID: {source_id}, Directory path: {directory_path}, "
|
|
f"Full path: {full_directory_path}"
|
|
)
|
|
if directory_path and file_name_map:
|
|
prefix = f"{directory_path.rstrip('/')}/"
|
|
keys_to_remove = [
|
|
key
|
|
for key in file_name_map.keys()
|
|
if key == directory_path or key.startswith(prefix)
|
|
]
|
|
if keys_to_remove:
|
|
for key in keys_to_remove:
|
|
file_name_map.pop(key, None)
|
|
with db_session() as conn:
|
|
SourcesRepository(conn).update(
|
|
resolved_source_id, owner,
|
|
{"file_name_map": dict(file_name_map)},
|
|
)
|
|
|
|
# Trigger re-ingestion pipeline
|
|
|
|
from application.api.user.tasks import reingest_source_task
|
|
|
|
task = reingest_source_task.apply_async(
|
|
kwargs={
|
|
"source_id": resolved_source_id,
|
|
"user": owner,
|
|
"idempotency_key": scoped_key or idempotency_key,
|
|
},
|
|
task_id=predetermined_task_id,
|
|
)
|
|
claim_transferred = True
|
|
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": True,
|
|
"message": f"Successfully removed directory: {directory_path}",
|
|
"removed_directory": directory_path,
|
|
"reingest_task_id": task.id,
|
|
"source_id": resolved_source_id,
|
|
}
|
|
),
|
|
200,
|
|
)
|
|
except Exception as err:
|
|
# Release the dedup claim only if it wasn't transferred to
|
|
# a worker. Without this, a same-key retry within the 24h
|
|
# TTL would 200-cache to a predetermined task_id whose
|
|
# ``apply_async`` never ran (or ran but the response builder
|
|
# blew up afterward — only the first case matters in
|
|
# practice; the flag protects both).
|
|
if scoped_key and not claim_transferred:
|
|
_release_claim(scoped_key)
|
|
error_context = f"operation={operation}, user={user}, source_id={source_id}"
|
|
if operation == "remove_directory":
|
|
directory_path = request.form.get("directory_path", "")
|
|
error_context += f", directory_path={directory_path}"
|
|
elif operation == "remove":
|
|
file_paths_str = request.form.get("file_paths", "")
|
|
error_context += f", file_paths={file_paths_str}"
|
|
elif operation == "add":
|
|
parent_dir = request.form.get("parent_dir", "")
|
|
error_context += f", parent_dir={parent_dir}"
|
|
current_app.logger.error(
|
|
f"Error managing source files: {err} ({error_context})", exc_info=True
|
|
)
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Operation failed"}), 500
|
|
)
|
|
|
|
|
|
@sources_upload_ns.route("/task_status")
|
|
class TaskStatus(Resource):
|
|
task_status_model = api.model(
|
|
"TaskStatusModel",
|
|
{"task_id": fields.String(required=True, description="Task ID")},
|
|
)
|
|
|
|
@api.expect(task_status_model)
|
|
@api.doc(description="Get celery job status")
|
|
def get(self):
|
|
task_id = request.args.get("task_id")
|
|
if not task_id:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Task ID is required"}), 400
|
|
)
|
|
try:
|
|
from application.celery_init import celery
|
|
|
|
task = celery.AsyncResult(task_id)
|
|
task_meta = task.info
|
|
print(f"Task status: {task.status}")
|
|
|
|
if task.status == "PENDING":
|
|
inspect = celery.control.inspect()
|
|
active_workers = inspect.ping()
|
|
if not active_workers:
|
|
raise ConnectionError("Service unavailable")
|
|
|
|
if not isinstance(
|
|
task_meta, (dict, list, str, int, float, bool, type(None))
|
|
):
|
|
task_meta = str(task_meta) # Convert to a string representation
|
|
except ConnectionError as err:
|
|
current_app.logger.error(f"Connection error getting task status: {err}")
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Service unavailable"}), 503
|
|
)
|
|
except Exception as err:
|
|
current_app.logger.error(f"Error getting task status: {err}", exc_info=True)
|
|
return make_response(jsonify({"success": False}), 400)
|
|
return make_response(jsonify({"status": task.status, "result": task_meta}), 200)
|