DB conn optimizations

changelog.md

@@ -2,6 +2,17 @@
 
 <!-- towncrier release notes start -->
 
+## [3.0.1] - 2024-01-xx
+
+### DB Connection Optimizations
+
+### Bug fixes
+
+## [3.0] - 2023-12-28
+
+Brand new Papermerge, version 3.0 breaks compatibility with previous
+versions
+
 ## [2.1.9] - 2023-04-01
 
 ### Fixed

papermerge/core/auth/__init__.py

@@ -1,27 +1,104 @@
+import logging
+import os
 
-from django.contrib.auth import get_user_model
-from rest_framework.permissions import DjangoModelPermissions
-
-# custom user is used - papermerge.core.models.User
-User = get_user_model()
-
-
-class CustomModelPermissions(DjangoModelPermissions):
-    """
-    The request is authenticated using `django.contrib.auth` permissions.
-    See: https://docs.djangoproject.com/en/dev/topics/auth/#permissions
-
-    It ensures that the user is authenticated, and has the appropriate
-    `view`/`add`/`change`/`delete` permissions on the model.
-    """
-    # Overrides perms_map of `DjangoModelPermissions` with value for 'GET',
-    # 'HEAD' and 'OPTIONS' keys
-    perms_map = {
-        'GET': ['%(app_label)s.view_%(model_name)s'],
-        'OPTIONS': ['%(app_label)s.view_%(model_name)s'],
-        'HEAD': ['%(app_label)s.view_%(model_name)s'],
-        'POST': ['%(app_label)s.add_%(model_name)s'],
-        'PUT': ['%(app_label)s.change_%(model_name)s'],
-        'PATCH': ['%(app_label)s.change_%(model_name)s'],
-        'DELETE': ['%(app_label)s.delete_%(model_name)s'],
-    }
+from fastapi import (Depends, Header, HTTPException, WebSocket,
+                     WebSocketException, status)
+from fastapi.security import OAuth2PasswordBearer
+from sqlalchemy import Engine
+
+from papermerge.core import db, schemas
+from papermerge.core.db import exceptions as db_exc
+from papermerge.core.utils import base64
+
+oauth2_scheme = OAuth2PasswordBearer(
+    tokenUrl="auth/token/",
+    auto_error=False
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_user_id_from_token(token: str = Depends(oauth2_scheme)) -> str | None:
+    if '.' in token:
+        _, payload, _ = token.split('.')
+        data = base64.decode(payload)
+        user_id = data.get("user_id")
+
+        return user_id
+
+    return None
+
+
+def get_current_user(
+    x_remote_user: str | None = Header(default=None),
+    token: str | None = Depends(oauth2_scheme),
+    engine: Engine = Depends(db.get_engine)
+) -> schemas.User:
+
+    if token:  # token found
+        user_id = get_user_id_from_token(token)
+        if user_id is not None:
+            try:
+                user = db.get_user(engine, user_id)
+            except db_exc.UserNotFound:
+                raise HTTPException(
+                    status_code=401,
+                    detail="User ID not found"
+                )
+    elif x_remote_user:  # get user from X_REMOTE_USER header
+        logger.debug(f"x_remote_user={x_remote_user}")
+        user = db.get_user(engine, x_remote_user)
+
+    remote_user_env_var = os.environ.get("REMOTE_USER")
+
+    if user is None and remote_user_env_var:
+        user = db.get_user(engine, remote_user_env_var)
+
+    if user is None:
+        raise HTTPException(
+            status_code=401,
+            detail="No credentials provided"
+        )
+
+    return user
+
+
+def get_ws_current_user(
+    websocket: WebSocket,
+    engine: Engine = Depends(db.get_engine)
+) -> schemas.User:
+    token = None
+    authorization_header = websocket.headers.get('authorization', None)
+    cookie = websocket.cookies.get('access_token')
+
+    if authorization_header:
+        parts = authorization_header.split(' ')
+        if len(parts) == 2:
+            token = parts[1].strip()
+    elif cookie:
+        token = cookie
+
+    if token is None:
+        raise WebSocketException(
+            code=status.WS_1008_POLICY_VIOLATION,
+            reason="token is missing"
+        )
+
+    user_id = get_user_id_from_token(token)
+
+    if user_id is None:
+        raise WebSocketException(
+            code=status.WS_1008_POLICY_VIOLATION,
+            reason="user_id is missing"
+        )
+
+    try:
+        user = db.get_user(engine, user_id)
+    except db_exc.UserNotFound:
+        raise HTTPException(
+            status_code=401,
+            detail="Remote user not found"
+        )
+
+    return user

papermerge/core/db/__init__.py

@@ -0,0 +1,18 @@
+from sqlalchemy import Engine
+
+from .doc_ver import get_last_doc_ver
+from .engine import get_engine
+from .folders import get_folder
+from .nodes import get_paginated_nodes
+from .pages import get_first_page
+from .users import get_user
+
+__all__ = [
+    'get_engine',
+    'get_user',
+    'get_folder',
+    'get_first_page',
+    'get_last_doc_ver',
+    'get_paginated_nodes',
+    'Engine'
+]

papermerge/core/db/doc_ver.py

@@ -0,0 +1,29 @@
+from uuid import UUID
+
+from sqlalchemy import Engine, select
+from sqlalchemy.orm import Session
+
+from papermerge.core import schemas
+from papermerge.core.db.models import Document, DocumentVersion
+
+
+def get_last_doc_ver(
+    engine: Engine,
+    user_id: UUID,
+    doc_id: UUID  # noqa
+) -> schemas.DocumentVersion:
+    """
+    Returns last version of the document
+    identified by doc_id
+    """
+    with Session(engine) as session:  # noqa
+        stmt = select(DocumentVersion).join(Document).where(
+            DocumentVersion.document_id == doc_id,
+            Document.user_id == user_id
+        ).order_by(
+            DocumentVersion.number.desc()
+        ).limit(1)
+        db_doc_ver = session.scalars(stmt).one()
+        model_doc_ver = schemas.DocumentVersion.model_validate(db_doc_ver)
+
+    return model_doc_ver

papermerge/core/db/engine.py

@@ -0,0 +1,19 @@
+import os
+
+from sqlalchemy import Engine, create_engine
+
+SQLALCHEMY_DATABASE_URL = os.environ.get(
+    'PAPERMERGE__DATABASE__URL',
+    'sqlite:////db/db.sqlite3'
+)
+connect_args = {}
+
+if SQLALCHEMY_DATABASE_URL.startswith('sqlite'):
+    # sqlite specific connection args
+    connect_args = {"check_same_thread": False}
+
+engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args=connect_args)
+
+
+def get_engine() -> Engine:
+    return engine

papermerge/core/db/exceptions.py

@@ -0,0 +1,6 @@
+class UserNotFound(Exception):
+    pass
+
+
+class PageNotFound(Exception):
+    pass

papermerge/core/db/folders.py

@@ -0,0 +1,71 @@
+from typing import List, Tuple
+from uuid import UUID
+
+from sqlalchemy import Engine, select, text
+from sqlalchemy.orm import Session
+
+from papermerge.core import schemas
+from papermerge.core.db.models import Folder, Node
+
+
+def get_folder(
+    engine: Engine,
+    folder_id: UUID,
+    user_id: UUID
+):
+    with Session(engine) as session:
+        breadcrumb = get_ancestors(session, folder_id)
+        stmt = select(Folder).where(
+            Folder.id == folder_id,
+            Node.user_id == user_id
+        )
+        db_model = session.scalars(stmt).one()
+        model = schemas.Folder.model_validate(db_model)
+        model.breadcrumb = breadcrumb
+
+    return model
+
+
+def get_ancestors(
+    db: Session,
+    node_id: UUID,
+    include_self=True
+) -> List[Tuple[str, str]]:
+    """Returns all ancestors of the node"""
+    if include_self:
+        stmt = text('''
+            WITH RECURSIVE tree AS (
+                SELECT nodes.id, nodes.title, nodes.parent_id, 0 as level
+                FROM core_basetreenode AS nodes
+                WHERE id = :node_id
+                UNION ALL
+                SELECT nodes.id, nodes.title, nodes.parent_id, level + 1
+                FROM core_basetreenode AS nodes, tree
+                WHERE nodes.id = tree.parent_id
+            )
+            SELECT id, title
+            FROM tree
+            ORDER BY level DESC
+        ''')
+    else:
+        stmt = text('''
+            WITH RECURSIVE tree AS (
+                SELECT nodes.id, nodes.title, nodes.parent_id, 0 as level
+                FROM core_basetreenode AS nodes
+                WHERE id = :node_id
+                UNION ALL
+                SELECT nodes.id, nodes.title, nodes.parent_id, level + 1
+                FROM core_basetreenode AS nodes, tree
+                WHERE nodes.id = tree.parent_id
+            )
+            SELECT id, title
+            FROM tree
+            WHERE NOT id = node_id
+            ORDER BY level DESC
+        ''')
+
+    result = db.execute(stmt, {"node_id": node_id})
+
+    items = list([(id, title) for id, title in result])
+
+    return items