Merge pull request #3712 from dimagi/cs/export_apis_followup

feat: expand team-export surface to remaining team models

Author: SmittieC

api-schemas/export.yml

@@ -3,7 +3,12 @@
 
 from functools import cache
 
+import pydantic
+from django.contrib.contenttypes.models import ContentType
+from django.db import models
+from django_pydantic_field.fields import PydanticSchemaField
 from drf_spectacular.utils import OpenApiResponse, extend_schema_field
+from pgvector.django import HalfVectorField, VectorField
 from rest_framework import serializers
 
 from apps.teams.export.manifest import (
@@ -13,20 +18,64 @@
     TEAM_MODEL,
     ManifestEntry,
     entry_model,
+    generic_fk_fields,
     model_has_team_field,
 )
 from apps.teams.export.seal import seal
 from apps.teams.models import Flag
 
 
+def _json_safe(value):
+    """Reduce a pydantic object (e.g. CollectionFile.metadata, DocumentSource.config) to plain JSON
+    data. Non-pydantic values are already JSON-native and pass through."""
+    return value.model_dump(mode="json") if isinstance(value, pydantic.BaseModel) else value
+
+
+class _RelativeFileField(serializers.FileField):
+    """Serialize a file as its stored relative path (``value.name``) rather than the default
+    ``/media/...`` URL. The URL is a presentation form; the importer needs the raw name so it
+    assigns straight back into a FileField -- an absolute ``/media/`` path is rejected by Django's
+    storage safe-join (SuspiciousFileOperation)."""
+
+    def __init__(self, *args, **kwargs):
+        kwargs.setdefault("use_url", False)
+        super().__init__(*args, **kwargs)
+
+
+@extend_schema_field(serializers.ListField(child=serializers.FloatField()))
+class _VectorField(serializers.Field):
+    """Serialize a pgvector column as a plain list of floats. Its DB string form (e.g. ``[0.1,0.2]``)
+    isn't accepted as input, so the importer needs the list to assign it straight back."""
+
+    def to_representation(self, value):
+        # A HalfVectorField reads back as a pgvector HalfVector object, which isn't iterable but
+        # exposes .to_list(); a VectorField reads back as an (iterable) numpy array.
+        if hasattr(value, "to_list"):
+            return value.to_list()
+        return [float(x) for x in value]
+
+    def to_internal_value(self, data):
+        return data
+
+
+class _PydanticSchemaField(serializers.JSONField):
+    """Serialize a django_pydantic_field SchemaField to a plain JSON object. The default JSONField
+    mapping leaves the pydantic model instance in place, which DRF's JSON encoder then mangles into a
+    list of (field, value) pairs (it falls back to iterating the model) that the importer can't read
+    back into the schema. Subclasses JSONField so DRF's encoder/decoder kwargs are still accepted."""
+
+    def to_representation(self, value):
+        return _json_safe(value)
+
+
 class _SecretMixin:
     secret_fields: list[str] = []
 
     def to_representation(self, instance):
         data = super().to_representation(instance)
         public_key = self.context.get("public_key")
         for field in self.secret_fields:
-            data[field] = seal(getattr(instance, field), public_key)
+            data[field] = seal(_json_safe(getattr(instance, field)), public_key)
         return data
 
 
@@ -50,6 +99,23 @@ def _team_role_groups(self, user):
     return []
 
 
+def _content_type_label_resolver(ct_field: str):
+    """Build the method that emits a generic FK's content type as its ``app_label.model`` label
+    instead of the source ContentType pk (which differs between servers). Reads the ``*_id`` attribute
+    and looks the type up by id (cached), so it never triggers a related-object query."""
+    attname = f"{ct_field}_id"
+
+    @extend_schema_field(serializers.CharField())
+    def get_content_type_label(self, instance) -> str | None:
+        type_id = getattr(instance, attname)
+        if type_id is None:
+            return None
+        ct = ContentType.objects.get_for_id(type_id)
+        return f"{ct.app_label}.{ct.model}"
+
+    return get_content_type_label
+
+
 def _is_global_resolver(null_field: str):
     """Build the ``is_global`` method: a row is global when its scoping field (team or voice
     provider) is null. Reads the ``*_id`` attribute so it never triggers a related-object query."""
@@ -107,6 +173,14 @@ def build_resource_serializer(model):
     attrs: dict[str, object] = {
         "Meta": type("Meta", (), meta_attrs),
         "secret_fields": secret_fields,
+        "serializer_field_mapping": {
+            **serializers.ModelSerializer.serializer_field_mapping,
+            models.FileField: _RelativeFileField,
+            models.ImageField: _RelativeFileField,
+            HalfVectorField: _VectorField,
+            VectorField: _VectorField,
+            PydanticSchemaField: _PydanticSchemaField,
+        },
     }
     # Secret fields are sealed to a base64 string at runtime (see _SecretMixin), so document them as
     # that string rather than their raw model field type.
@@ -120,6 +194,12 @@ def build_resource_serializer(model):
         attrs["is_global"] = serializers.SerializerMethodField()
         attrs["get_is_global"] = _is_global_resolver(spec.null_field)
 
+    # A generic FK's content type goes out as its model label, not the source ContentType pk. The
+    # object-id column keeps its default (the source target-row pk); the importer resolves both.
+    for ct_field, _fk_field in generic_fk_fields(model):
+        attrs[ct_field] = serializers.SerializerMethodField()
+        attrs[f"get_{ct_field}"] = _content_type_label_resolver(ct_field)
+
     return type(f"{component_name(model)}DetailSerializer", (_SecretMixin, serializers.ModelSerializer), attrs)
 
 

apps/api/export/serializers.py

@@ -1,31 +1,41 @@
+import json
+
 import pytest
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
+from django.conf import settings
 from django.contrib.auth.models import Group
 from django.db import connection
 from django.test.utils import CaptureQueriesContext
 from rest_framework import serializers
+from rest_framework.renderers import JSONRenderer
 
 from apps.api.export.serializers import (
     ManifestEntrySerializer,
     ManifestSerializer,
     build_resource_response_serializer,
     build_resource_serializer,
 )
+from apps.assessments.models import Score
 from apps.chat.models import Chat
+from apps.documents.models import CollectionFile
 from apps.experiments.models import Experiment, ExperimentSession, ParticipantData
+from apps.files.models import File, FileChunkEmbedding
 from apps.service_providers.models import LlmProvider, LlmProviderModel
 from apps.teams.export import seal as seal_mod
 from apps.teams.export.manifest import build_manifest, entry_model, get_manifest_entry, team_scoped_queryset
 from apps.teams.models import Membership, Team
 from apps.users.models import CustomUser
+from apps.utils.factories.assessments import ScoreFactory
+from apps.utils.factories.documents import CollectionFileFactory
 from apps.utils.factories.experiment import (
     ChatFactory,
     ConsentFormFactory,
     ExperimentFactory,
     ExperimentSessionFactory,
     ParticipantFactory,
 )
+from apps.utils.factories.files import FileChunkEmbeddingFactory, FileFactory
 from apps.utils.factories.service_provider_factories import LlmProviderFactory
 from apps.utils.factories.team import TeamFactory
 from apps.utils.factories.user import UserFactory
@@ -188,6 +198,16 @@ def test_global_able_model_exposes_is_global_flag_instead_of_team(is_global):
     assert "team" not in data
 
 
+def test_generic_fk_content_type_serializes_as_model_label():
+    """A generic FK's content type is emitted as its ``app_label.model`` label, not the source
+    ContentType pk (which is server-specific). The object id stays the source target-row pk; the
+    importer resolves both on the target."""
+    score = ScoreFactory()  # target is an ExperimentSession
+    data = _serialize(Score, score)
+    assert data["target_content_type"] == "experiments.experimentsession"
+    assert data["target_object_id"] == score.target_object_id
+
+
 def test_response_envelope_has_cursor_has_more_and_results():
     fields = build_resource_response_serializer(get_manifest_entry("users"))().fields
     assert set(fields) == {"cursor", "has_more", "results"}
@@ -208,3 +228,34 @@ def test_manifest_serializer_matches_build_manifest_payload():
     declared = set(ManifestEntrySerializer().fields)
     actual_keys = set().union(*(entry.keys() for entry in manifest["entries"]))
     assert declared == actual_keys
+
+
+def test_file_field_serializes_as_relative_name_not_media_url():
+    """A FileField serializes to its stored relative name so the importer can assign it straight
+    back; the default ``/media/...`` URL is an absolute path Django's storage rejects on save."""
+    file = FileFactory()
+    data = build_resource_serializer(File)(file, context={"public_key": None}).data
+    assert data["file"] == file.file.name
+    assert not data["file"].startswith("/")
+
+
+def test_halfvector_field_serializes_as_list_of_floats():
+    """A pgvector column reads back as a non-iterable HalfVector object; the serializer must still
+    emit a plain list of floats the importer can assign straight back."""
+    size = settings.EMBEDDING_VECTOR_SIZE
+    created = FileChunkEmbeddingFactory(embedding=[0.5] * size)
+    instance = FileChunkEmbedding.objects.get(pk=created.pk)  # re-fetch so embedding is the DB object
+    data = build_resource_serializer(FileChunkEmbedding)(instance, context={"public_key": None}).data
+    assert data["embedding"] == [0.5] * size
+
+
+def test_pydantic_schema_field_serializes_as_json_object():
+    """A SchemaField (pydantic) serializes to a JSON object. The default JSONField mapping leaves the
+    pydantic instance in place, which DRF's JSON encoder mangles into (field, value) pairs the
+    importer can't read back. The corruption is at encode time, so render to JSON to catch it."""
+    metadata = {"chunking_strategy": {"chunk_size": 400, "chunk_overlap": 200}}
+    cf = CollectionFileFactory(metadata=metadata)
+    instance = CollectionFile.objects.get(pk=cf.pk)
+    data = build_resource_serializer(CollectionFile)(instance, context={"public_key": None}).data
+    rendered = json.loads(JSONRenderer().render(data))
+    assert rendered["metadata"] == metadata

apps/api/export/tests/test_serializers.py

@@ -9,9 +9,11 @@
 from django.utils import timezone
 from rest_framework.test import APIClient
 
+from apps.documents.datamodels import DocumentSourceConfig, GitHubSourceConfig
 from apps.experiments.models import Participant
 from apps.pipelines.models import Pipeline
 from apps.teams.export import seal as seal_mod
+from apps.utils.factories.documents import CollectionFactory, DocumentSourceFactory
 from apps.utils.factories.experiment import ConsentFormFactory, ParticipantFactory
 from apps.utils.factories.pipelines import PipelineFactory
 from apps.utils.factories.service_provider_factories import LlmProviderFactory, LlmProviderModelFactory
@@ -240,3 +242,27 @@ def test_updated_at_id_cursor_pages_ties_without_skip_or_repeat():
             break
 
     assert sorted(collected) == sorted(p.id for p in parts)  # each tied row exactly once
+
+
+def test_document_sources_resource_seals_config():
+    """A document_sources row seals its config field through the HTTP layer and is team-scoped."""
+    private = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    public_pem = private.public_key().public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+    team = TeamWithUsersFactory(public_key=public_pem.decode())
+    collection = CollectionFactory(team=team)
+    secret_token = "ghp_supersecrettoken"
+    config = DocumentSourceConfig(
+        github=GitHubSourceConfig(repo_url="https://github.com/dimagi/open-chat-studio", branch=secret_token)
+    )
+    DocumentSourceFactory(team=team, collection=collection, config=config)
+    client = ApiTestClient(_admin(team), team)
+    response = client.get(_resource_url("document_sources"))
+    assert response.status_code == 200
+    sealed = response.json()["results"][0]["config"]
+    # The response must be an opaque sealed token, not plaintext config.
+    assert secret_token not in str(sealed)
+    unsealed = seal_mod.unseal(sealed, private)
+    assert unsealed["github"]["branch"] == secret_token