Skip to content

Commit 42f4bb7

Browse files
chloediaZewedStanGirardAmineDiro
authored
feat(integration): Notion (#3173)
# Description Fix multiple notion bugs 👍 -> Delete your notion sync and all the notion files from the db -> Ensure a sync is not already running before launching a sync. -> Add a status to subscribe to for user_sync --------- Co-authored-by: Antoine Dewez <[email protected]> Co-authored-by: Stan Girard <[email protected]> Co-authored-by: aminediro <[email protected]> Co-authored-by: Stan Girard <[email protected]>
1 parent 9c6d998 commit 42f4bb7

36 files changed

+755
-237
lines changed

‎backend/api/quivr_api/modules/knowledge/tests/test_knowledge_service.py

+1
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,7 @@ async def test_should_process_knowledge_prev_error(
527527
assert new.file_sha1
528528

529529

530+
@pytest.mark.skip(reason="Bug: UnboundLocalError: cannot access local variable 'response'")
530531
@pytest.mark.asyncio(loop_scope="session")
531532
async def test_get_knowledge_storage_path(session: AsyncSession, test_data: TestData):
532533
_, [knowledge, _] = test_data

‎backend/api/quivr_api/modules/rag_service/utils.py

+19-15
Original file line numberDiff line numberDiff line change
@@ -65,24 +65,28 @@ async def generate_source(
6565
source_url = doc.metadata["original_file_name"]
6666
else:
6767
# Check if the URL has already been generated
68-
file_name = doc.metadata["file_name"]
69-
file_path = await knowledge_service.get_knowledge_storage_path(
68+
try:
69+
file_name = doc.metadata["file_name"]
70+
file_path = await knowledge_service.get_knowledge_storage_path(
7071
file_name=file_name, brain_id=brain_id
71-
)
72-
if file_path in generated_urls:
73-
source_url = generated_urls[file_path]
74-
else:
75-
# Generate the URL
76-
if file_path in sources_url_cache:
77-
source_url = sources_url_cache[file_path]
72+
)
73+
if file_path in generated_urls:
74+
source_url = generated_urls[file_path]
7875
else:
79-
generated_url = generate_file_signed_url(file_path)
80-
if generated_url is not None:
81-
source_url = generated_url.get("signedURL", "")
76+
# Generate the URL
77+
if file_path in sources_url_cache:
78+
source_url = sources_url_cache[file_path]
8279
else:
83-
source_url = ""
84-
# Store the generated URL
85-
generated_urls[file_path] = source_url
80+
generated_url = generate_file_signed_url(file_path)
81+
if generated_url is not None:
82+
source_url = generated_url.get("signedURL", "")
83+
else:
84+
source_url = ""
85+
# Store the generated URL
86+
generated_urls[file_path] = source_url
87+
except Exception as e:
88+
logger.error(f"Error generating file signed URL: {e}")
89+
continue
8690

8791
# Append a new Sources object to the list
8892
sources_list.append(

‎backend/api/quivr_api/modules/sync/controller/azure_sync_routes.py

+9-2
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,11 @@
77

88
from quivr_api.logger import get_logger
99
from quivr_api.middlewares.auth import AuthBearer, get_current_user
10-
from quivr_api.modules.sync.dto.inputs import SyncsUserInput, SyncUserUpdateInput
10+
from quivr_api.modules.sync.dto.inputs import (
11+
SyncsUserInput,
12+
SyncsUserStatus,
13+
SyncUserUpdateInput,
14+
)
1115
from quivr_api.modules.sync.service.sync_service import SyncService, SyncUserService
1216
from quivr_api.modules.user.entity.user_identity import UserIdentity
1317

@@ -70,6 +74,7 @@ def authorize_azure(
7074
credentials={},
7175
state={"state": state},
7276
additional_data={"flow": flow},
77+
status=str(SyncsUserStatus.SYNCING),
7378
)
7479
sync_user_service.create_sync_user(sync_user_input)
7580
return {"authorization_url": flow["auth_uri"]}
@@ -138,7 +143,9 @@ def oauth2callback_azure(request: Request):
138143
logger.info(f"Retrieved email for user: {current_user} - {user_email}")
139144

140145
sync_user_input = SyncUserUpdateInput(
141-
credentials=result, state={}, email=user_email
146+
credentials=result,
147+
email=user_email,
148+
status=str(SyncsUserStatus.SYNCED),
142149
)
143150

144151
sync_user_service.update_sync_user(current_user, state_dict, sync_user_input)

‎backend/api/quivr_api/modules/sync/controller/dropbox_sync_routes.py

+9-2
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,11 @@
77

88
from quivr_api.logger import get_logger
99
from quivr_api.middlewares.auth import AuthBearer, get_current_user
10-
from quivr_api.modules.sync.dto.inputs import SyncsUserInput, SyncUserUpdateInput
10+
from quivr_api.modules.sync.dto.inputs import (
11+
SyncsUserInput,
12+
SyncsUserStatus,
13+
SyncUserUpdateInput,
14+
)
1115
from quivr_api.modules.sync.service.sync_service import SyncService, SyncUserService
1216
from quivr_api.modules.user.entity.user_identity import UserIdentity
1317

@@ -72,6 +76,7 @@ def authorize_dropbox(
7276
credentials={},
7377
state={"state": state},
7478
additional_data={},
79+
status=str(SyncsUserStatus.SYNCING),
7580
)
7681
sync_user_service.create_sync_user(sync_user_input)
7782
return {"authorization_url": authorize_url}
@@ -147,9 +152,11 @@ def oauth2callback_dropbox(request: Request):
147152

148153
sync_user_input = SyncUserUpdateInput(
149154
credentials=result,
150-
state={},
155+
# state={},
151156
email=user_email,
157+
status=str(SyncsUserStatus.SYNCED),
152158
)
159+
assert current_user
153160
sync_user_service.update_sync_user(current_user, state_dict, sync_user_input)
154161
logger.info(f"DropBox sync created successfully for user: {current_user}")
155162
return HTMLResponse(successfullConnectionPage)

‎backend/api/quivr_api/modules/sync/controller/github_sync_routes.py

+10-2
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,11 @@
66

77
from quivr_api.logger import get_logger
88
from quivr_api.middlewares.auth import AuthBearer, get_current_user
9-
from quivr_api.modules.sync.dto.inputs import SyncsUserInput, SyncUserUpdateInput
9+
from quivr_api.modules.sync.dto.inputs import (
10+
SyncsUserInput,
11+
SyncsUserStatus,
12+
SyncUserUpdateInput,
13+
)
1014
from quivr_api.modules.sync.service.sync_service import SyncService, SyncUserService
1115
from quivr_api.modules.user.entity.user_identity import UserIdentity
1216

@@ -61,6 +65,7 @@ def authorize_github(
6165
provider="GitHub",
6266
credentials={},
6367
state={"state": state},
68+
status=str(SyncsUserStatus.SYNCING),
6469
)
6570
sync_user_service.create_sync_user(sync_user_input)
6671
return {"authorization_url": authorization_url}
@@ -148,7 +153,10 @@ def oauth2callback_github(request: Request):
148153
logger.info(f"Retrieved email for user: {current_user} - {user_email}")
149154

150155
sync_user_input = SyncUserUpdateInput(
151-
credentials=result, state={}, email=user_email
156+
credentials=result,
157+
# state={},
158+
email=user_email,
159+
status=str(SyncsUserStatus.SYNCED),
152160
)
153161

154162
sync_user_service.update_sync_user(current_user, state_dict, sync_user_input)

‎backend/api/quivr_api/modules/sync/controller/google_sync_routes.py

+8-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,11 @@
99

1010
from quivr_api.logger import get_logger
1111
from quivr_api.middlewares.auth import AuthBearer, get_current_user
12-
from quivr_api.modules.sync.dto.inputs import SyncsUserInput, SyncUserUpdateInput
12+
from quivr_api.modules.sync.dto.inputs import (
13+
SyncsUserInput,
14+
SyncsUserStatus,
15+
SyncUserUpdateInput,
16+
)
1317
from quivr_api.modules.sync.service.sync_service import SyncService, SyncUserService
1418
from quivr_api.modules.user.entity.user_identity import UserIdentity
1519

@@ -101,6 +105,7 @@ def authorize_google(
101105
credentials={},
102106
state={"state": state},
103107
additional_data={},
108+
status=str(SyncsUserStatus.SYNCED),
104109
)
105110
sync_user_service.create_sync_user(sync_user_input)
106111
return {"authorization_url": authorization_url}
@@ -156,8 +161,9 @@ def oauth2callback_google(request: Request):
156161

157162
sync_user_input = SyncUserUpdateInput(
158163
credentials=json.loads(creds.to_json()),
159-
state={},
164+
# state={},
160165
email=user_email,
166+
status=str(SyncsUserStatus.SYNCED),
161167
)
162168
sync_user_service.update_sync_user(current_user, state_dict, sync_user_input)
163169
logger.info(f"Google Drive sync created successfully for user: {current_user}")

‎backend/api/quivr_api/modules/sync/controller/notion_sync_routes.py

+13-3
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,11 @@
1010
from quivr_api.celery_config import celery
1111
from quivr_api.logger import get_logger
1212
from quivr_api.middlewares.auth import AuthBearer, get_current_user
13-
from quivr_api.modules.sync.dto.inputs import SyncsUserInput, SyncUserUpdateInput
13+
from quivr_api.modules.sync.dto.inputs import (
14+
SyncsUserInput,
15+
SyncsUserStatus,
16+
SyncUserUpdateInput,
17+
)
1418
from quivr_api.modules.sync.service.sync_service import SyncService, SyncUserService
1519
from quivr_api.modules.user.entity.user_identity import UserIdentity
1620

@@ -65,6 +69,7 @@ def authorize_notion(
6569
provider="Notion",
6670
credentials={},
6771
state={"state": state},
72+
status=str(SyncsUserStatus.SYNCING),
6873
)
6974
sync_user_service.create_sync_user(sync_user_input)
7075
return {"authorization_url": authorize_url}
@@ -145,15 +150,20 @@ def oauth2callback_notion(request: Request, background_tasks: BackgroundTasks):
145150

146151
sync_user_input = SyncUserUpdateInput(
147152
credentials=result,
148-
state={},
153+
# state={},
149154
email=user_email,
155+
status=str(SyncsUserStatus.SYNCING),
150156
)
151157
sync_user_service.update_sync_user(current_user, state_dict, sync_user_input)
152158
logger.info(f"Notion sync created successfully for user: {current_user}")
153159
# launch celery task to sync notion data
154160
celery.send_task(
155161
"fetch_and_store_notion_files_task",
156-
kwargs={"access_token": access_token, "user_id": current_user},
162+
kwargs={
163+
"access_token": access_token,
164+
"user_id": current_user,
165+
"sync_user_id": sync_user_state.id,
166+
},
157167
)
158168
return HTMLResponse(successfullConnectionPage)
159169

‎backend/api/quivr_api/modules/sync/dto/inputs.py

+19-1
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,23 @@
1+
import enum
12
from typing import List, Optional
23

34
from pydantic import BaseModel
45

56

7+
class SyncsUserStatus(enum.Enum):
8+
"""
9+
Enum for the status of a sync user.
10+
"""
11+
12+
SYNCED = "SYNCED"
13+
SYNCING = "SYNCING"
14+
ERROR = "ERROR"
15+
REMOVED = "REMOVED"
16+
17+
def __str__(self):
18+
return self.value
19+
20+
621
class SyncsUserInput(BaseModel):
722
"""
823
Input model for creating a new sync user.
@@ -17,10 +32,12 @@ class SyncsUserInput(BaseModel):
1732

1833
user_id: str
1934
name: str
35+
email: str | None = None
2036
provider: str
2137
credentials: dict
2238
state: dict
2339
additional_data: dict = {}
40+
status: str
2441

2542

2643
class SyncUserUpdateInput(BaseModel):
@@ -33,8 +50,9 @@ class SyncUserUpdateInput(BaseModel):
3350
"""
3451

3552
credentials: dict
36-
state: dict
53+
state: dict | None = None
3754
email: str
55+
status: str
3856

3957

4058
class SyncActiveSettings(BaseModel):

‎backend/api/quivr_api/modules/sync/entity/notion_page.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ class NotionPage(BaseModel):
9797
cover: Cover | None
9898
icon: Icon | None
9999
properties: PageProps
100+
sync_user_id: UUID | None = Field(default=None, foreign_key="syncs_user.id") # type: ignore
100101

101102
# TODO: Fix UUID in table NOTION
102103
def _get_parent_id(self) -> UUID | None:
@@ -110,7 +111,7 @@ def _get_parent_id(self) -> UUID | None:
110111
case BlockParent():
111112
return None
112113

113-
def to_syncfile(self, user_id: UUID):
114+
def to_syncfile(self, user_id: UUID, sync_user_id: int) -> NotionSyncFile:
114115
name = (
115116
self.properties.title.title[0].text.content if self.properties.title else ""
116117
)
@@ -125,6 +126,7 @@ def to_syncfile(self, user_id: UUID):
125126
last_modified=self.last_edited_time,
126127
type="page",
127128
user_id=user_id,
129+
sync_user_id=sync_user_id,
128130
)
129131

130132

‎backend/api/quivr_api/modules/sync/entity/sync_models.py

+6
Original file line numberDiff line numberDiff line change
@@ -54,10 +54,12 @@ class SyncsUser(BaseModel):
5454
id: int
5555
user_id: UUID
5656
name: str
57+
email: str | None = None
5758
provider: str
5859
credentials: dict
5960
state: dict
6061
additional_data: dict
62+
status: str
6163

6264

6365
class SyncsActive(BaseModel):
@@ -114,3 +116,7 @@ class NotionSyncFile(SQLModel, table=True):
114116
description="The ID of the user who owns the file",
115117
)
116118
user: User = Relationship(back_populates="notion_syncs")
119+
sync_user_id: int = Field(
120+
# foreign_key="syncs_user.id",
121+
description="The ID of the sync user associated with the file",
122+
)

‎backend/api/quivr_api/modules/sync/repository/sync_files.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
from quivr_api.logger import get_logger
22
from quivr_api.modules.dependencies import get_supabase_client
3-
from quivr_api.modules.sync.dto.inputs import SyncFileInput, SyncFileUpdateInput
3+
from quivr_api.modules.sync.dto.inputs import (
4+
SyncFileInput,
5+
SyncFileUpdateInput,
6+
)
47
from quivr_api.modules.sync.entity.sync_models import DBSyncFile, SyncFile, SyncsActive
58
from quivr_api.modules.sync.repository.sync_interfaces import SyncFileInterface
69

‎backend/api/quivr_api/modules/sync/repository/sync_repository.py

+13-4
Original file line numberDiff line numberDiff line change
@@ -212,8 +212,13 @@ def __init__(self, session: AsyncSession):
212212
self.session = session
213213
self.db = get_supabase_client()
214214

215-
async def get_user_notion_files(self, user_id: UUID) -> Sequence[NotionSyncFile]:
216-
query = select(NotionSyncFile).where(NotionSyncFile.user_id == user_id)
215+
async def get_user_notion_files(
216+
self, user_id: UUID, sync_user_id: int
217+
) -> Sequence[NotionSyncFile]:
218+
query = select(NotionSyncFile).where(
219+
NotionSyncFile.user_id == user_id
220+
and NotionSyncFile.sync_user_id == sync_user_id
221+
)
217222
response = await self.session.exec(query)
218223
return response.all()
219224

@@ -275,9 +280,13 @@ async def get_notion_files_by_ids(self, ids: List[str]) -> Sequence[NotionSyncFi
275280
return response.all()
276281

277282
async def get_notion_files_by_parent_id(
278-
self, parent_id: str | None
283+
self, parent_id: str | None, sync_user_id: int
279284
) -> Sequence[NotionSyncFile]:
280-
query = select(NotionSyncFile).where(NotionSyncFile.parent_id == parent_id)
285+
query = (
286+
select(NotionSyncFile)
287+
.where(NotionSyncFile.parent_id == parent_id)
288+
.where(NotionSyncFile.sync_user_id == sync_user_id)
289+
)
281290
response = await self.session.exec(query)
282291
return response.all()
283292

0 commit comments

Comments
 (0)