683 include thumbnail ids in elasticsearch index (#756)

* adding method stub

* adding index for thumbnails

* index worked

* thumbnail now has downloads field
incremented False by default
using that in es indexing

* formatting

* codegen

Author: tcnichol

backend/app/models/thumbnails.py

@@ -23,6 +23,7 @@ class ThumbnailDB(Document, ThumbnailBase):
     modified: datetime = Field(default_factory=datetime.utcnow)
     bytes: int = 0
     content_type: ContentType = ContentType()
+    downloads: int = 0
 
     class Settings:
         name = "thumbnails"

backend/app/routers/files.py

@@ -42,7 +42,7 @@
     insert_record,
     update_record,
 )
-from app.search.index import index_file
+from app.search.index import index_file, index_thumbnail
 
 router = APIRouter()
 security = HTTPBearer()
@@ -515,6 +515,7 @@ async def add_file_thumbnail(
     file_id: str,
     thumbnail_id: str,
     allow: bool = Depends(FileAuthorization("editor")),
+    es: Elasticsearch = Depends(dependencies.get_elasticsearchclient),
 ):
     if (file := await FileDB.get(PydanticObjectId(file_id))) is not None:
         if (
@@ -523,6 +524,9 @@ async def add_file_thumbnail(
             # TODO: Should we garbage collect existing thumbnail if nothing else points to it?
             file.thumbnail_id = thumbnail_id
             await file.save()
+            await index_thumbnail(
+                es, thumbnail_id, str(file.id), str(file.dataset_id), False
+            )
             return file.dict()
         else:
             raise HTTPException(

backend/app/routers/thumbnails.py

@@ -1,6 +1,7 @@
 from typing import List, Optional
 
 from beanie import PydanticObjectId
+from beanie.odm.operators.update.general import Inc
 from fastapi import APIRouter, HTTPException, Depends
 from fastapi import File, UploadFile
 from fastapi.security import HTTPBearer
@@ -72,7 +73,9 @@ async def remove_thumbnail(thumb_id: str, fs: Minio = Depends(dependencies.get_f
 
 @router.get("/{thumbnail_id}")
 async def download_thumbnail(
-    thumbnail_id: str, fs: Minio = Depends(dependencies.get_fs)
+    thumbnail_id: str,
+    fs: Minio = Depends(dependencies.get_fs),
+    increment: Optional[bool] = False,
 ):
     # If thumbnail exists in MongoDB, download from Minio
     if (thumbnail := await ThumbnailDB.get(PydanticObjectId(thumbnail_id))) is not None:
@@ -81,6 +84,9 @@ async def download_thumbnail(
         # Get content type & open file stream
         response = StreamingResponse(content.stream(settings.MINIO_UPLOAD_CHUNK_SIZE))
         response.headers["Content-Disposition"] = "attachment; filename=%s" % "thumb"
+        if increment:
+            # Increment download count
+            await thumbnail.update(Inc({ThumbnailDB.downloads: 1}))
         return response
     else:
         raise HTTPException(

backend/app/search/index.py

@@ -2,11 +2,12 @@
 
 from bson import ObjectId
 from elasticsearch import Elasticsearch, NotFoundError
-
+from beanie import PydanticObjectId
 from app.config import settings
 from app.models.authorization import AuthorizationDB
-from app.models.datasets import DatasetOut
-from app.models.files import FileOut
+from app.models.datasets import DatasetOut, DatasetDB
+from app.models.files import FileOut, FileDB
+from app.models.thumbnails import ThumbnailOut, ThumbnailDB
 from app.models.metadata import MetadataDB
 from app.models.search import (
     ElasticsearchEntry,
@@ -107,3 +108,56 @@ async def index_file(
             insert_record(es, settings.elasticsearch_index, doc, file.id)
     else:
         insert_record(es, settings.elasticsearch_index, doc, file.id)
+
+
+async def index_thumbnail(
+    es: Elasticsearch,
+    thumbnail_id: str,
+    file_id: str,
+    dataset_id: str,
+    update: bool = False,
+):
+    """Create or update an Elasticsearch entry for the file. user_ids is the list of users
+    with permission to at least view the file's dataset, it will be queried if not provided.
+    """
+    if (file := await FileDB.get(PydanticObjectId(file_id))) is not None:
+        if (
+            thumbnail := await ThumbnailDB.get(PydanticObjectId(thumbnail_id))
+        ) is not None:
+            # Get authorized users from db
+            authorized_user_ids = []
+            async for auth in AuthorizationDB.find(
+                AuthorizationDB.dataset_id == ObjectId(dataset_id)
+            ):
+                authorized_user_ids += auth.user_ids
+            # Get full metadata from db (granular updates possible but complicated)
+            metadata = []
+            async for md in MetadataDB.find(
+                MetadataDB.resource.resource_id == ObjectId(file.id)
+            ):
+                metadata.append(md.content)
+            # Add en entry to the file index
+            doc = ElasticsearchEntry(
+                resource_type="thumbnail",
+                name=file.name,
+                creator=thumbnail.creator.email,
+                created=thumbnail.created,
+                user_ids=authorized_user_ids,
+                content_type=thumbnail.content_type.content_type,
+                content_type_main=thumbnail.content_type.main_type,
+                file_id=str(file.id),
+                dataset_id=str(file.dataset_id),
+                folder_id=str(file.folder_id),
+                bytes=thumbnail.bytes,
+                metadata=metadata,
+                downloads=thumbnail.downloads,
+            ).dict()
+            if update:
+                try:
+                    update_record(
+                        es, settings.elasticsearch_index, {"doc": doc}, file.id
+                    )
+                except NotFoundError:
+                    insert_record(es, settings.elasticsearch_index, doc, file.id)
+            else:
+                insert_record(es, settings.elasticsearch_index, doc, file.id)

frontend/src/openapi/v2/models/ThumbnailOut.ts

@@ -25,4 +25,5 @@ export type ThumbnailOut = {
     modified?: string;
     bytes?: number;
     content_type?: ContentType;
+    downloads?: number;
 }

frontend/src/openapi/v2/services/ThumbnailsService.ts

@@ -32,15 +32,20 @@ export class ThumbnailsService {
     /**
      * Download Thumbnail
      * @param thumbnailId
+     * @param increment
      * @returns any Successful Response
      * @throws ApiError
      */
     public static downloadThumbnailApiV2ThumbnailsThumbnailIdGet(
         thumbnailId: string,
+        increment: boolean = false,
     ): CancelablePromise<any> {
         return __request({
             method: 'GET',
             path: `/api/v2/thumbnails/${thumbnailId}`,
+            query: {
+                'increment': increment,
+            },
             errors: {
                 422: `Validation Error`,
             },