From d5ab198c6f6268aa0abc5cba50a6a00c2ac3578a Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Tue, 21 Jan 2025 18:07:55 +0530
Subject: [PATCH] feat: handle document limit in the model

---
 .../api/v1/accounts/captain/documents_controller.rb         | 6 ------
 enterprise/app/models/captain/document.rb                   | 6 ++++++
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb b/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb
index ba488269b..9743c3892 100644
--- a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb
+++ b/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb
@@ -3,7 +3,6 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC
   before_action -> { check_authorization(Captain::Assistant) }
 
   before_action :set_current_page, only: [:index]
-  before_action :ensure_within_plan_limit, only: [:create]
   before_action :set_documents, except: [:create]
   before_action :set_document, only: [:show, :destroy]
   before_action :set_assistant, only: [:create]
@@ -49,11 +48,6 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC
     @current_page = permitted_params[:page] || 1
   end
 
-  def ensure_within_plan_limit
-    limits = Current.account.usage_limits[:captain][:documents]
-    return render_could_not_create_error('Document limit exceeded') unless limits[:current_available].positive?
-  end
-
   def permitted_params
     params.permit(:assistant_id, :page, :id, :account_id)
   end
diff --git a/enterprise/app/models/captain/document.rb b/enterprise/app/models/captain/document.rb
index 703724607..3332f863a 100644
--- a/enterprise/app/models/captain/document.rb
+++ b/enterprise/app/models/captain/document.rb
@@ -35,6 +35,7 @@ class Captain::Document < ApplicationRecord
     available: 1
   }
 
+  before_create :ensure_within_plan_limit
   after_create_commit :enqueue_crawl_job
   after_create_commit :update_document_usage
   after_destroy :update_document_usage
@@ -65,4 +66,9 @@ class Captain::Document < ApplicationRecord
   def ensure_account_id
     self.account_id = assistant&.account_id
   end
+
+  def ensure_within_plan_limit
+    limits = account.usage_limits[:captain][:documents]
+    raise 'Document limit exceeded' unless limits[:current_available].positive?
+  end
 end