google-cloud-firestore 2.11.0 → 2.13.0

data/lib/google/cloud/firestore/bulk_writer_operation.rb

@@ -0,0 +1,126 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "concurrent"
+require "google/cloud/firestore/bulk_writer_exception"
+
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      #
+      # @private A BulkWriterOperation object refers to a write operation and contains
+      # all the necessary information for a specific write task, including meta
+      # information like the current number of attempts
+      #
+      class BulkWriterOperation
+        attr_reader :retry_time
+        attr_reader :result
+        attr_reader :completion_event
+        attr_reader :write
+
+        ##
+        # Initialize the object
+        def initialize write, retries
+          @write = write
+          @failed_attempts = 0
+          @retries = retries
+          @retry_time = Time.now
+          @completion_event = Concurrent::Event.new
+        end
+
+        ##
+        # Processing to be done when the response is a success.
+        # Updates the result and set the completion event.
+        #
+        # @param [Google::Cloud::Firestore::V1::WriteResult] result The result returned in the response.
+        def on_success result
+          begin
+            @result = WriteResult.new result
+          rescue StandardError => e
+            raise BulkWriterOperationError, e
+          ensure
+            @completion_event.set
+          end
+        end
+
+        ##
+        # Processing to be done when the response is a failure.
+        # Updates the failure attempts. If the retry count reaches
+        # the upper threshold, operations will be marked
+        # as failure and the completion event will be set.
+        #
+        # @param [Google::Rpc::Status] status The status received in the response.
+        #
+        def on_failure status
+          @failed_attempts += 1
+          if @failed_attempts == @retries + 1
+            begin
+              @result = BulkWriterException.new status
+            rescue StandardError => e
+              raise BulkWriterOperationError, e
+            ensure
+              @completion_event.set
+            end
+          else
+            backoff_duration
+          end
+        end
+
+        ##
+        # Exponentially increases the waiting time for retry.
+        #
+        def backoff_duration
+          @retry_time = Time.now + (@failed_attempts**2)
+        end
+
+        ##
+        # Represents the result of applying a write.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #   bw = firestore.bulk_writer
+        #
+        #   # Set the data for NYC
+        #   result = bw.set("cities/NYC", { name: "New York City" })
+        #
+        #   result.wait!
+        #
+        #   puts result.value
+        #
+        class WriteResult
+          ##
+          # The last update time of the document after applying the write. Set to
+          # nil for a +delete+ mutation.
+          #
+          # If the write did not actually change the document, this will be
+          # the previous update_time.
+          #
+          # @return [Time] The last update time.
+          attr_reader :update_time
+
+          ##
+          # @private
+          def initialize result
+            @update_time = Convert.timestamp_to_time result.update_time
+          end
+        end
+      end
+    end
+  end
+end

data/lib/google/cloud/firestore/bulk_writer_scheduler.rb

@@ -0,0 +1,164 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License")
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "concurrent"
+require "google/cloud/firestore/errors"
+require "google/cloud/firestore/bulk_writer_operation"
+require "google/cloud/firestore/rate_limiter"
+require "google/cloud/firestore/bulk_commit_batch"
+require "google/cloud/firestore/bulk_writer_exception"
+require "google/cloud/firestore/bulk_writer_scheduler"
+
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      #
+      # @private Accumulate BulkWriterOperations from the BulkWriter, schedules them
+      # in accordance with 555 rule and retry the failed operations from the BulkCommitBatch.
+      #
+      class BulkWriterScheduler
+        MAX_BATCH_SIZE = 20
+        BATCH_THREAD_COUNT = 4
+
+        ##
+        # Initialize the attributes and start the schedule_operations job
+        #
+        def initialize client, service, batch_threads
+          @client = client
+          @service = service
+          @rate_limiter = RateLimiter.new
+          @buffered_operations = []
+          @batch_threads = (batch_threads || BATCH_THREAD_COUNT).to_i
+          @batch_thread_pool = Concurrent::ThreadPoolExecutor.new max_threads: @batch_threads,
+                                                                  max_queue: 0,
+                                                                  auto_terminate: true
+          @retry_operations = []
+          @mutex = Mutex.new
+          start_scheduling_operations
+        end
+
+        def start_scheduling_operations
+          Concurrent::Promises.future_on @batch_thread_pool do
+            begin
+              schedule_operations
+            rescue StandardError
+              # TODO: Log the error when logging is available
+              retry
+            end
+          end
+        end
+
+        def add_operation operation
+          @mutex.synchronize { @buffered_operations << operation }
+        end
+
+        ##
+        # Closes the scheduler object.
+        # Waits for the enqueued tasks to complete
+        # before closing down.
+        #
+        # @return [nil]
+        def close
+          @mutex.synchronize do
+            @batch_thread_pool.shutdown
+            @batch_thread_pool.wait_for_termination 1
+            @batch_thread_pool.kill unless @batch_thread_pool.shutdown?
+          end
+        end
+
+        private
+
+        ##
+        # @private Adds failed operations in the retry heap.
+        #
+        def post_commit_batch bulk_commit_batch
+          @mutex.synchronize do
+            bulk_commit_batch.operations.each do |operation|
+              unless operation.completion_event.set?
+                @retry_operations << operation
+              end
+            end
+            @retry_operations.sort_by!(&:retry_time)
+          end
+        end
+
+        ##
+        # @private Commits a batch of scheduled operations.
+        # Batch size <= 20 to match the constraint of request size < 9.8 MB
+        #
+        # @return [nil]
+        def commit_batch bulk_commit_batch
+          begin
+            Concurrent::Promises.future_on @batch_thread_pool, bulk_commit_batch do |batch|
+              begin
+                batch.commit
+              rescue StandardError
+                # TODO: Log the errors while committing a batch
+              ensure
+                post_commit_batch bulk_commit_batch
+              end
+            end
+          rescue StandardError => e
+            post_commit_batch bulk_commit_batch
+            raise BulkWriterSchedulerError, e.message
+          end
+        end
+
+        ##
+        # @private Schedule the enqueued operations in batches.
+        #
+        # @return [nil]
+        def schedule_operations
+          loop do
+            break if @batch_thread_pool.shuttingdown?
+            dequeue_retry_operations
+            batch_size = [MAX_BATCH_SIZE, @buffered_operations.length].min
+            if batch_size.zero?
+              sleep 0.001
+              next
+            end
+            @rate_limiter.wait_for_tokens batch_size
+            operations = dequeue_buffered_operations batch_size
+            commit_batch BulkCommitBatch.new(@service, operations)
+          end
+        end
+
+        ##
+        # @private Removes BulkWriterOperations from the buffered queue to scheduled in
+        # the current batch
+        #
+        def dequeue_buffered_operations size
+          @mutex.synchronize do
+            @buffered_operations.shift size
+          end
+        end
+
+        ##
+        # @private Removes BulkWriterOperations from the retry queue to scheduled in
+        # the current batch
+        #
+        def dequeue_retry_operations
+          @mutex.synchronize do
+            while @retry_operations.length.positive? && @retry_operations.first.retry_time <= Time.now
+              @buffered_operations << @retry_operations.shift
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/google/cloud/firestore/client.rb

@@ -23,6 +23,8 @@ require "google/cloud/firestore/document_snapshot"
 require "google/cloud/firestore/collection_group"
 require "google/cloud/firestore/batch"
 require "google/cloud/firestore/transaction"
+require "google/cloud/firestore/bulk_writer"
+require "google/cloud/firestore/filter"
 
 module Google
   module Cloud
@@ -185,6 +187,56 @@ module Google
         end
         alias collection_group col_group
 
+        ##
+        # Creates a filter object.
+        #
+        # @param field [FieldPath, String, Symbol] A field path to filter
+        #   results with.
+        #   If a {FieldPath} object is not provided then the field will be
+        #   treated as a dotted string, meaning the string represents individual
+        #   fields joined by ".". Fields containing `~`, `*`, `/`, `[`, `]`, and
+        #   `.` cannot be in a dotted string, and should provided using a
+        #   {FieldPath} object instead.
+        #
+        # @param operator [String, Symbol] The operation to compare the field
+        #   to. Acceptable values include:
+        #   * less than: `<`, `lt`
+        #   * less than or equal: `<=`, `lte`
+        #   * greater than: `>`, `gt`
+        #   * greater than or equal: `>=`, `gte`
+        #   * equal: `=`, `==`, `eq`, `eql`, `is`
+        #   * not equal: `!=`
+        #   * in: `in`
+        #   * not in: `not-in`, `not_in`
+        #   * array contains: `array-contains`, `array_contains`
+        #
+        # @param value [Object] The value to compare the property to. Defaults to nil.
+        #   Possible values are:
+        #   * Integer
+        #   * Float/BigDecimal
+        #   * String
+        #   * Boolean
+        #   * Array
+        #   * Date/Time
+        #   * StringIO
+        #   * Google::Cloud::Datastore::Key
+        #   * Google::Cloud::Datastore::Entity
+        #   * nil
+        #
+        # @return [Google::Cloud::Firestore::Filter] New filter object.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #
+        #   # Create a filter
+        #   filter = firestore.filter(:population, :>=, 1000000)
+        #
+        def filter field, operator, value
+          Filter.new field, operator, value
+        end
+
         ##
         # Retrieves a document reference.
         #
@@ -738,6 +790,35 @@ module Google
           yield transaction
         end
 
+        ##
+        # Create a bulk writer to perform multiple writes that are
+        # executed parallely.
+        #
+        # @param [Integer] request_threads The number of threads used for handling
+        #  requests. Default is 2. Optional.
+        # @param [Integer] batch_threads The number of threads used for processing
+        #  batches. Default is 4. Optional.
+        # @param [Integer] retries The number of times a failed write request will
+        # be retried (with exponential delay) before being marked as failure. Max
+        # attempts are 15. Optional
+        #
+        # @return [Google::Cloud::Firestore::BulkWriter] Returns an object of
+        #   bulk writer.
+        #
+        # @example Initializing a BulkWriter with all the configurations.
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #
+        #   bw = firestore.bulk_writer
+        #
+        #   bulk_write_result = bw.create "doc_ref", request_threads: 4, batch_threads: 10, retries: 10
+        #
+        def bulk_writer request_threads: nil, batch_threads: nil, retries: nil
+          BulkWriter.new self, @service, request_threads: request_threads,
+                         batch_threads: batch_threads, retries: retries
+        end
+
         # @!endgroup
 
         # @private

data/lib/google/cloud/firestore/errors.rb

@@ -0,0 +1,60 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "google/cloud/errors"
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      # Indicates that the an error was reported while scheduling
+      # BulkWriter operations.
+      #
+      class BulkWriterSchedulerError < Google::Cloud::Error
+        def initialize message
+          super "BulkWriterSchedulerError : #{message}"
+        end
+      end
+
+      ##
+      # Indicates that the an error was reported while committing a
+      # batch of operations.
+      #
+      class BulkCommitBatchError < Google::Cloud::Error
+        def initialize message
+          super "BulkCommitBatchError : #{message}"
+        end
+      end
+
+      ##
+      # Indicates that the an error was reported while parsing response for
+      # BulkWriterOperation.
+      #
+      class BulkWriterOperationError < Google::Cloud::Error
+        def initialize message
+          super "BulkWriterOperationError : #{message}"
+        end
+      end
+
+      ##
+      # Indicates that the an error was reported in BulkWriter.
+      #
+      class BulkWriterError < Google::Cloud::Error
+        def initialize message
+          super "BulkWriterError : #{message}"
+        end
+      end
+    end
+  end
+end