Backend fixes for queries and showing entity info per mediaItem

prabhuomkar · Oct 18, 2021 · aa7f486 · aa7f486
1 parent d9e69f4
commit aa7f486
Show file tree

Hide file tree

Showing 6 changed files with 102 additions and 14 deletions.
diff --git a/api/internal/graph/generated/generated.go b/api/internal/graph/generated/generated.go
diff --git a/api/internal/graph/resolvers/schema.resolvers.go b/api/internal/graph/resolvers/schema.resolvers.go
@@ -82,6 +82,26 @@ func (r *entityResolver) MediaItems(ctx context.Context, obj *models.Entity, pag
 	}, nil
 }
 
+func (r *mediaItemResolver) Entities(ctx context.Context, obj *models.MediaItem) ([]*models.Entity, error) {
+	entityIDs := make([]primitive.ObjectID, len(obj.Entities))
+	for idx, strID := range obj.Entities {
+		oid, _ := primitive.ObjectIDFromHex(strID)
+		entityIDs[idx] = oid
+	}
+
+	cur, err := r.DB.Collection(models.ColEntity).Find(ctx, bson.M{"_id": bson.M{"$in": entityIDs}})
+	if err != nil {
+		return nil, err
+	}
+
+	var result []*models.Entity
+	if err = cur.All(ctx, &result); err != nil {
+		return nil, err
+	}
+
+	return result, nil
+}
+
 func (r *mutationResolver) Upload(ctx context.Context, file graphql.Upload) (bool, error) {
 	result, err := r.CDN.Upload(file.File, file.Filename, file.Size, "", "")
 	if err != nil {
@@ -302,12 +322,16 @@ func (r *queryResolver) Entity(ctx context.Context, id string) (*models.Entity,
 // Entity returns generated.EntityResolver implementation.
 func (r *Resolver) Entity() generated.EntityResolver { return &entityResolver{r} }
 
+// MediaItem returns generated.MediaItemResolver implementation.
+func (r *Resolver) MediaItem() generated.MediaItemResolver { return &mediaItemResolver{r} }
+
 // Mutation returns generated.MutationResolver implementation.
 func (r *Resolver) Mutation() generated.MutationResolver { return &mutationResolver{r} }
 
 // Query returns generated.QueryResolver implementation.
 func (r *Resolver) Query() generated.QueryResolver { return &queryResolver{r} }
 
 type entityResolver struct{ *Resolver }
+type mediaItemResolver struct{ *Resolver }
 type mutationResolver struct{ *Resolver }
 type queryResolver struct{ *Resolver }
diff --git a/api/internal/models/entity.go b/api/internal/models/entity.go
@@ -3,10 +3,11 @@ package models
 const ColEntity = "entities"
 
 type Entity struct {
-	ID         string `json:"id" bson:"_id"`
-	Name       string `json:"name"`
-	ImageURL   string `json:"imageUrl"`
-	EntityType string `json:"entityType"`
-	CreatedAt  string `json:"createdAt"`
-	UpdatedAt  string `json:"updatedAt"`
+	ID         string   `json:"id" bson:"_id"`
+	Name       string   `json:"name"`
+	ImageURL   string   `json:"imageUrl"`
+	EntityType string   `json:"entityType"`
+	MediaItems []string `json:"mediaItems"`
+	CreatedAt  string   `json:"createdAt"`
+	UpdatedAt  string   `json:"updatedAt"`
 }
diff --git a/api/internal/models/mediaitem.go b/api/internal/models/mediaitem.go
@@ -13,6 +13,7 @@ type (
 		FileName      string         `json:"fileName"`
 		FileSize      int64          `json:"fileSize"`
 		MediaMetadata *MediaMetaData `json:"mediaMetadata"`
+		Entities      []string       `json:"entities"`
 		CreatedAt     time.Time      `json:"createdAt"`
 		UpdatedAt     time.Time      `json:"updatedAt"`
 	}

diff --git a/api/schema.graphql b/api/schema.graphql
@@ -9,6 +9,7 @@ type MediaItem {
   fileName: String!
   fileSize: Int!
   mediaMetadata: MediaMetaData
+  entities: [Entity!]
   createdAt: Time!
   updatedAt: Time!
 }

diff --git a/worker/pipeline/places.py b/worker/pipeline/places.py
@@ -2,6 +2,7 @@
 import json
 import urllib.request
 import exiftool
+from bson.objectid import ObjectId
 from pymongo import ReturnDocument
 from .component import Component
 
@@ -23,6 +24,10 @@ def upsert_entity(self, data):
       upsert=True,
       return_document=ReturnDocument.AFTER
     )
+    self.db['entities'].update_one(
+      {'_id': result['_id']},
+      {'$addToSet': {'mediaItems': ObjectId(self.oid)}},
+    )
     print(f'[place]: {result}')
     return result['_id']