mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-24 21:38:09 +02:00
feat: Implement notification system with real-time updates and Electric SQL integration
- Added notifications table to the database schema with replication support for Electric SQL. - Developed NotificationService to manage indexing notifications, including creation, updates, and status tracking. - Introduced NotificationButton and NotificationPopup components for displaying notifications in the UI. - Enhanced useNotifications hook for real-time notification syncing using PGlite live queries. - Updated package dependencies for Electric SQL and improved error handling in notification processes.
This commit is contained in:
Anish Sarkar
parent
f441c7b0ce
commit
93d17b51f5
10 changed files with 1062 additions and 103 deletions
|
|
@ -41,12 +41,27 @@ def upgrade() -> None:
|
|||
op.create_index("ix_notifications_user_read", "notifications", ["user_id", "read"])
|
||||
|
||||
# Set REPLICA IDENTITY FULL (required by Electric SQL for replication)
|
||||
# This allows Electric SQL to track all column values for updates/deletes
|
||||
op.execute("ALTER TABLE notifications REPLICA IDENTITY FULL;")
|
||||
|
||||
# Note: ElectricSQL 1.x dynamically adds tables to the publication when
|
||||
# clients subscribe to shapes. No need to manually create publications.
|
||||
# Grant SELECT to electric user for Electric SQL replication
|
||||
# This is needed because ALTER DEFAULT PRIVILEGES only applies during initial DB setup
|
||||
op.execute("GRANT SELECT ON notifications TO electric;")
|
||||
|
||||
# Add notifications table to Electric SQL publication for replication
|
||||
# This is required for Electric SQL to sync the table
|
||||
op.execute("""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_publication_tables
|
||||
WHERE pubname = 'electric_publication_default'
|
||||
AND tablename = 'notifications'
|
||||
) THEN
|
||||
ALTER PUBLICATION electric_publication_default ADD TABLE notifications;
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
""")
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema - remove notifications table."""
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from app.db import (
|
|||
async_session_maker,
|
||||
get_async_session,
|
||||
)
|
||||
from app.services.notification_service import NotificationService
|
||||
from app.schemas import (
|
||||
GoogleDriveIndexRequest,
|
||||
SearchSourceConnectorBase,
|
||||
|
|
@ -973,6 +974,118 @@ async def run_slack_indexing(
|
|||
logger.error(f"Error in background Slack indexing task: {e!s}")
|
||||
|
||||
|
||||
async def _run_indexing_with_notifications(
|
||||
session: AsyncSession,
|
||||
connector_id: int,
|
||||
search_space_id: int,
|
||||
user_id: str,
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
indexing_function,
|
||||
update_timestamp_func=None,
|
||||
):
|
||||
"""
|
||||
Generic helper to run indexing with real-time notifications.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
connector_id: ID of the connector
|
||||
search_space_id: ID of the search space
|
||||
user_id: ID of the user
|
||||
start_date: Start date for indexing
|
||||
end_date: End date for indexing
|
||||
indexing_function: Async function that performs the indexing
|
||||
update_timestamp_func: Optional function to update connector timestamp
|
||||
"""
|
||||
from uuid import UUID
|
||||
|
||||
notification = None
|
||||
try:
|
||||
# Get connector info for notification
|
||||
connector_result = await session.execute(
|
||||
select(SearchSourceConnector).where(SearchSourceConnector.id == connector_id)
|
||||
)
|
||||
connector = connector_result.scalar_one_or_none()
|
||||
|
||||
if connector:
|
||||
# Create notification when indexing starts
|
||||
notification = await NotificationService.connector_indexing.notify_indexing_started(
|
||||
session=session,
|
||||
user_id=UUID(user_id),
|
||||
connector_id=connector_id,
|
||||
connector_name=connector.name,
|
||||
connector_type=connector.connector_type.value,
|
||||
search_space_id=search_space_id,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
)
|
||||
|
||||
# Run the indexing function
|
||||
documents_processed, error_or_warning = await indexing_function(
|
||||
session=session,
|
||||
connector_id=connector_id,
|
||||
search_space_id=search_space_id,
|
||||
user_id=user_id,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
update_last_indexed=False,
|
||||
)
|
||||
|
||||
# Update connector timestamp if function provided and indexing was successful
|
||||
if documents_processed > 0 and update_timestamp_func:
|
||||
await update_timestamp_func(session, connector_id)
|
||||
logger.info(
|
||||
f"Indexing completed successfully: {documents_processed} documents processed"
|
||||
)
|
||||
|
||||
# Update notification on success
|
||||
if notification:
|
||||
await NotificationService.connector_indexing.notify_indexing_completed(
|
||||
session=session,
|
||||
notification=notification,
|
||||
indexed_count=documents_processed,
|
||||
error_message=None,
|
||||
)
|
||||
elif documents_processed > 0:
|
||||
# Success but no timestamp update function
|
||||
logger.info(
|
||||
f"Indexing completed successfully: {documents_processed} documents processed"
|
||||
)
|
||||
if notification:
|
||||
await NotificationService.connector_indexing.notify_indexing_completed(
|
||||
session=session,
|
||||
notification=notification,
|
||||
indexed_count=documents_processed,
|
||||
error_message=None,
|
||||
)
|
||||
else:
|
||||
# Failure or no documents processed
|
||||
logger.error(
|
||||
f"Indexing failed or no documents processed: {error_or_warning}"
|
||||
)
|
||||
if notification:
|
||||
await NotificationService.connector_indexing.notify_indexing_completed(
|
||||
session=session,
|
||||
notification=notification,
|
||||
indexed_count=0,
|
||||
error_message=error_or_warning or "No documents processed",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in indexing task: {e!s}", exc_info=True)
|
||||
|
||||
# Update notification on exception
|
||||
if notification:
|
||||
try:
|
||||
await NotificationService.connector_indexing.notify_indexing_completed(
|
||||
session=session,
|
||||
notification=notification,
|
||||
indexed_count=0,
|
||||
error_message=str(e),
|
||||
)
|
||||
except Exception as notif_error:
|
||||
logger.error(f"Failed to update notification: {notif_error!s}")
|
||||
|
||||
|
||||
async def run_notion_indexing_with_new_session(
|
||||
connector_id: int,
|
||||
search_space_id: int,
|
||||
|
|
@ -985,8 +1098,15 @@ async def run_notion_indexing_with_new_session(
|
|||
This prevents session leaks by creating a dedicated session for the background task.
|
||||
"""
|
||||
async with async_session_maker() as session:
|
||||
await run_notion_indexing(
|
||||
session, connector_id, search_space_id, user_id, start_date, end_date
|
||||
await _run_indexing_with_notifications(
|
||||
session=session,
|
||||
connector_id=connector_id,
|
||||
search_space_id=search_space_id,
|
||||
user_id=user_id,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
indexing_function=index_notion_pages,
|
||||
update_timestamp_func=_update_connector_timestamp_by_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1009,30 +1129,16 @@ async def run_notion_indexing(
|
|||
start_date: Start date for indexing
|
||||
end_date: End date for indexing
|
||||
"""
|
||||
try:
|
||||
# Index Notion pages without updating last_indexed_at (we'll do it separately)
|
||||
documents_processed, error_or_warning = await index_notion_pages(
|
||||
session=session,
|
||||
connector_id=connector_id,
|
||||
search_space_id=search_space_id,
|
||||
user_id=user_id,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
update_last_indexed=False, # Don't update timestamp in the indexing function
|
||||
)
|
||||
|
||||
# Only update last_indexed_at if indexing was successful (either new docs or updated docs)
|
||||
if documents_processed > 0:
|
||||
await _update_connector_timestamp_by_id(session, connector_id)
|
||||
logger.info(
|
||||
f"Notion indexing completed successfully: {documents_processed} documents processed"
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"Notion indexing failed or no documents processed: {error_or_warning}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in background Notion indexing task: {e!s}")
|
||||
await _run_indexing_with_notifications(
|
||||
session=session,
|
||||
connector_id=connector_id,
|
||||
search_space_id=search_space_id,
|
||||
user_id=user_id,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
indexing_function=index_notion_pages,
|
||||
update_timestamp_func=_update_connector_timestamp_by_id,
|
||||
)
|
||||
|
||||
|
||||
# Add new helper functions for GitHub indexing
|
||||
|
|
|
|||
|
|
@ -1,18 +1,343 @@
|
|||
"""Service for creating and managing notifications."""
|
||||
"""Service for creating and managing notifications with Electric SQL sync."""
|
||||
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm.attributes import flag_modified
|
||||
|
||||
from app.db import Notification
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseNotificationHandler:
|
||||
"""Base class for notification handlers - provides common functionality."""
|
||||
|
||||
def __init__(self, notification_type: str):
|
||||
"""
|
||||
Initialize the notification handler.
|
||||
|
||||
Args:
|
||||
notification_type: Type of notification (e.g., 'connector_indexing', 'document_processing')
|
||||
"""
|
||||
self.notification_type = notification_type
|
||||
|
||||
async def find_notification_by_operation(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
user_id: UUID,
|
||||
operation_id: str,
|
||||
search_space_id: int | None = None,
|
||||
) -> Notification | None:
|
||||
"""
|
||||
Find an existing notification by operation ID.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
user_id: User ID
|
||||
operation_id: Unique operation identifier
|
||||
search_space_id: Optional search space ID
|
||||
|
||||
Returns:
|
||||
Notification if found, None otherwise
|
||||
"""
|
||||
query = select(Notification).where(
|
||||
Notification.user_id == user_id,
|
||||
Notification.type == self.notification_type,
|
||||
Notification.notification_metadata["operation_id"].astext == operation_id,
|
||||
)
|
||||
if search_space_id is not None:
|
||||
query = query.where(Notification.search_space_id == search_space_id)
|
||||
|
||||
result = await session.execute(query)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def find_or_create_notification(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
user_id: UUID,
|
||||
operation_id: str,
|
||||
title: str,
|
||||
message: str,
|
||||
search_space_id: int | None = None,
|
||||
initial_metadata: dict[str, Any] | None = None,
|
||||
) -> Notification:
|
||||
"""
|
||||
Find an existing notification or create a new one.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
user_id: User ID
|
||||
operation_id: Unique operation identifier
|
||||
title: Notification title
|
||||
message: Notification message
|
||||
search_space_id: Optional search space ID
|
||||
initial_metadata: Initial metadata dictionary
|
||||
|
||||
Returns:
|
||||
Notification: The found or created notification
|
||||
"""
|
||||
# Try to find existing notification
|
||||
notification = await self.find_notification_by_operation(
|
||||
session, user_id, operation_id, search_space_id
|
||||
)
|
||||
|
||||
if notification:
|
||||
# Update existing notification
|
||||
notification.title = title
|
||||
notification.message = message
|
||||
if initial_metadata:
|
||||
notification.notification_metadata = {
|
||||
**notification.notification_metadata,
|
||||
**initial_metadata,
|
||||
}
|
||||
# Mark JSONB column as modified so SQLAlchemy detects the change
|
||||
flag_modified(notification, "notification_metadata")
|
||||
await session.commit()
|
||||
await session.refresh(notification)
|
||||
logger.info(f"Updated notification {notification.id} for operation {operation_id}")
|
||||
return notification
|
||||
|
||||
# Create new notification
|
||||
metadata = initial_metadata or {}
|
||||
metadata["operation_id"] = operation_id
|
||||
metadata["status"] = "in_progress"
|
||||
metadata["started_at"] = datetime.now(UTC).isoformat()
|
||||
|
||||
notification = Notification(
|
||||
user_id=user_id,
|
||||
search_space_id=search_space_id,
|
||||
type=self.notification_type,
|
||||
title=title,
|
||||
message=message,
|
||||
notification_metadata=metadata,
|
||||
)
|
||||
session.add(notification)
|
||||
await session.commit()
|
||||
await session.refresh(notification)
|
||||
logger.info(f"Created notification {notification.id} for operation {operation_id}")
|
||||
return notification
|
||||
|
||||
async def update_notification(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
notification: Notification,
|
||||
title: str | None = None,
|
||||
message: str | None = None,
|
||||
status: str | None = None,
|
||||
metadata_updates: dict[str, Any] | None = None,
|
||||
) -> Notification:
|
||||
"""
|
||||
Update an existing notification.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
notification: Notification to update
|
||||
title: New title (optional)
|
||||
message: New message (optional)
|
||||
status: New status (optional)
|
||||
metadata_updates: Additional metadata to merge (optional)
|
||||
|
||||
Returns:
|
||||
Updated notification
|
||||
"""
|
||||
if title is not None:
|
||||
notification.title = title
|
||||
if message is not None:
|
||||
notification.message = message
|
||||
|
||||
if status is not None:
|
||||
notification.notification_metadata["status"] = status
|
||||
if status in ("completed", "failed"):
|
||||
notification.notification_metadata["completed_at"] = (
|
||||
datetime.now(UTC).isoformat()
|
||||
)
|
||||
# Mark JSONB column as modified so SQLAlchemy detects the change
|
||||
flag_modified(notification, "notification_metadata")
|
||||
|
||||
if metadata_updates:
|
||||
notification.notification_metadata = {
|
||||
**notification.notification_metadata,
|
||||
**metadata_updates,
|
||||
}
|
||||
# Mark JSONB column as modified
|
||||
flag_modified(notification, "notification_metadata")
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(notification)
|
||||
logger.info(f"Updated notification {notification.id}")
|
||||
return notification
|
||||
|
||||
|
||||
class ConnectorIndexingNotificationHandler(BaseNotificationHandler):
|
||||
"""Handler for connector indexing notifications."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__("connector_indexing")
|
||||
|
||||
def _generate_operation_id(
|
||||
self, connector_id: int, start_date: str | None = None, end_date: str | None = None
|
||||
) -> str:
|
||||
"""
|
||||
Generate a unique operation ID for a connector indexing operation.
|
||||
|
||||
Args:
|
||||
connector_id: Connector ID
|
||||
start_date: Start date (optional)
|
||||
end_date: End date (optional)
|
||||
|
||||
Returns:
|
||||
Unique operation ID string
|
||||
"""
|
||||
timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
|
||||
date_range = ""
|
||||
if start_date or end_date:
|
||||
date_range = f"_{start_date or 'none'}_{end_date or 'none'}"
|
||||
return f"connector_{connector_id}_{timestamp}{date_range}"
|
||||
|
||||
async def notify_indexing_started(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
user_id: UUID,
|
||||
connector_id: int,
|
||||
connector_name: str,
|
||||
connector_type: str,
|
||||
search_space_id: int,
|
||||
start_date: str | None = None,
|
||||
end_date: str | None = None,
|
||||
) -> Notification:
|
||||
"""
|
||||
Create or update notification when connector indexing starts.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
user_id: User ID
|
||||
connector_id: Connector ID
|
||||
connector_name: Connector name
|
||||
connector_type: Connector type
|
||||
search_space_id: Search space ID
|
||||
start_date: Start date for indexing
|
||||
end_date: End date for indexing
|
||||
|
||||
Returns:
|
||||
Notification: The created or updated notification
|
||||
"""
|
||||
operation_id = self._generate_operation_id(connector_id, start_date, end_date)
|
||||
title = f"Indexing: {connector_name}"
|
||||
message = f'Indexing "{connector_name}" in progress...'
|
||||
|
||||
metadata = {
|
||||
"connector_id": connector_id,
|
||||
"connector_name": connector_name,
|
||||
"connector_type": connector_type,
|
||||
"start_date": start_date,
|
||||
"end_date": end_date,
|
||||
"indexed_count": 0,
|
||||
}
|
||||
|
||||
return await self.find_or_create_notification(
|
||||
session=session,
|
||||
user_id=user_id,
|
||||
operation_id=operation_id,
|
||||
title=title,
|
||||
message=message,
|
||||
search_space_id=search_space_id,
|
||||
initial_metadata=metadata,
|
||||
)
|
||||
|
||||
async def notify_indexing_progress(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
notification: Notification,
|
||||
indexed_count: int,
|
||||
total_count: int | None = None,
|
||||
) -> Notification:
|
||||
"""
|
||||
Update notification with indexing progress.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
notification: Notification to update
|
||||
indexed_count: Number of items indexed so far
|
||||
total_count: Total number of items (optional)
|
||||
|
||||
Returns:
|
||||
Updated notification
|
||||
"""
|
||||
connector_name = notification.notification_metadata.get("connector_name", "Connector")
|
||||
progress_msg = f'Indexing "{connector_name}": {indexed_count} items'
|
||||
if total_count is not None:
|
||||
progress_msg += f" of {total_count}"
|
||||
progress_msg += " indexed..."
|
||||
|
||||
metadata_updates = {"indexed_count": indexed_count}
|
||||
if total_count is not None:
|
||||
metadata_updates["total_count"] = total_count
|
||||
progress_percent = int((indexed_count / total_count) * 100)
|
||||
metadata_updates["progress_percent"] = progress_percent
|
||||
|
||||
return await self.update_notification(
|
||||
session=session,
|
||||
notification=notification,
|
||||
message=progress_msg,
|
||||
status="in_progress",
|
||||
metadata_updates=metadata_updates,
|
||||
)
|
||||
|
||||
async def notify_indexing_completed(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
notification: Notification,
|
||||
indexed_count: int,
|
||||
error_message: str | None = None,
|
||||
) -> Notification:
|
||||
"""
|
||||
Update notification when connector indexing completes.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
notification: Notification to update
|
||||
indexed_count: Total number of items indexed
|
||||
error_message: Error message if indexing failed (optional)
|
||||
|
||||
Returns:
|
||||
Updated notification
|
||||
"""
|
||||
connector_name = notification.notification_metadata.get("connector_name", "Connector")
|
||||
|
||||
if error_message:
|
||||
title = f"Indexing failed: {connector_name}"
|
||||
message = f'Indexing "{connector_name}" failed: {error_message}'
|
||||
status = "failed"
|
||||
else:
|
||||
title = f"Indexing completed: {connector_name}"
|
||||
message = f'Indexing "{connector_name}" completed successfully. {indexed_count} items indexed.'
|
||||
status = "completed"
|
||||
|
||||
metadata_updates = {
|
||||
"indexed_count": indexed_count,
|
||||
"error_message": error_message,
|
||||
}
|
||||
|
||||
return await self.update_notification(
|
||||
session=session,
|
||||
notification=notification,
|
||||
title=title,
|
||||
message=message,
|
||||
status=status,
|
||||
metadata_updates=metadata_updates,
|
||||
)
|
||||
|
||||
|
||||
class NotificationService:
|
||||
"""Service for creating notifications that sync via Electric SQL."""
|
||||
"""Service for creating and managing notifications that sync via Electric SQL."""
|
||||
|
||||
# Handler instances
|
||||
connector_indexing = ConnectorIndexingNotificationHandler()
|
||||
|
||||
@staticmethod
|
||||
async def create_notification(
|
||||
|
|
@ -105,6 +430,7 @@ class NotificationService:
|
|||
) -> Notification:
|
||||
"""
|
||||
Create notification when connector indexing completes.
|
||||
DEPRECATED: Use NotificationService.connector_indexing methods instead.
|
||||
|
||||
Args:
|
||||
session: Database session
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue