Track document attachments

app/models/document.rb

@@ -1,16 +1,24 @@
 class Document < ApplicationRecord
   belongs_to :project
+  has_one :owner, through: :project
+
   has_many :documents_tags, dependent: :destroy
   has_many :tags, through: :documents_tags
   accepts_nested_attributes_for :tags
 
+  has_many :documents_s3_files, dependent: :destroy
+  has_many :s3_files, through: :documents_s3_files
+
   scope :blank, -> { where(blank: true) }
   scope :not_blank, -> { where(blank: false) }
   scope :pinned, -> { where.not(pinned_at: nil) }
 
   include Queryable.permit(*%i[id title safe_title preview body body_type tags blank updated_by created_at updated_at pinned_at locked_at])
   include Listenable
 
+  after_create :update_linked_s3_files
+  after_update :update_linked_s3_files
+
   after_commit :upsert_to_typesense, on: %i[create update]
   after_destroy :destroy_from_typesense
 
@@ -34,6 +42,10 @@ def preview
     (preview.presence || plain_body.slice(0, 100)).strip
   end
 
+  def slate?
+    body_type == 'json/slate'
+  end
+
   def was_updated_on_server
     self.updated_by = 'server'
   end
@@ -64,6 +76,22 @@ def tags_attributes=(tags_attributes)
     end
   end
 
+  def update_linked_s3_files
+    return unless slate?
+
+    s3_file_ids = []
+
+    SlateUtils::SlateNode.parse(body).traverse do |node|
+      if node.type == 'attachment'
+        s3_file_ids << node.attributes[:s3FileId]
+      end
+    end
+
+    self.s3_files = s3_file_ids.uniq.map do |s3_file_id|
+      owner.s3_files.find_by(id: s3_file_id)
+    end.compact
+  end
+
   def upsert_to_typesense(collection: self.class.typesense_collection)
     collection.documents.upsert(
       id: id.to_s,

app/models/documents_s3_file.rb

@@ -0,0 +1,4 @@
+class DocumentsS3File < ApplicationRecord
+  belongs_to :document
+  belongs_to :s3_file
+end

app/models/replace_in_document.rb

@@ -2,49 +2,31 @@ class ReplaceInDocument
   def self.perform(document:, find:, replace:)
     raise ArgumentError, 'find cannot be blank' if find.blank?
 
-    return 0 unless document.body_type.starts_with?('json/')
+    return 0 unless document.slate?
 
-    document_root = JSON.parse(document.body)
+    document_root = SlateUtils::SlateNode.parse(document.body)
 
     replace_count = 0
 
-    new_document_root = document_root.map do |block|
-      replace_in_node(block) do |text|
-        text.tap do |replaced_text|
-          offset = 0
+    document_root.traverse do |node|
+      if node.text?
+        offset = 0
 
-          while (index = text.downcase.index(find.downcase, offset))
-            replaced_text[index, find.length] = replace
-            offset = index + replace.length
-            replace_count += 1
-          end
+        while (index = node.text.downcase.index(find.downcase, offset))
+          node.text[index, find.length] = replace
+          offset = index + replace.length
+          replace_count += 1
         end
       end
     end
 
     if replace_count > 0
-      document.body = JSON.generate(new_document_root)
+      document.body = document_root.to_json
       document.plain_body.gsub!(find, replace)
       document.was_updated_on_server
       document.save!
     end
 
     replace_count
   end
-
-  private
-
-  def self.replace_in_node(node, &block)
-    if node.key?('text')
-      node['text'] = yield(node['text'])
-    end
-
-    if node.key?('children')
-      node['children'].map! do |child|
-        replace_in_node(child, &block)
-      end
-    end
-
-    node
-  end
 end

app/models/s3_file.rb

@@ -3,6 +3,9 @@ class S3File < ApplicationRecord
   belongs_to :original_project, class_name: 'Project', inverse_of: :s3_files
   has_many :used_as_image_in_projects, class_name: 'Project', foreign_key: 'image_id', dependent: :nullify
 
+  has_many :documents_s3_files, dependent: :destroy
+  has_many :documents, through: :documents_s3_files
+
   validates :role, presence: true
   validates :s3_key, presence: true
   validates :filename, presence: true

app/models/slate_utils.rb

@@ -0,0 +1,75 @@
+module SlateUtils
+  class SlateChildren < Array
+    def to_json
+      map(&:to_h).to_json
+    end
+
+    def self.from_a(array)
+      new(array.map { |h| SlateNode.from_h(h) })
+    end
+
+    def traverse(&block)
+      each do |node|
+        node.traverse(&block)
+      end
+    end
+  end
+
+  class SlateNode
+    SPECIAL_ATTRIBUTES = %i[type text children].freeze
+
+    attr_accessor :type, :text, :children, :attributes
+
+    def initialize(type: nil, text: nil, attributes: {}, children: nil)
+      @type = type
+      @text = text
+      @attributes = attributes
+      @children = children
+    end
+
+    def element?
+      !@type.nil?
+    end
+
+    def text?
+      !@text.nil?
+    end
+
+    def traverse(&block)
+      yield self
+      children&.traverse(&block)
+    end
+
+    def self.from_h(hash)
+      new(
+        type: hash[:type],
+        text: hash[:text],
+        attributes: hash.reject { |k, _| SPECIAL_ATTRIBUTES.include?(k) },
+        children: hash[:children]&.then { |children| SlateChildren.from_a(children) }
+      )
+    end
+
+    def self.parse(json)
+      parsed = JSON.parse(json, symbolize_names: true)
+
+      if parsed.is_a?(Array)
+        SlateChildren.from_a(parsed)
+      else
+        from_h(parsed)
+      end
+    end
+
+    def to_h
+      {}.tap do |h|
+        h[:type] = type unless type.nil?
+        h[:text] = text unless text.nil?
+        h.merge!(attributes)
+        h[:children] = children.map(&:to_h) unless children.nil?
+      end
+    end
+
+    def to_json
+      to_h.to_json
+    end
+  end
+end

db/migrate/20231125151117_create_join_table_documents_s3_files.rb

@@ -0,0 +1,16 @@
+class CreateJoinTableDocumentsS3Files < ActiveRecord::Migration[7.0]
+  def up
+    create_join_table :documents, :s3_files do |t|
+      t.index [:document_id, :s3_file_id]
+      t.index [:s3_file_id, :document_id]
+    end
+
+    Document.find_each do |document|
+      document.update_linked_s3_files
+    end
+  end
+
+  def down
+    drop_join_table :documents, :s3_files
+  end
+end

test/models/documents_s3_file_test.rb

@@ -0,0 +1,72 @@
+require 'test_helper'
+
+class DocumentsS3FileTest < ActiveSupport::TestCase
+  include SlateJSONHelper
+
+  def setup
+    @user_1 = create(:user)
+    @user_2 = create(:user)
+
+    @s3_file_1 = create(:s3_file, owner: @user_1)
+    @s3_file_2 = create(:s3_file, owner: @user_1)
+
+    @user_1_project = create(:project, owner: @user_1)
+    @user_2_project = create(:project, owner: @user_2)
+  end
+
+  test 'document with no attachments has no linked s3_files' do
+    document = create_document_with_attachment_ids([], project: @user_1_project)
+    assert_empty document.s3_files
+  end
+
+  test 'document with one attachment has one linked s3_file' do
+    document = create_document_with_attachment_ids([@s3_file_1.id], project: @user_1_project)
+    assert_equal [@s3_file_1], document.s3_files
+  end
+
+  test 'linked s3_files are updated when document is updated' do
+    document = create_document_with_attachment_ids([@s3_file_1.id], project: @user_1_project)
+    document.update!(body: document_body_with_attachment_ids([@s3_file_2.id]))
+    assert_equal [@s3_file_2], document.s3_files
+  end
+
+  test 'no error when s3 file does not exist' do
+    fake_s3_file_id = 123456789
+    assert_nil S3File.find_by(id: fake_s3_file_id)
+    document = create_document_with_attachment_ids([fake_s3_file_id], project: @user_1_project)
+    assert_empty document.s3_files
+  end
+
+  test 'cannot link s3 file belonging to another user' do
+    document = create_document_with_attachment_ids([@s3_file_1.id], project: @user_2_project)
+    assert_empty document.s3_files
+  end
+
+  private
+
+  def create_document_with_attachment_ids(s3_file_ids, **args)
+    create(
+      :document,
+      body: document_body_with_attachment_ids(s3_file_ids),
+      body_type: 'json/slate',
+      **args
+    )
+  end
+
+  def document_body_with_attachment_ids(s3_file_ids)
+    create_document_body do
+      p do
+        text 'Hello, world!'
+      end
+
+      s3_file_ids.each do |s3_file_id|
+        # Future proofing: Make sure attachments in descendants are counted
+        p do
+          p do
+            attachment s3_file_id: s3_file_id, filename: 'test.txt'
+          end
+        end
+      end
+    end
+  end
+end

test/models/slate_utils_test.rb

@@ -0,0 +1,73 @@
+require 'test_helper'
+
+class SlateUtilsTest < ActiveSupport::TestCase
+  include SlateJSONHelper
+
+  test 'SlateNode parses and serializes document bodies' do
+    original_json = create_document_body do
+      p do
+        text 'Hello, world!', bold: true
+      end
+
+      blockquote do
+        p do
+          text 'This is a blockquote'
+        end
+      end
+    end
+
+    expected_json = create_document_body do
+      p do
+        text 'Goodbye, world!', bold: true
+      end
+
+      blockquote do
+        p do
+          text 'This is a blockquote'
+        end
+      end
+    end
+
+    slate_nodes = SlateUtils::SlateNode.parse(original_json)
+    slate_nodes[0].children[0].text = 'Goodbye, world!'
+    slate_nodes_json = slate_nodes.to_json
+
+    assert_equal JSON.parse(expected_json), JSON.parse(slate_nodes_json)
+  end
+
+  test 'traverses document bodies' do
+    json = create_document_body do
+      p do
+        text 'Hello, world!', bold: true
+      end
+
+      blockquote do
+        p do
+          text 'This is a blockquote'
+        end
+      end
+    end
+
+    slate_nodes = SlateUtils::SlateNode.parse(json)
+
+    nodes = []
+
+    slate_nodes.traverse do |node|
+      nodes << node
+    end
+
+    first_paragraph = nodes[0]
+    first_text = nodes[1]
+    blockquote = nodes[2]
+    second_paragraph = nodes[3]
+    second_text = nodes[4]
+
+    assert_equal 'p', first_paragraph.type
+    assert_equal 'Hello, world!', first_text.text
+    assert first_text.attributes[:bold], 'first_text should be bold'
+    assert_equal 'blockquote', blockquote.type
+    assert_equal 'p', second_paragraph.type
+    assert_equal 'This is a blockquote', second_text.text
+    assert_nil second_text.attributes[:bold], 'second_text should not be bold'
+  end
+end