Add Session labels to Project#client and #batch_client

Convert labels argument from Google::Protobuf::Map.
Rescue GRPC error in Results.from_enum.
Update expected errors in acceptance tests.
Add error case unit tests for Session#reload! and #keepalive!

[closes googleapis#1817]

google-cloud-spanner/acceptance/spanner/client/edge_test.rb

@@ -26,13 +26,13 @@
   end
 
   it "reads with invalid table fails" do
-    assert_raises GRPC::NotFound do
+    assert_raises Google::Cloud::NotFoundError do
       db.read "invalid_table", [:id, :int]
     end
   end
 
   it "reads with invalid column fails" do
-    assert_raises GRPC::NotFound do
+    assert_raises Google::Cloud::NotFoundError do
       db.read table_name, [:id, :invalid]
     end
   end
@@ -59,19 +59,19 @@
   end
 
   it "queries to a non-existing table fails" do
-    assert_raises GRPC::InvalidArgument do
+    assert_raises Google::Cloud::InvalidArgumentError do
       db.execute "SELECT id, name FROM invalid_table"
     end
   end
 
   it "queries to a non-existing column fails" do
-    assert_raises GRPC::InvalidArgument do
+    assert_raises Google::Cloud::InvalidArgumentError do
       db.execute "SELECT id, name FROM #{table_name}"
     end
   end
 
   it "queries with bad SQL fails" do
-    assert_raises GRPC::InvalidArgument do
+    assert_raises Google::Cloud::InvalidArgumentError do
       db.execute "SELECT Apples AND Oranges"
     end
   end

google-cloud-spanner/lib/google/cloud/spanner/batch_client.rb

@@ -65,10 +65,11 @@ module Spanner
       class BatchClient
         ##
         # @private Creates a new Spanner BatchClient instance.
-        def initialize project, instance_id, database_id
+        def initialize project, instance_id, database_id, session_labels: nil
           @project = project
           @instance_id = instance_id
           @database_id = database_id
+          @session_labels = session_labels
         end
 
         # The unique identifier for the project.
@@ -406,7 +407,8 @@ def session
           grpc = @project.service.create_session \
             Admin::Database::V1::DatabaseAdminClient.database_path(
               project_id, instance_id, database_id
-            )
+            ),
+            labels: @session_labels
           Session.from_grpc(grpc, @project.service)
         end
 

google-cloud-spanner/lib/google/cloud/spanner/client.rb

@@ -52,11 +52,13 @@ module Spanner
       class Client
         ##
         # @private Creates a new Spanner Client instance.
-        def initialize project, instance_id, database_id, opts = {}
+        def initialize project, instance_id, database_id, session_labels: nil,
+                       pool_opts: {}
           @project = project
           @instance_id = instance_id
           @database_id = database_id
-          @pool = Pool.new self, opts
+          @session_labels = session_labels
+          @pool = Pool.new self, pool_opts
         end
 
         # The unique identifier for the project.
@@ -1091,7 +1093,8 @@ def create_new_session
           grpc = @project.service.create_session \
             Admin::Database::V1::DatabaseAdminClient.database_path(
               project_id, instance_id, database_id
-            )
+            ),
+            labels: @session_labels
           Session.from_grpc(grpc, @project.service)
         end
 

google-cloud-spanner/lib/google/cloud/spanner/project.rb

@@ -219,6 +219,8 @@ def instance instance_id
         def create_instance instance_id, name: nil, config: nil, nodes: nil,
                             labels: nil
           config = config.path if config.respond_to? :path
+          # Convert from possible Google::Protobuf::Map
+          labels = Hash[labels.map { |k, v| [String(k), String(v)] }] if labels
           grpc = service.create_instance \
             instance_id, name: name, config: config, nodes: nodes,
                          labels: labels
@@ -435,6 +437,21 @@ def create_database instance_id, database_id, statements: []
         #     becomes available. The default is `true`.
         #   * `:threads` (Integer) The number of threads in the thread pool. The
         #     default is twice the number of available CPUs.
+        # @param [Hash] labels The labels to be applied to all sessions
+        #   created by the client. Cloud Labels are a flexible and lightweight
+        #   mechanism for organizing cloud resources into groups that reflect a
+        #   customer's organizational needs and deployment strategies. Cloud
+        #   Labels can be used to filter collections of resources. They can be
+        #   used to control how resource metrics are aggregated. And they can be
+        #   used as arguments to policy management rules (e.g. route, firewall,
+        #   load balancing, etc.). Optional. The default is `nil`.
+        #
+        #   * Label keys must be between 1 and 63 characters long and must
+        #     conform to the following regular expression:
+        #     `[a-z]([-a-z0-9]*[a-z0-9])?`.
+        #   * Label values must be between 0 and 63 characters long and must
+        #     conform to the regular expression `([a-z]([-a-z0-9]*[a-z0-9])?)?`.
+        #   * No more than 64 labels can be associated with a given resource.
         #
         # @return [Client] The newly created client.
         #
@@ -453,9 +470,12 @@ def create_database instance_id, database_id, statements: []
         #     end
         #   end
         #
-        def client instance_id, database_id, pool: {}
+        def client instance_id, database_id, pool: {}, labels: nil
+          # Convert from possible Google::Protobuf::Map
+          labels = Hash[labels.map { |k, v| [String(k), String(v)] }] if labels
           Client.new self, instance_id, database_id,
-                     valid_session_pool_options(pool)
+                     session_labels: labels,
+                     pool_opts: valid_session_pool_options(pool)
         end
 
         ##
@@ -466,6 +486,21 @@ def client instance_id, database_id, pool: {}
         #   Required.
         # @param [String] database_id The unique identifier for the database.
         #   Required.
+        # @param [Hash] labels The labels to be applied to all sessions
+        #   created by the batch client. Labels are a flexible and lightweight
+        #   mechanism for organizing cloud resources into groups that reflect a
+        #   customer's organizational needs and deployment strategies. Cloud
+        #   Labels can be used to filter collections of resources. They can be
+        #   used to control how resource metrics are aggregated. And they can be
+        #   used as arguments to policy management rules (e.g. route, firewall,
+        #   load balancing, etc.). Optional. The default is `nil`.
+        #
+        #   * Label keys must be between 1 and 63 characters long and must
+        #     conform to the following regular expression:
+        #     `[a-z]([-a-z0-9]*[a-z0-9])?`.
+        #   * Label values must be between 0 and 63 characters long and must
+        #     conform to the regular expression `([a-z]([-a-z0-9]*[a-z0-9])?)?`.
+        #   * No more than 64 labels can be associated with a given resource.
         #
         # @return [Client] The newly created client.
         #
@@ -494,8 +529,10 @@ def client instance_id, database_id, pool: {}
         #   results = new_batch_snapshot.execute_partition \
         #     new_partition
         #
-        def batch_client instance_id, database_id
-          BatchClient.new self, instance_id, database_id
+        def batch_client instance_id, database_id, labels: nil
+          # Convert from possible Google::Protobuf::Map
+          labels = Hash[labels.map { |k, v| [String(k), String(v)] }] if labels
+          BatchClient.new self, instance_id, database_id, session_labels: labels
         end
 
         protected

google-cloud-spanner/lib/google/cloud/spanner/results.rb

@@ -206,6 +206,8 @@ def self.from_enum enum, service
             results.instance_variable_set :@enum,     enum
             results.instance_variable_set :@service,  service
           end
+        rescue GRPC::BadStatus => e
+          raise Google::Cloud::Error.from_error(e)
         end
 
         # @private

google-cloud-spanner/lib/google/cloud/spanner/service.rb

@@ -241,10 +241,12 @@ def get_session session_name
           end
         end
 
-        def create_session database_name
+        def create_session database_name, labels: nil
           opts = default_options_from_session database_name
+          session = Google::Spanner::V1::Session.new(labels: labels) if labels
           execute do
-            service.create_session database_name, options: opts
+            service.create_session database_name, session: session,
+                                                  options: opts
           end
         end
 

google-cloud-spanner/lib/google/cloud/spanner/session.rb

@@ -597,10 +597,12 @@ def reload!
           @last_updated_at = Time.now
           return self
         rescue Google::Cloud::NotFoundError
+          labels = @grpc.labels.to_h unless @grpc.labels.to_h.empty?
           @grpc = service.create_session \
             Admin::Database::V1::DatabaseAdminClient.database_path(
               project_id, instance_id, database_id
-            )
+            ),
+            labels: labels
           @last_updated_at = Time.now
           return self
         end
@@ -613,10 +615,12 @@ def keepalive!
           execute "SELECT 1"
           return true
         rescue Google::Cloud::NotFoundError
+          labels = @grpc.labels.to_h unless @grpc.labels.to_h.empty?
           @grpc = service.create_session \
             Admin::Database::V1::DatabaseAdminClient.database_path(
               project_id, instance_id, database_id
-            )
+            ),
+            labels: labels
           return false
         end
 

google-cloud-spanner/test/google/cloud/spanner/batch_client_test.rb

@@ -31,6 +31,9 @@
   let(:default_options) { Google::Gax::CallOptions.new kwargs: { "google-cloud-resource-prefix" => database_path(instance_id, database_id) } }
   let(:batch_client) { spanner.batch_client instance_id, database_id }
 
+  let(:labels) { { "env" => "production" } }
+  let(:batch_client_labels) { spanner.batch_client instance_id, database_id, labels: labels }
+
   it "knows its project_id" do
     batch_client.project_id.must_equal project
   end
@@ -80,7 +83,7 @@
 
   it "creates a batch_snapshot" do
     mock = Minitest::Mock.new
-    mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+    mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
     mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
     spanner.service.mocked_service = mock
 
@@ -93,12 +96,29 @@
     batch_snapshot.session.path.must_equal session.path
   end
 
+  it "creates a batch_snapshot with session labels" do
+    mock = Minitest::Mock.new
+    session_labels_grpc = Google::Spanner::V1::Session.new labels: labels
+    session_labels_resp_grpc = Google::Spanner::V1::Session.new name: session_path(instance_id, database_id, session_id), labels: labels
+    mock.expect :create_session, session_labels_resp_grpc, [database_path(instance_id, database_id), session: session_labels_grpc, options: default_options]
+    mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
+    spanner.service.mocked_service = mock
+
+    batch_snapshot = batch_client_labels.batch_snapshot
+
+    mock.verify
+
+    batch_snapshot.transaction_id.must_equal transaction_id
+    batch_snapshot.timestamp.must_equal timestamp_time
+    batch_snapshot.session.path.must_equal session.path
+  end
+
   describe :strong do
     let(:snp_opts) { Google::Spanner::V1::TransactionOptions::ReadOnly.new strong: true, return_read_timestamp: true }
 
     it "creates a batch_snapshot with strong timestamp bound" do
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -120,7 +140,7 @@
 
     it "creates a batch_snapshot with timestamp option (Time)" do
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -135,7 +155,7 @@
 
     it "creates a batch_snapshot with read_timestamp option (Time)" do
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -150,7 +170,7 @@
 
     it "creates a batch_snapshot with timestamp option (DateTime)" do
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -165,7 +185,7 @@
 
     it "creates a batch_snapshot with read_timestamp option (DateTime)" do
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -186,7 +206,7 @@
 
     it "creates a batch_snapshot with the staleness option" do
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -201,7 +221,7 @@
 
     it "creates a batch_snapshot with the exact_staleness option" do
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :begin_transaction, transaction_grpc, [session_grpc.name, tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 

google-cloud-spanner/test/google/cloud/spanner/client/commit_field_values_test.rb

@@ -56,7 +56,7 @@
       ]
 
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :commit, commit_resp, [session_grpc.name, mutations, transaction_id: nil, single_use_transaction: tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -84,7 +84,7 @@
       ]
 
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :commit, commit_resp, [session_grpc.name, mutations, transaction_id: nil, single_use_transaction: tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -107,7 +107,7 @@
       ]
 
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :commit, commit_resp, [session_grpc.name, mutations, transaction_id: nil, single_use_transaction: tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -130,7 +130,7 @@
       ]
 
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :commit, commit_resp, [session_grpc.name, mutations, transaction_id: nil, single_use_transaction: tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -153,7 +153,7 @@
       ]
 
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :commit, commit_resp, [session_grpc.name, mutations, transaction_id: nil, single_use_transaction: tx_opts, options: default_options]
       spanner.service.mocked_service = mock
 
@@ -176,7 +176,7 @@
       ]
 
       mock = Minitest::Mock.new
-      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), options: default_options]
+      mock.expect :create_session, session_grpc, [database_path(instance_id, database_id), session: nil, options: default_options]
       mock.expect :commit, commit_resp, [session_grpc.name, mutations, transaction_id: nil, single_use_transaction: tx_opts, options: default_options]
       spanner.service.mocked_service = mock