MONGOID-5592 Allow key vault client configuration (mongodb#5602)

lib/mongoid/clients/factory.rb

@@ -58,15 +58,8 @@ def create_client(configuration)
         database = config.delete(:database)
         hosts = config.delete(:hosts)
         opts = config.delete(:options) || {}
-        if opts[:auto_encryption_options]
-          if opts[:auto_encryption_options].key?(:schema_map)
-            Mongoid.logger.warn(
-              'The :schema_map is configured in the :auto_encryption_options for the client;' +
-              ' encryption setting in Mongoid documents will be ignored.'
-            )
-          else
-            opts[:auto_encryption_options][:schema_map] = Mongoid.config.encryption_schema_map(database)
-          end
+        if opts.key?(:auto_encryption_options)
+          opts[:auto_encryption_options] = build_auto_encryption_options(opts, database)
         end
         unless config.empty?
           default_logger.warn("Unknown config options detected: #{config}.")
@@ -81,6 +74,36 @@ def create_client(configuration)
         end
       end
 
+      # Build auto encryption options for the client based on the options
+      # provided in the Mongoid client configuration and the encryption
+      # schema map for the database.
+      #
+      # @param [ Hash ] opts Options from the Mongoid client configuration.
+      # @param [ String ] database Database name to use for encryption schema map.
+      #
+      # @return [ Hash | nil ] Auto encryption options for the client.
+      #
+      # @api private
+      def build_auto_encryption_options(opts, database)
+        return nil unless opts[:auto_encryption_options]
+
+        opts[:auto_encryption_options].dup.tap do |auto_encryption_options|
+          if auto_encryption_options.key?(:schema_map)
+            default_logger.warn(
+              'The :schema_map is configured in the :auto_encryption_options for the client;' +
+              ' encryption setting in Mongoid documents will be ignored.'
+            )
+          else
+            auto_encryption_options[:schema_map] = Mongoid.config.encryption_schema_map(database)
+          end
+          if auto_encryption_options.key?(:key_vault_client)
+            auto_encryption_options[:key_vault_client] = Mongoid::Clients.with_name(
+              auto_encryption_options[:key_vault_client]
+            )
+          end
+        end
+      end
+
       MONGOID_WRAPPING_LIBRARY = {
         name: 'Mongoid',
         version: VERSION,

spec/integration/encryption_spec.rb

@@ -10,11 +10,13 @@
   let(:config) do
     {
       default: { hosts: SpecConfig.instance.addresses, database: database_id },
+      key_vault: { hosts: SpecConfig.instance.addresses, database: :key_vault },
       encrypted: {
         hosts: SpecConfig.instance.addresses,
         database: database_id,
         options: {
           auto_encryption_options: {
+            key_vault_client: :key_vault,
             kms_providers: kms_providers,
             key_vault_namespace: key_vault_namespace,
             extra_options: extra_options
@@ -29,11 +31,11 @@
   end
 
   around do |example|
-    key_vault_client[key_vault_collection].drop
     Mongoid.default_client[Crypt::Patient.collection_name].drop
     existing_key_id = Crypt::Patient.encrypt_metadata[:key_id]
     Crypt::Patient.set_key_id(data_key_id)
     Mongoid::Config.send(:clients=, config)
+    Mongoid::Clients.with_name(:key_vault)[key_vault_collection].drop
     Crypt::Patient.store_in(client: :encrypted)
 
     example.run

spec/mongoid/clients/factory_spec.rb

@@ -252,23 +252,6 @@
         restore_config_clients
         include_context 'with encryption'
 
-        let(:config) do
-          {
-            default: { hosts: SpecConfig.instance.addresses, database: database_id },
-            encrypted: {
-              hosts: SpecConfig.instance.addresses,
-              database: database_id,
-              options: {
-                auto_encryption_options: {
-                  kms_providers: kms_providers,
-                  key_vault_namespace: key_vault_namespace,
-                  extra_options: extra_options
-                }
-              }
-            }
-          }
-        end
-
         before do
           Mongoid::Config.send(:clients=, config)
           key_vault_client[key_vault_collection].drop
@@ -282,12 +265,60 @@
           described_class.create(:encrypted)
         end
 
-        it "returns a client" do
-          expect(client).to be_a(Mongo::Client)
+        context 'when no key vault client is provided' do
+          let(:config) do
+            {
+              default: { hosts: SpecConfig.instance.addresses, database: database_id },
+              encrypted: {
+                hosts: SpecConfig.instance.addresses,
+                database: database_id,
+                options: {
+                  auto_encryption_options: {
+                    kms_providers: kms_providers,
+                    key_vault_namespace: key_vault_namespace,
+                    extra_options: extra_options
+                  }
+                }
+              }
+            }
+          end
+
+          it "returns a client" do
+            expect(client).to be_a(Mongo::Client)
+          end
+
+          it 'sets schema_map for the client' do
+            expect(client.options[:auto_encryption_options][:schema_map]).not_to be_nil
+          end
         end
 
-        it 'sets schema_map for the client' do
-          expect(client.options[:auto_encryption_options][:schema_map]).not_to be_nil
+        context 'when a key vault client is provided' do
+          let(:config) do
+            {
+              default: { hosts: SpecConfig.instance.addresses, database: database_id },
+              key_vault: { hosts: SpecConfig.instance.addresses, database: database_id },
+              encrypted: {
+                hosts: SpecConfig.instance.addresses,
+                database: database_id,
+                options: {
+                  auto_encryption_options: {
+                    key_vault_client: :key_vault,
+                    kms_providers: kms_providers,
+                    key_vault_namespace: key_vault_namespace,
+                    extra_options: extra_options
+                  }
+                }
+              }
+            }
+          end
+
+          it "returns a client" do
+            expect(client).to be_a(Mongo::Client)
+          end
+
+          it 'sets key_vault_client option for the client' do
+            expect(client.options[:auto_encryption_options][:key_vault_client]).to eq(Mongoid::Clients.with_name(:key_vault))
+          end
         end
       end
     end