feat(session): add label support

src/database.js

@@ -344,16 +344,23 @@ class Database extends ServiceObject {
       callback = options;
       options = {};
     }
-    options = options || {};
+
+    const gaxOpts = extend({}, options);
     const reqOpts = {
       database: this.formattedName_,
     };
+
+    if (gaxOpts.labels) {
+      reqOpts.session = {labels: gaxOpts.labels};
+      delete gaxOpts.labels;
+    }
+
     this.request(
       {
         client: 'SpannerClient',
         method: 'createSession',
         reqOpts: reqOpts,
-        gaxOpts: options,
+        gaxOpts: gaxOpts,
       },
       (err, resp) => {
         if (err) {

src/session-pool.js

@@ -31,6 +31,7 @@ const DEFAULTS = {
   fail: false,
   idlesAfter: 10,
   keepAlive: 30,
+  labels: {},
   max: 100,
   maxIdle: 1,
   min: 0,
@@ -109,14 +110,16 @@ class WritePercentError extends TypeError {
  *     there are no available sessions for a request.
  * @property {number} [idlesAfter=10] How long until a resource becomes idle, in
  *     minutes.
+ * @property {number} [keepAlive=50] How often to ping idle sessions, in
+ *     minutes. Must be less than 1 hour.
+ * @property {Object<string, string>} [labels] Labels to apply to any session
+ *     created by the pool.
  * @property {number} [max=100] Maximum number of resources to create at any
  *     given time.
  * @property {number} [maxIdle=1] Maximum number of idle resources to keep in
  *     the pool at any given time.
  * @property {number} [min=0] Minimum number of resources to keep in the pool at
  *     any given time.
- * @property {number} [keepAlive=50] How often to ping idle sessions, in
- *     minutes. Must be less than 1 hour.
  * @property {number} [writes=0.0] Percentage of sessions to be pre-allocated as
  *     write sessions represented as a float.
  */
@@ -481,12 +484,13 @@ class SessionPool extends EventEmitter {
    */
   _createSession(type) {
     const session = this.database.session();
+    const labels = this.options.labels;
 
     this._inventory.borrowed.add(session);
 
     return this._requests
       .add(() => {
-        return session.create().then(() => {
+        return session.create({labels}).then(() => {
           if (type === READWRITE) {
             return this._prepareTransaction(session).catch(
               () => (type = READONLY)

src/session.js

@@ -172,12 +172,18 @@ class Session extends ServiceObject {
        */
       id: name,
       methods: methods,
-      createMethod: (options, callback) => {
-        database.createSession(options, (err, session, apiResponse) => {
+      createMethod: (_, options, callback) => {
+        if (is.fn(options)) {
+          callback = options;
+          options = {};
+        }
+
+        return database.createSession(options, (err, session, apiResponse) => {
           if (err) {
             callback(err, null, apiResponse);
             return;
           }
+
           extend(this, session);
           callback(null, this, apiResponse);
         });

test/database.js

@@ -1461,7 +1461,7 @@ describe('Database', () => {
           database: database.formattedName_,
         });
 
-        assert.strictEqual(config.gaxOpts, OPTIONS);
+        assert.deepStrictEqual(config.gaxOpts, OPTIONS);
 
         done();
       };
@@ -1483,6 +1483,20 @@ describe('Database', () => {
       database.createSession(assert.ifError);
     });
 
+    it('should send labels correctly', done => {
+      const labels = {a: 'b'};
+      const options = {a: 'b', labels};
+      const originalOptions = extend(true, {}, options);
+
+      database.request = function(config) {
+        assert.deepStrictEqual(config.reqOpts.session, {labels});
+        assert.deepStrictEqual(options, originalOptions);
+        done();
+      };
+
+      database.createSession({labels}, assert.ifError);
+    });
+
     describe('error', () => {
       const ERROR = new Error('Error.');
       const API_RESPONSE = {};

test/session-pool.js

@@ -49,8 +49,9 @@ class FakeSession {
   beginTransaction(options) {
     return new FakeTransaction(options).begin();
   }
-  create() {
+  create(options) {
     this.created = true;
+    this.createOptions = options;
     return Promise.resolve();
   }
   delete() {
@@ -235,6 +236,7 @@ describe('SessionPool', () => {
         assert.strictEqual(sessionPool.options.fail, false);
         assert.strictEqual(sessionPool.options.idlesAfter, 10);
         assert.strictEqual(sessionPool.options.keepAlive, 30);
+        assert.deepStrictEqual(sessionPool.options.labels, {});
         assert.strictEqual(sessionPool.options.max, 100);
         assert.strictEqual(sessionPool.options.maxIdle, 1);
         assert.strictEqual(sessionPool.options.min, 0);
@@ -866,6 +868,17 @@ describe('SessionPool', () => {
       });
     });
 
+    it('should pass along the session labels', () => {
+      const labels = {a: 'b'};
+
+      sessionPool.options.labels = labels;
+
+      return sessionPool._createSession('readonly').then(() => {
+        const session = sessionPool._inventory.readonly[0];
+        assert.deepStrictEqual(session.createOptions, {labels});
+      });
+    });
+
     it('should discard the session if unable to create', () => {
       const error = new Error('err');
 

test/session.js

@@ -135,15 +135,38 @@ describe('Session', () => {
         const session = new Session(databaseInstance, NAME);
         assert(session instanceof FakeGrpcServiceObject);
 
-        session.calledWith_[0].createMethod(options, (err, sess, resp) => {
-          assert.ifError(err);
+        session.calledWith_[0].createMethod(
+          null,
+          options,
+          (err, sess, resp) => {
+            assert.ifError(err);
 
-          assert.strictEqual(sess, session);
+            assert.strictEqual(sess, session);
 
-          assert.strictEqual(session.uniqueProperty, true);
+            assert.strictEqual(session.uniqueProperty, true);
 
-          assert.strictEqual(resp, apiResponse);
+            assert.strictEqual(resp, apiResponse);
+
+            done();
+          }
+        );
+      });
 
+      it('should check for options', done => {
+        const databaseInstance = extend({}, DATABASE, {
+          createSession: function(options, callback) {
+            assert.deepStrictEqual(options, {});
+            callback(null, {}, apiResponse);
+          },
+        });
+
+        const session = new Session(databaseInstance, NAME);
+        const apiResponse = {};
+
+        session.calledWith_[0].createMethod(null, (err, sess, resp) => {
+          assert.ifError(err);
+          assert.strictEqual(sess, session);
+          assert.strictEqual(resp, apiResponse);
           done();
         });
       });
@@ -161,7 +184,7 @@ describe('Session', () => {
         const session = new Session(databaseInstance, NAME);
         assert(session instanceof FakeGrpcServiceObject);
 
-        session.calledWith_[0].createMethod({}, (err, sess, resp) => {
+        session.calledWith_[0].createMethod(null, {}, (err, sess, resp) => {
           assert.strictEqual(err, error);
           assert.strictEqual(sess, null);
           assert.strictEqual(resp, apiResponse);