✨(backend) handle uploaded images in Notion import

Author: emersion

src/backend/core/api/viewsets.py

@@ -2111,13 +2111,33 @@ def notion_import_callback(request):
     return redirect(f"{settings.FRONTEND_URL}/import-notion/")
 
 
-def _import_notion_child_page(imported_doc, parent_doc, user, imported_docs_by_page_id):
-    document_content = YdocConverter().convert_blocks(imported_doc.blocks)
+def _import_notion_doc_content(imported_doc, obj, user):
+    for att in imported_doc.attachments:
+        extra_args = {
+            "Metadata": {
+                "owner": str(user.id),
+                "status": enums.DocumentAttachmentStatus.READY,  # TODO
+            },
+        }
+        file_id = uuid.uuid4()
+        key = f"{obj.key_base}/{enums.ATTACHMENTS_FOLDER:s}/{file_id!s}.raw"
+        with requests.get(att.file.file["url"], stream=True) as resp:
+            default_storage.connection.meta.client.upload_fileobj(
+                resp.raw, default_storage.bucket_name, key
+            )
+        obj.attachments.append(key)
+        att.block["props"]["url"] = (
+            f"{settings.MEDIA_BASE_URL}{settings.MEDIA_URL}{key}"
+        )
+
+    obj.content = YdocConverter().convert_blocks(imported_doc.blocks)
+    obj.save()
 
+
+def _import_notion_child_page(imported_doc, parent_doc, user, imported_docs_by_page_id):
     obj = parent_doc.add_child(
         creator=user,
         title=imported_doc.page.get_title() or "Child page",
-        content=document_content,
     )
 
     models.DocumentAccess.objects.create(
@@ -2126,21 +2146,20 @@ def _import_notion_child_page(imported_doc, parent_doc, user, imported_docs_by_p
         role=models.RoleChoices.OWNER,
     )
 
+    _import_notion_doc_content(imported_doc, obj, user)
+
     imported_docs_by_page_id[imported_doc.page.id] = obj
 
     for child in imported_doc.children:
         _import_notion_child_page(child, obj, user, imported_docs_by_page_id)
 
 
 def _import_notion_root_page(imported_doc, user, imported_docs_by_page_id):
-    document_content = YdocConverter().convert_blocks(imported_doc.blocks)
-
     obj = models.Document.add_root(
         depth=1,
         creator=user,
         title=imported_doc.page.get_title() or "Imported Notion page",
         link_reach=models.LinkReachChoices.RESTRICTED,
-        content=document_content,
     )
 
     models.DocumentAccess.objects.create(
@@ -2149,13 +2168,15 @@ def _import_notion_root_page(imported_doc, user, imported_docs_by_page_id):
         role=models.RoleChoices.OWNER,
     )
 
+    _import_notion_doc_content(imported_doc, obj, user)
+
     imported_docs_by_page_id[imported_doc.page.id] = obj
 
     for child in imported_doc.children:
         _import_notion_child_page(child, obj, user, imported_docs_by_page_id)
 
 
-@drf.decorators.api_view(["POST"])
+@drf.decorators.api_view(["GET", "POST"])  # TODO: drop GET (used for testing)
 def notion_import_run(request):
     if "notion_token" not in request.session:
         raise drf.exceptions.PermissionDenied()

src/backend/core/notion_schemas/notion_block.py

@@ -166,20 +166,12 @@ class NotionEmbed(BaseModel):
     url: str
 
 
-class NotionFileType(StrEnum):
-    FILE = "file"
-    EXTERNAL = "external"
-    FILE_UPLOAD = "file_upload"
-
-
-class NotionFile(BaseModel):
+class NotionBlockFile(BaseModel):
     # FIXME: this is actually another occurrence of type discriminating
     """https://developers.notion.com/reference/block#file"""
 
     block_type: Literal[NotionBlockType.FILE] = NotionBlockType.FILE
-    caption: list[NotionRichText]
-    type: NotionFileType
-    ...
+    # TODO: NotionFile
 
 
 class NotionImage(BaseModel):
@@ -188,6 +180,13 @@ class NotionImage(BaseModel):
     block_type: Literal[NotionBlockType.IMAGE] = NotionBlockType.IMAGE
     file: NotionFile
 
+    @model_validator(mode="before")
+    @classmethod
+    def move_file_type_inward_and_rename(cls, data: Any) -> Any:
+        if not isinstance(data, dict):
+            return data
+        return {"block_type": "image", "file": data}
+
 
 class NotionVideo(BaseModel):
     """https://developers.notion.com/reference/block#video"""
@@ -280,7 +279,7 @@ def put_all_in_raw(cls, data: Any) -> Any:
         | NotionColumnList
         | NotionDivider
         | NotionEmbed
-        | NotionFile
+        | NotionBlockFile
         | NotionImage
         | NotionVideo
         | NotionLinkPreview

src/backend/core/services/notion_import.py

@@ -18,13 +18,15 @@
     NotionHeading1,
     NotionHeading2,
     NotionHeading3,
+    NotionImage,
     NotionNumberedListItem,
     NotionParagraph,
     NotionTable,
     NotionTableRow,
     NotionToDo,
     NotionUnsupported,
 )
+from ..notion_schemas.notion_file import NotionFileExternal, NotionFileHosted
 from ..notion_schemas.notion_page import (
     NotionPage,
     NotionParentBlock,
@@ -149,7 +151,42 @@ def convert_rich_text(rich_text: NotionRichText) -> dict[str, Any]:
     }
 
 
-def convert_block(block: NotionBlock) -> list[dict[str, Any]]:
+class ImportedAttachment(BaseModel):
+    block: Any
+    file: NotionFileHosted
+
+
+def convert_image(
+    image: NotionImage, attachments: list[ImportedAttachment]
+) -> list[dict[str, Any]]:
+    # TODO: NotionFileUpload
+    match image.file:
+        case NotionFileExternal():
+            return [
+                {
+                    "type": "image",
+                    "props": {
+                        "url": image.file.external["url"],
+                    },
+                }
+            ]
+        case NotionFileHosted():
+            block = {
+                "type": "image",
+                "props": {
+                    "url": "about:blank",  # populated later on
+                },
+            }
+            attachments.append(ImportedAttachment(block=block, file=image.file))
+
+            return [block]
+        case _:
+            return [{"paragraph": {"content": "Unsupported image type"}}]
+
+
+def convert_block(
+    block: NotionBlock, attachments: list[ImportedAttachment]
+) -> list[dict[str, Any]]:
     match block.specific:
         case NotionColumnList():
             columns_content = []
@@ -170,6 +207,8 @@ def convert_block(block: NotionBlock) -> list[dict[str, Any]]:
                     "content": content,
                 }
             ]
+        case NotionImage():
+            return convert_image(block.specific, attachments)
         case NotionHeading1() | NotionHeading2() | NotionHeading3():
             return [
                 {
@@ -328,17 +367,20 @@ def convert_annotations(annotations: NotionRichTextAnnotation) -> dict[str, str]
     return res
 
 
-def convert_block_list(blocks: list[NotionBlock]) -> list[dict[str, Any]]:
+def convert_block_list(
+    blocks: list[NotionBlock], attachments: list[ImportedAttachment]
+) -> list[dict[str, Any]]:
     converted_blocks = []
     for block in blocks:
-        converted_blocks.extend(convert_block(block))
+        converted_blocks.extend(convert_block(block, attachments))
     return converted_blocks
 
 
 class ImportedDocument(BaseModel):
     page: NotionPage
     blocks: list[dict[str, Any]] = []
     children: list["ImportedDocument"] = []
+    attachments: list[ImportedAttachment] = []
 
 
 def find_block_child_page(block_id: str, all_pages: list[NotionPage]):
@@ -386,10 +428,13 @@ def import_page(
     blocks = fetch_block_children(session, page.id)
     logger.info(f"Page {page.get_title()} (id {page.id})")
     logger.info(blocks)
+    attachments = []
+    converted_blocks = convert_block_list(blocks, attachments)
     return ImportedDocument(
         page=page,
-        blocks=convert_block_list(blocks),
+        blocks=converted_blocks,
         children=convert_child_pages(session, page, blocks, all_pages),
+        attachments=attachments,
     )