mongoid 8.1.6 → 9.0.0

data/lib/mongoid/tasks/encryption.rake

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require 'optparse'
+
+def parse_data_key_options(argv = ARGV)
+  # The only way to use OptionParser to parse custom options in rake is
+  # to pass an empty argument ("--") before specifying them, e.g.:
+  #
+  #    rake db:mongoid:encryption:create_data_key -- --client default
+  #
+  # Otherwise, rake complains about an unknown option. Thus, we can tell
+  # if the argument list is valid for us to parse by detecting this empty
+  # argument.
+  #
+  # (This works around an issue in the tests, where the specs are loading
+  # the tasks directly to test them, but the option parser is then picking
+  # up rspec command-line arguments and raising an exception).
+  return {} unless argv.include?('--')
+
+  {}.tap do |options|
+    parser = OptionParser.new do |opts|
+      opts.on('-c', '--client CLIENT', 'Name of the client to use') do |v|
+        options[:client_name] = v
+      end
+      opts.on('-p', '--provider PROVIDER', 'KMS provider to use') do |v|
+        options[:kms_provider_name] = v
+      end
+      opts.on('-n', '--key-alt-name KEY_ALT_NAME', 'Alternate name for the key') do |v|
+        options[:key_alt_name] = v
+      end
+    end
+    # rubocop:disable Lint/EmptyBlock
+    parser.parse!(parser.order!(argv) {})
+    # rubocop:enable Lint/EmptyBlock
+  end
+end
+
+namespace :db do
+  namespace :mongoid do
+    namespace :encryption do
+      desc 'Create encryption key'
+      task create_data_key: [ :environment ] do
+        options = parse_data_key_options
+        result = Mongoid::Tasks::Encryption.create_data_key(
+          client_name: options[:client_name],
+          kms_provider_name: options[:kms_provider_name],
+          key_alt_name: options[:key_alt_name]
+        )
+        output = [].tap do |lines|
+          lines << "Created data key with id: '#{result[:key_id]}'"
+          lines << "with key alt name: '#{result[:key_alt_name]}'" if result[:key_alt_name]
+          lines << "for kms provider: '#{result[:kms_provider]}'"
+          lines << "in key vault: '#{result[:key_vault_namespace]}'."
+        end
+        puts output.join(' ')
+      end
+    end
+  end
+end

data/lib/mongoid/tasks/encryption.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+# rubocop:todo all
+
+module Mongoid
+  module Tasks
+    # This module contains helper methods for data encryption.
+    module Encryption
+      extend self
+
+      # Create a data encryption key for the given kms provider using the
+      # auto_encryption_options from the client's configuration.
+      #
+      # @param [ String | nil ] kms_provider_name The name of the kms provider
+      #   to use. If not provided, the first provider in the client's
+      #   auto_encryption_options will be used.
+      # @param [ String | nil ] client_name The name of the client to take
+      #   auto_encryption_options from. If not provided, the default client
+      #   will be used.
+      # @param [ String | nil ] key_alt_name The alternate name of the key.
+      #
+      # @return [ Hash ] A hash containing the key id as :key_id,
+      #   kms provider name as :kms_provider, and key vault namespace as
+      #   :key_vault_namespace.
+      def create_data_key(client_name: nil, kms_provider_name: nil, key_alt_name: nil)
+        kms_provider_name, kms_providers, key_vault_namespace = prepare_arguments(
+          kms_provider_name,
+          client_name
+        )
+        key_vault_client = Mongoid::Clients.default.with(database: key_vault_namespace.split('.').first)
+        client_encryption = Mongo::ClientEncryption.new(
+          key_vault_client,
+          key_vault_namespace: key_vault_namespace,
+          kms_providers: kms_providers
+        )
+        client_encryption_opts = {}.tap do |opts|
+          opts[:key_alt_names] = [key_alt_name] if key_alt_name
+        end
+        data_key_id = client_encryption.create_data_key(kms_provider_name, client_encryption_opts)
+        {
+          key_id: Base64.strict_encode64(data_key_id.data),
+          kms_provider: kms_provider_name,
+          key_vault_namespace: key_vault_namespace,
+          key_alt_name: key_alt_name
+        }.compact
+      end
+
+      private
+
+      # Prepare arguments needed to create a data key from the client's
+      # auto_encryption_options.
+      #
+      # @param [ String | nil ] kms_provider_name The name of the kms provider.
+      # @param [ String | nil ] client_name The name of the client.
+      #
+      # @return [ Array<String, Hash, String> ] An array containing the
+      #   normalized kms provider name, the kms providers hash, and the key
+      #   vault namespace.
+      def prepare_arguments(kms_provider_name, client_name)
+        client = (client_name || 'default').to_s
+        client_options = Mongoid.clients[client]
+        unless client_options.is_a?(Hash)
+          raise Errors::NoClientConfig.new(client)
+        end
+        auto_encryption_options = client_options.dig(:options, :auto_encryption_options)
+        unless auto_encryption_options.is_a?(Hash)
+          raise Errors::InvalidAutoEncryptionConfiguration.new(client)
+        end
+        key_vault_namespace = auto_encryption_options[:key_vault_namespace]
+        unless key_vault_namespace.is_a?(String)
+          raise Errors::InvalidAutoEncryptionConfiguration.new(client)
+        end
+        kms_providers = auto_encryption_options[:kms_providers]
+        unless kms_providers.is_a?(Hash)
+          raise Errors::InvalidAutoEncryptionConfiguration.new(client)
+        end
+        valid_kms_provider_name = get_kms_provider_name(kms_provider_name, kms_providers)
+        unless kms_providers.key?(valid_kms_provider_name)
+          raise Errors::InvalidAutoEncryptionConfiguration.new(client, valid_kms_provider_name)
+        end
+
+        [valid_kms_provider_name, kms_providers, key_vault_namespace]
+      end
+
+      # Get kms provider name to use for creating a data key.
+      #
+      # If kms_provider_name is provided, it will be used. Otherwise, if there
+      # is only one kms provider, that provider will be used. Otherwise, an
+      # error will be raised.
+      #
+      # @param [ String | nil ] kms_provider_name The name of the kms provider
+      #   as provided by the user.
+      # @param [ Hash ] kms_providers The kms providers hash from the client's
+      #   auto_encryption_options.
+      #
+      # @return [ String ] The kms provider name to use for creating a data
+      #   key.
+      def get_kms_provider_name(kms_provider_name, kms_providers)
+        if kms_provider_name
+          kms_provider_name
+        elsif kms_providers.keys.length == 1
+          kms_providers.keys.first
+        else
+          raise ArgumentError, 'kms_provider_name must be provided when there are multiple kms providers'
+        end
+      end
+    end
+  end
+end

data/lib/mongoid/threaded/lifecycle.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 module Mongoid
   module Threaded

data/lib/mongoid/threaded.rb

@@ -1,31 +1,35 @@
 # frozen_string_literal: true
 
-require "mongoid/threaded/lifecycle"
+require 'mongoid/threaded/lifecycle'
 
 module Mongoid
-
   # This module contains logic for easy access to objects that have a lifecycle
   # on the current thread.
   module Threaded
-
-    DATABASE_OVERRIDE_KEY = "[mongoid]:db-override"
+    DATABASE_OVERRIDE_KEY = '[mongoid]:db-override'
 
     # Constant for the key to store clients.
-    CLIENTS_KEY = "[mongoid]:clients"
+    CLIENTS_KEY = '[mongoid]:clients'
 
     # The key to override the client.
-    CLIENT_OVERRIDE_KEY = "[mongoid]:client-override"
+    CLIENT_OVERRIDE_KEY = '[mongoid]:client-override'
 
     # The key for the current thread's scope stack.
-    CURRENT_SCOPE_KEY = "[mongoid]:current-scope"
+    CURRENT_SCOPE_KEY = '[mongoid]:current-scope'
 
-    AUTOSAVES_KEY = "[mongoid]:autosaves"
-    VALIDATIONS_KEY = "[mongoid]:validations"
+    AUTOSAVES_KEY = '[mongoid]:autosaves'
+    VALIDATIONS_KEY = '[mongoid]:validations'
 
     STACK_KEYS = Hash.new do |hash, key|
       hash[key] = "[mongoid]:#{key}-stack"
     end
 
+    # The key for the current thread's sessions.
+    SESSIONS_KEY = '[mongoid]:sessions'
+
+    # The key for storing documents modified inside transactions.
+    MODIFIED_DOCUMENTS_KEY = '[mongoid]:modified-documents'
+
     # The key storing the default value for whether or not callbacks are
     # executed on documents.
     EXECUTE_CALLBACKS = '[mongoid]:execute-callbacks'
@@ -230,10 +234,7 @@ module Mongoid
     # @return [ Criteria ] The scope.
     def set_current_scope(scope, klass)
       if scope.nil?
-        if Thread.current[CURRENT_SCOPE_KEY]
-          Thread.current[CURRENT_SCOPE_KEY].delete(klass)
-          Thread.current[CURRENT_SCOPE_KEY] = nil if Thread.current[CURRENT_SCOPE_KEY].empty?
-        end
+        unset_current_scope(klass)
       else
         Thread.current[CURRENT_SCOPE_KEY] ||= {}
         Thread.current[CURRENT_SCOPE_KEY][klass] = scope
@@ -307,6 +308,7 @@ module Mongoid
     def autosaves_for(klass)
       autosaves[klass] ||= []
     end
+
     # Get all validations on the current thread for the class.
     #
     # @example Get all validations.
@@ -319,36 +321,64 @@ module Mongoid
       validations[klass] ||= []
     end
 
-    # Cache a session for this thread.
+    # Cache a session for this thread for a client.
     #
-    # @example Save a session for this thread.
-    #   Threaded.set_session(session)
+    # @note For backward compatibility it is allowed to call this method without
+    # specifying `client` parameter.
     #
     # @param [ Mongo::Session ] session The session to save.
-    def set_session(session)
-      Thread.current["[mongoid]:session"] = session
+    # @param [ Mongo::Client | nil ] client The client to cache the session for.
+    def set_session(session, client: nil)
+      sessions[client] = session
     end
 
-    # Get the cached session for this thread.
+    # Get the cached session for this thread for a client.
+    #
+    # @note For backward compatibility it is allowed to call this method without
+    # specifying `client` parameter.
     #
-    # @example Get the session for this thread.
-    #   Threaded.get_session
+    # @param [ Mongo::Client | nil ] client The client to cache the session for.
     #
     # @return [ Mongo::Session | nil ] The session cached on this thread or nil.
-    def get_session
-      Thread.current["[mongoid]:session"]
+    def get_session(client: nil)
+      sessions[client]
     end
 
-    # Clear the cached session for this thread.
+    # Clear the cached session for this thread for a client.
     #
-    # @example Clear this thread's session.
-    #   Threaded.clear_session
+    # @note For backward compatibility it is allowed to call this method without
+    # specifying `client` parameter.
+    #
+    # @param [ Mongo::Client | nil ] client The client to clear the session for.
     #
     # @return [ nil ]
-    def clear_session
-      session = get_session
-      session.end_session if session
-      Thread.current["[mongoid]:session"] = nil
+    def clear_session(client: nil)
+      sessions.delete(client)&.end_session
+    end
+
+    # Store a reference to the document that was modified inside a transaction
+    # associated with the session.
+    #
+    # @param [ Mongo::Session ] session Session in scope of which the document
+    #   was modified.
+    # @param [ Mongoid::Document ] document Mongoid document that was modified.
+    def add_modified_document(session, document)
+      return unless session&.in_transaction?
+
+      modified_documents[session] << document
+    end
+
+    # Clears the set of modified documents for the given session, and return the
+    # content of the set before the clearance.
+    # @param [ Mongo::Session ] session Session for which the modified documents
+    #   set should be cleared.
+    #
+    # @return [ Set<Mongoid::Document> ] Collection of modified documents before
+    #   it was cleared.
+    def clear_modified_documents(session)
+      modified_documents[session].dup
+    ensure
+      modified_documents[session].clear
     end
 
     # Queries whether document callbacks should be executed by default for the
@@ -376,5 +406,39 @@ module Mongoid
     def execute_callbacks=(flag)
       Thread.current[EXECUTE_CALLBACKS] = flag
     end
+
+    # Returns the thread store of sessions.
+    #
+    # @return [ Hash<Integer, Set> ] The sessions indexed by client object ID.
+    #
+    # @api private
+    def sessions
+      Thread.current[SESSIONS_KEY] ||= {}.compare_by_identity
+    end
+
+    # Returns the thread store of modified documents.
+    #
+    # @return [ Hash<Mongo::Session, Set<Mongoid::Document>> ] The modified
+    #   documents indexed by session.
+    #
+    # @api private
+    def modified_documents
+      Thread.current[MODIFIED_DOCUMENTS_KEY] ||= Hash.new do |h, k|
+        h[k] = Set.new
+      end
+    end
+
+    private
+
+    # Removes the given klass from the current scope, and tidies the current
+    # scope list.
+    #
+    # @param klass [ Class ] the class to remove from the current scope.
+    def unset_current_scope(klass)
+      return unless Thread.current[CURRENT_SCOPE_KEY]
+
+      Thread.current[CURRENT_SCOPE_KEY].delete(klass)
+      Thread.current[CURRENT_SCOPE_KEY] = nil if Thread.current[CURRENT_SCOPE_KEY].empty?
+    end
   end
 end

data/lib/mongoid/timestamps/created/short.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 module Mongoid
   module Timestamps

data/lib/mongoid/timestamps/created.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 require "mongoid/timestamps/created/short"
 
@@ -23,9 +24,9 @@ module Mongoid
       #   person.set_created_at
       def set_created_at
         if !timeless? && !created_at
-          time = Time.configured.now
-          self.updated_at = time if is_a?(Updated) && !updated_at_changed?
-          self.created_at = time
+          now = Time.current
+          self.updated_at = now if is_a?(Updated) && !updated_at_changed?
+          self.created_at = now
         end
         clear_timeless_option
       end

data/lib/mongoid/timestamps/short.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 module Mongoid
   module Timestamps

data/lib/mongoid/timestamps/timeless.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 module Mongoid
   module Timestamps
@@ -34,6 +35,10 @@ module Mongoid
         self
       end
 
+      # Returns whether the document should skip timestamping.
+      #
+      # @return [ true | false ] Whether the document should
+      #   skip timestamping.
       def timeless?
         self.class.timeless?
       end
@@ -41,6 +46,12 @@ module Mongoid
       class << self
         extend Forwardable
 
+        # Returns the in-memory thread cache of classes
+        # for which to skip timestamping.
+        #
+        # @return [ Hash ] The timeless table.
+        #
+        # @api private
         def timeless_table
           Thread.current['[mongoid]:timeless'] ||= Hash.new
         end
@@ -66,6 +77,9 @@ module Mongoid
           self
         end
 
+        # Removes the timeless option on the current class.
+        #
+        # @return [ true ] Always true.
         def clear_timeless_option
           if counter = Timeless[name]
             counter -= 1
@@ -74,6 +88,10 @@ module Mongoid
           true
         end
 
+        # Sets to remove the timeless option when the next
+        # instance of the current class is updated.
+        #
+        # @return [ true ] Always true.
         def clear_timeless_option_on_update
           if counter = Timeless[name]
             counter -= 1 if self < Mongoid::Timestamps::Created
@@ -82,10 +100,21 @@ module Mongoid
           end
         end
 
+        # Clears the timeless counter for the current class
+        # if the value has reached zero.
+        #
+        # @param [ Integer ] counter The counter value.
+        #
+        # @return [ Integer | nil ] The counter value, or nil
+        #   if the counter was cleared.
         def set_timeless_counter(counter)
           Timeless[name] = (counter == 0) ? nil : counter
         end
 
+        # Returns whether the current class should skip timestamping.
+        #
+        # @return [ true | false ] Whether the current class should
+        #   skip timestamping.
         def timeless?
           !!Timeless[name]
         end

data/lib/mongoid/timestamps/updated/short.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 module Mongoid
   module Timestamps

data/lib/mongoid/timestamps/updated.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 require "mongoid/timestamps/updated/short"
 
@@ -24,7 +25,7 @@ module Mongoid
       #   person.set_updated_at
       def set_updated_at
         if able_to_set_updated_at?
-          self.updated_at = Time.configured.now unless updated_at_changed?
+          self.updated_at = Time.current unless updated_at_changed?
         end
 
         clear_timeless_option

data/lib/mongoid/timestamps.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+# rubocop:todo all
 
 require "mongoid/timestamps/timeless"
 require "mongoid/timestamps/created"