diff --git a/lib/activerecord_spanner_adapter/connection.rb b/lib/activerecord_spanner_adapter/connection.rb index e82e8e11..d359d0be 100644 --- a/lib/activerecord_spanner_adapter/connection.rb +++ b/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 diff --git a/lib/spanner_client_ext.rb b/lib/spanner_client_ext.rb index fcc90797..d5cc5348 100644 --- a/lib/spanner_client_ext.rb +++ b/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 diff --git a/test/activerecord_spanner_adapter/connection_mock_server_test.rb b/test/activerecord_spanner_adapter/connection_mock_server_test.rb index 2106cc2d..89e74eb8 100644 --- a/test/activerecord_spanner_adapter/connection_mock_server_test.rb +++ b/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 diff --git a/test/activerecord_spanner_adapter/connection_test.rb b/test/activerecord_spanner_adapter/connection_test.rb index dbf33735..06304a67 100644 --- a/test/activerecord_spanner_adapter/connection_test.rb +++ b/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 diff --git a/test/mock_server/spanner_mock_server.rb b/test/mock_server/spanner_mock_server.rb index 94014529..39397e91 100644 --- a/test/mock_server/spanner_mock_server.rb +++ b/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 diff --git a/test/test_helper.rb b/test/test_helper.rb index fd80dd4d..1e689652 100644 --- a/test/test_helper.rb +++ b/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