39
39
from core .services .ai_services import AIService
40
40
from core .services .collaboration_services import CollaborationService
41
41
from core .services .converter_services import YdocConverter
42
- from core .services .notion_import import import_notion
42
+ from core .services .notion_import import (
43
+ ImportedDocument ,
44
+ build_notion_session ,
45
+ fetch_all_pages ,
46
+ import_page ,
47
+ link_child_page_to_parent ,
48
+ )
43
49
from core .tasks .mail import send_ask_for_access_mail
44
50
from core .utils import extract_attachments , filter_descendants
45
51
52
+ from ..notion_schemas .notion_page import NotionPage
46
53
from . import permissions , serializers , utils
47
54
from .filters import DocumentFilter , ListDocumentFilter
48
55
@@ -2134,7 +2141,7 @@ def _import_notion_doc_content(imported_doc, obj, user):
2134
2141
obj .save ()
2135
2142
2136
2143
2137
- def _import_notion_child_page (imported_doc , parent_doc , user , imported_docs_by_page_id ):
2144
+ def _import_notion_child_page (imported_doc , parent_doc , user , imported_ids ):
2138
2145
obj = parent_doc .add_child (
2139
2146
creator = user ,
2140
2147
title = imported_doc .page .get_title () or "Child page" ,
@@ -2148,13 +2155,13 @@ def _import_notion_child_page(imported_doc, parent_doc, user, imported_docs_by_p
2148
2155
2149
2156
_import_notion_doc_content (imported_doc , obj , user )
2150
2157
2151
- imported_docs_by_page_id [ imported_doc .page .id ] = obj
2158
+ imported_ids . append ( imported_doc .page .id )
2152
2159
2153
2160
for child in imported_doc .children :
2154
- _import_notion_child_page (child , obj , user , imported_docs_by_page_id )
2161
+ _import_notion_child_page (child , obj , user , imported_ids )
2155
2162
2156
2163
2157
- def _import_notion_root_page (imported_doc , user , imported_docs_by_page_id ) :
2164
+ def _import_notion_root_page (imported_doc , user ) -> list [ str ] :
2158
2165
obj = models .Document .add_root (
2159
2166
depth = 1 ,
2160
2167
creator = user ,
@@ -2168,23 +2175,81 @@ def _import_notion_root_page(imported_doc, user, imported_docs_by_page_id):
2168
2175
role = models .RoleChoices .OWNER ,
2169
2176
)
2170
2177
2171
- _import_notion_doc_content ( imported_doc , obj , user )
2178
+ imported_ids = [ imported_doc . page . id ]
2172
2179
2173
- imported_docs_by_page_id [ imported_doc . page . id ] = obj
2180
+ _import_notion_doc_content ( imported_doc , obj , user )
2174
2181
2175
2182
for child in imported_doc .children :
2176
- _import_notion_child_page (child , obj , user , imported_docs_by_page_id )
2183
+ _import_notion_child_page (child , obj , user , imported_ids )
2177
2184
2185
+ return imported_ids
2178
2186
2179
- @drf .decorators .api_view (["GET" , "POST" ]) # TODO: drop GET (used for testing)
2180
- def notion_import_run (request ):
2181
- if "notion_token" not in request .session :
2182
- raise drf .exceptions .PermissionDenied ()
2183
2187
2184
- imported_docs = import_notion (request .session ["notion_token" ])
2188
+ def _generate_notion_progress (
2189
+ all_pages : list [NotionPage ], page_statuses : dict [str , str ]
2190
+ ) -> str :
2191
+ raw = json .dumps (
2192
+ [
2193
+ {
2194
+ "title" : page .get_title (),
2195
+ "status" : page_statuses [page .id ],
2196
+ }
2197
+ for page in all_pages
2198
+ ]
2199
+ )
2200
+ return f"data: { raw } \n \n "
2201
+
2185
2202
2186
- imported_docs_by_page_id = {}
2187
- for imported_doc in imported_docs :
2188
- _import_notion_root_page ( imported_doc , request . user , imported_docs_by_page_id )
2203
+ def _notion_import_event_stream ( request ):
2204
+ session = build_notion_session ( request . session [ "notion_token" ])
2205
+ all_pages = fetch_all_pages ( session )
2189
2206
2190
- return drf .response .Response ({})
2207
+ page_statuses = {}
2208
+ for page in all_pages :
2209
+ page_statuses [page .id ] = "pending"
2210
+
2211
+ yield _generate_notion_progress (all_pages , page_statuses )
2212
+
2213
+ docs_by_page_id : dict [str , ImportedDocument ] = {}
2214
+ child_page_blocs_ids_to_parent_page_ids : dict [str , str ] = {}
2215
+
2216
+ for page in all_pages :
2217
+ docs_by_page_id [page .id ] = import_page (
2218
+ session , page , child_page_blocs_ids_to_parent_page_ids
2219
+ )
2220
+ page_statuses [page .id ] = "fetched"
2221
+ yield _generate_notion_progress (all_pages , page_statuses )
2222
+
2223
+ for page in all_pages :
2224
+ link_child_page_to_parent (
2225
+ page , docs_by_page_id , child_page_blocs_ids_to_parent_page_ids
2226
+ )
2227
+
2228
+ root_docs = [doc for doc in docs_by_page_id .values () if doc .page .is_root ()]
2229
+
2230
+ for root_doc in root_docs :
2231
+ imported_ids = _import_notion_root_page (root_doc , request .user )
2232
+ for imported_id in imported_ids :
2233
+ page_statuses [imported_id ] = "imported"
2234
+
2235
+ yield _generate_notion_progress (all_pages , page_statuses )
2236
+
2237
+
2238
+ class IgnoreClientContentNegotiation (drf .negotiation .BaseContentNegotiation ):
2239
+ def select_parser (self , request , parsers ):
2240
+ return parsers [0 ]
2241
+
2242
+ def select_renderer (self , request , renderers , format_suffix ):
2243
+ return (renderers [0 ], renderers [0 ].media_type )
2244
+
2245
+
2246
+ class NotionImportRunView (drf .views .APIView ):
2247
+ content_negotiation_class = IgnoreClientContentNegotiation
2248
+
2249
+ def get (self , request , format = None ):
2250
+ if "notion_token" not in request .session :
2251
+ raise drf .exceptions .PermissionDenied ()
2252
+
2253
+ return StreamingHttpResponse (
2254
+ _notion_import_event_stream (request ), content_type = "text/event-stream"
2255
+ )
0 commit comments