feat: support for common interface to get a session

src/database.ts

@@ -102,6 +102,7 @@ import Policy = google.iam.v1.Policy;
 import FieldMask = google.protobuf.FieldMask;
 import IDatabase = google.spanner.admin.database.v1.IDatabase;
 import snakeCase = require('lodash.snakecase');
+import {SessionFactory, SessionFactoryInterface} from './session-factory';
 import {
   ObservabilityOptions,
   Span,
@@ -338,7 +339,7 @@ export interface RestoreOptions {
 class Database extends common.GrpcServiceObject {
   private instance: Instance;
   formattedName_: string;
-  pool_: SessionPoolInterface;
+  sessionFactory_: SessionFactoryInterface;
   queryOptions_?: spannerClient.spanner.v1.ExecuteSqlRequest.IQueryOptions;
   resourceHeader_: {[k: string]: string};
   request: DatabaseRequest;
@@ -450,22 +451,12 @@ class Database extends common.GrpcServiceObject {
       },
     } as {} as ServiceObjectConfig);
 
-    this.pool_ =
-      typeof poolOptions === 'function'
-        ? new (poolOptions as SessionPoolConstructor)(this, null)
-        : new SessionPool(this, poolOptions);
-    const sessionPoolInstance = this.pool_ as SessionPool;
-    if (sessionPoolInstance) {
-      sessionPoolInstance._observabilityOptions =
-        instance._observabilityOptions;
-    }
     if (typeof poolOptions === 'object') {
       this.databaseRole = poolOptions.databaseRole || null;
       this.labels = poolOptions.labels || null;
     }
     this.formattedName_ = formattedName_;
     this.instance = instance;
-    this._observabilityOptions = instance._observabilityOptions;
     this._traceConfig = {
       opts: this._observabilityOptions,
       dbName: this.formattedName_,
@@ -478,8 +469,14 @@ class Database extends common.GrpcServiceObject {
     this._observabilityOptions = instance._observabilityOptions;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     this.requestStream = instance.requestStream as any;
-    this.pool_.on('error', this.emit.bind(this, 'error'));
-    this.pool_.open();
+    this.sessionFactory_ = new SessionFactory(this, name, poolOptions);
+    this.pool_ = this.sessionFactory_.getPool();
+    this.multiplexedSession_ = this.sessionFactory_.getMultiplexedSession();
+    const sessionPoolInstance = this.pool_ as SessionPool;
+    if (sessionPoolInstance) {
+      sessionPoolInstance._observabilityOptions =
+        instance._observabilityOptions;
+    }
     this.queryOptions_ = Object.assign(
       Object.assign({}, queryOptions),
       Database.getEnvironmentQueryOptions()

src/multiplexed-session.ts

@@ -17,7 +17,7 @@
 import {EventEmitter} from 'events';
 import {Database} from './database';
 import {Session} from './session';
-import {GetSessionCallback} from './session-pool';
+import {GetSessionCallback} from './session-factory';
 import {
   ObservabilityOptions,
   getActiveOrNoopSpan,
@@ -38,7 +38,7 @@ export const MUX_SESSION_CREATE_ERROR = 'mux-session-create-error';
  * @constructs MultiplexedSessionInterface
  * @param {Database} database The database to create a multiplexed session for.
  */
-export interface MultiplexedSessionInterface {
+export interface MultiplexedSessionInterface extends EventEmitter {
   /**
    * When called creates a multiplexed session.
    *

src/session-factory.ts

@@ -0,0 +1,88 @@
+import {Database, Session, Transaction} from '.';
+import {
+  MultiplexedSession,
+  MultiplexedSessionInterface,
+} from './multiplexed-session';
+import {
+  SessionPool,
+  SessionPoolInterface,
+  SessionPoolOptions,
+} from './session-pool';
+import {SessionPoolConstructor} from './database';
+import {ServiceObjectConfig} from '@google-cloud/common';
+const common = require('./common-grpc/service-object');
+
+/**
+ * @callback GetSessionCallback
+ * @param {?Error} error Request error, if any.
+ * @param {Session} session The read-write session.
+ * @param {Transaction} transaction The transaction object.
+ */
+export interface GetSessionCallback {
+  (
+    err: Error | null,
+    session?: Session | null,
+    transaction?: Transaction | null
+  ): void;
+}
+
+export interface SessionFactoryInterface {
+  getSession(callback: GetSessionCallback): void;
+  getPool(): SessionPoolInterface;
+  getMultiplexedSession(): MultiplexedSessionInterface | undefined;
+  release(session: Session): void;
+}
+
+export class SessionFactory
+  extends common.GrpcServiceObject
+  implements SessionFactoryInterface
+{
+  multiplexedSession_?: MultiplexedSessionInterface;
+  pool_: SessionPoolInterface;
+  constructor(
+    database: Database,
+    name: String,
+    poolOptions?: SessionPoolConstructor | SessionPoolOptions
+  ) {
+    super({
+      parent: database,
+      id: name,
+    } as {} as ServiceObjectConfig);
+    this.pool_ =
+      typeof poolOptions === 'function'
+        ? new (poolOptions as SessionPoolConstructor)(database, null)
+        : new SessionPool(database, poolOptions);
+    this.multiplexedSession_ = new MultiplexedSession(database);
+    this.pool_.on('error', this.emit.bind(this, 'error'));
+    this.pool_.open();
+    // multiplexed session should only get created if the env varaible is enabled
+    if (process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true') {
+      this.multiplexedSession_.on('error', this.emit.bind(this, 'error'));
+      this.multiplexedSession_.createSession();
+    }
+  }
+
+  getSession(callback: GetSessionCallback): void {
+    if (process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true') {
+      this.multiplexedSession_?.getSession((err, session) => {
+        err ? callback(err, null) : callback(null, session);
+      });
+    } else {
+      this.pool_?.getSession((err, session) => {
+        err ? callback(err, null) : callback(null, session);
+      });
+    }
+  }
+
+  getPool(): SessionPoolInterface {
+    return this.pool_;
+  }
+
+  getMultiplexedSession(): MultiplexedSessionInterface | undefined {
+    return this.multiplexedSession_;
+  }
+
+  release(session: Session): void {
+    this.pool_.release(session);
+  }
+}

src/session-pool.ts

@@ -24,6 +24,7 @@ import {Transaction} from './transaction';
 import {NormalCallback} from './common';
 import {GoogleError, grpc, ServiceError} from 'google-gax';
 import trace = require('stack-trace');
+import {GetSessionCallback} from './session-factory';
 import {
   ObservabilityOptions,
   getActiveOrNoopSpan,
@@ -59,20 +60,6 @@ export interface GetWriteSessionCallback {
   ): void;
 }
 
-/**
- * @callback GetSessionCallback
- * @param {?Error} error Request error, if any.
- * @param {Session} session The read-write session.
- * @param {Transaction} transaction The transaction object.
- */
-export interface GetSessionCallback {
-  (
-    err: Error | null,
-    session?: Session | null,
-    transaction?: Transaction | null
-  ): void;
-}
-
 /**
  * Interface for implementing custom session pooling logic, it should extend the
  * {@link https://nodejs.org/api/events.html|EventEmitter} class and emit any

test/database.ts

@@ -78,7 +78,7 @@ class FakeBatchTransaction {
   }
 }
 
-class FakeGrpcServiceObject extends EventEmitter {
+export class FakeGrpcServiceObject extends EventEmitter {
   calledWith_: IArguments;
   constructor() {
     super();
@@ -109,7 +109,7 @@ class FakeSession {
   }
 }
 
-class FakeSessionPool extends EventEmitter {
+export class FakeSessionPool extends EventEmitter {
   calledWith_: IArguments;
   constructor() {
     super();
@@ -295,45 +295,6 @@ describe('Database', () => {
       assert(database.formattedName_, formattedName);
     });
 
-    it('should create a SessionPool object', () => {
-      assert(database.pool_ instanceof FakeSessionPool);
-      assert.strictEqual(database.pool_.calledWith_[0], database);
-      assert.strictEqual(database.pool_.calledWith_[1], POOL_OPTIONS);
-    });
-
-    it('should accept a custom Pool class', () => {
-      function FakePool() {}
-      FakePool.prototype.on = util.noop;
-      FakePool.prototype.open = util.noop;
-
-      const database = new Database(
-        INSTANCE,
-        NAME,
-        FakePool as {} as db.SessionPoolConstructor
-      );
-      assert(database.pool_ instanceof FakePool);
-    });
-
-    it('should re-emit SessionPool errors', done => {
-      const error = new Error('err');
-
-      database.on('error', err => {
-        assert.strictEqual(err, error);
-        done();
-      });
-
-      database.pool_.emit('error', error);
-    });
-
-    it('should open the pool', done => {
-      FakeSessionPool.prototype.open = () => {
-        FakeSessionPool.prototype.open = util.noop;
-        done();
-      };
-
-      new Database(INSTANCE, NAME);
-    });
-
     it('should inherit from ServiceObject', done => {
       const options = {};