chore: create multiplexed session at startup

Author: olavloite

lib/activerecord_spanner_adapter/connection.rb

@@ -17,14 +17,22 @@ class Connection
     attr_accessor :isolation_level
 
     def initialize config
+      config = config.symbolize_keys
+      @config = config
       @instance_id = config[:instance]
       @database_id = config[:database]
       @isolation_level = config[:isolation_level]
       @spanner = self.class.spanners config
+      begin
+        @mux_session = self.class.mux_sessions @spanner, config unless config[:skip_create_multiplexed_session]
+      rescue Google::Cloud::NotFoundError
+        @mux_session = nil
+      rescue Google::Cloud::UnimplementedError
+        @mux_session_unimplemented = true
+      end
     end
 
     def self.spanners config
-      config = config.symbolize_keys
       @spanners ||= {}
       @mutex ||= Mutex.new
       @mutex.synchronize do
@@ -40,6 +48,17 @@ def self.spanners config
       end
     end
 
+    def self.mux_sessions spanner, config
+      instance_id = config[:instance]
+      database_id = config[:database]
+      @mux_sessions ||= {}
+      @mutex ||= Mutex.new
+      @mutex.synchronize do
+        @mux_sessions[database_path(config)] ||=
+          spanner.create_multiplexed_session instance_id, database_id
+      end
+    end
+
     # Clears the cached information about the underlying information schemas.
     # Call this method if you drop and recreate a database with the same name
     # to prevent the cached information to be used for the new database.
@@ -96,6 +115,9 @@ def create_database
       job = spanner.create_database instance_id, database_id
       job.wait_until_done!
       raise Google::Cloud::Error.from_error job.error if job.error?
+      unless @config[:skip_create_multiplexed_session] || @mux_session_unimplemented
+        @mux_session = self.class.mux_sessions @spanner, @config
+      end
       job.database
     end
 

lib/spanner_client_ext.rb

@@ -20,6 +20,30 @@ def create_session instance_id, database_id, labels: nil
           )
           Session.from_grpc grpc, @service
         end
+
+        def create_multiplexed_session instance_id, database_id
+          ensure_service!
+
+          grpc = @service._create_multiplexed_session(
+            Admin::Database::V1::DatabaseAdmin::Paths.database_path(
+              project: project, instance: instance_id, database: database_id
+            )
+          )
+          Session.from_grpc grpc, @service
+        end
+      end
+
+      class Service
+        def _create_multiplexed_session database_name, database_role: nil
+          route_to_leader = LARHeaders.create_session
+          opts = default_options session_name: database_name,
+                                 route_to_leader: route_to_leader
+          session = Google::Cloud::Spanner::V1::Session.new multiplexed: true,
+                                                            creator_role: database_role
+          service.create_session(
+            { database: database_name, session: session }, opts
+          )
+        end
       end
 
       class Session

test/activerecord_spanner_adapter/connection_mock_server_test.rb

@@ -47,10 +47,11 @@ def create_connection
       _(row_count).must_equal 1
       connection.disconnect!
 
-      _(@mock.requests.length).must_equal 3
+      _(@mock.requests.length).must_equal 4
       _(@mock.requests[0]).must_be_kind_of Google::Cloud::Spanner::V1::CreateSessionRequest
-      _(@mock.requests[1]).must_be_kind_of Google::Cloud::Spanner::V1::ExecuteSqlRequest
-      _(@mock.requests[2]).must_be_kind_of Google::Cloud::Spanner::V1::DeleteSessionRequest
+      _(@mock.requests[1]).must_be_kind_of Google::Cloud::Spanner::V1::CreateSessionRequest
+      _(@mock.requests[2]).must_be_kind_of Google::Cloud::Spanner::V1::ExecuteSqlRequest
+      _(@mock.requests[3]).must_be_kind_of Google::Cloud::Spanner::V1::DeleteSessionRequest
     end
 
     it "can execute random query" do
@@ -84,10 +85,11 @@ def create_connection
       end
       connection.disconnect!
 
-      _(@mock.requests.length).must_equal 3
+      _(@mock.requests.length).must_equal 4
       _(@mock.requests[0]).must_be_kind_of Google::Cloud::Spanner::V1::CreateSessionRequest
-      _(@mock.requests[1]).must_be_kind_of Google::Cloud::Spanner::V1::ExecuteSqlRequest
-      _(@mock.requests[2]).must_be_kind_of Google::Cloud::Spanner::V1::DeleteSessionRequest
+      _(@mock.requests[1]).must_be_kind_of Google::Cloud::Spanner::V1::CreateSessionRequest
+      _(@mock.requests[2]).must_be_kind_of Google::Cloud::Spanner::V1::ExecuteSqlRequest
+      _(@mock.requests[3]).must_be_kind_of Google::Cloud::Spanner::V1::DeleteSessionRequest
     end
 
     it "can execute transaction" do

test/activerecord_spanner_adapter/connection_test.rb

@@ -13,7 +13,8 @@ def test_create_connection
         project: project_id,
         instance: instance_id,
         database: database_id,
-        credentials: credentials
+        credentials: credentials,
+        skip_create_multiplexed_session: true
       })
     end
 

test/mock_server/spanner_mock_server.rb

@@ -40,7 +40,8 @@ def push_error sql_or_method, error
 
   def create_session request, _unused_call
     @requests << request
-    do_create_session request.database
+    multiplexed = request.session&.multiplexed
+    do_create_session request.database, multiplexed: multiplexed
   end
 
   def batch_create_sessions request, _unused_call
@@ -199,9 +200,9 @@ def validate_session session
     end
   end
 
-  def do_create_session database
+  def do_create_session database, multiplexed: false
     name = "#{database}/sessions/#{SecureRandom.uuid}"
-    session = Google::Cloud::Spanner::V1::Session.new name: name
+    session = Google::Cloud::Spanner::V1::Session.new name: name, multiplexed: multiplexed
     @sessions[name] = session
     session
   end

test/test_helper.rb

@@ -57,6 +57,7 @@ def setup
           instance: instance_id,
           database: database_id,
           credentials: credentials,
+          skip_create_multiplexed_session: true
         )
       end
     end