From a5d966663f336134e38a972257b18f70af801f67 Mon Sep 17 00:00:00 2001
From: Alka Trivedi <alkatrivedi12dec@gmail.com>
Date: Thu, 26 Dec 2024 13:04:35 +0530
Subject: [PATCH] refactor: dynamic fetching of env to in-memory state

---
 src/session-factory.ts  | 13 +++++----
 test/session-factory.ts | 58 +++++++++++++++++++++++------------------
 test/spanner.ts         | 16 +++++-------
 3 files changed, 47 insertions(+), 40 deletions(-)

diff --git a/src/session-factory.ts b/src/session-factory.ts
index 14c7af571..786e7d492 100644
--- a/src/session-factory.ts
+++ b/src/session-factory.ts
@@ -70,6 +70,7 @@ export class SessionFactory
 {
   multiplexedSession_?: MultiplexedSessionInterface;
   pool_: SessionPoolInterface;
+  isMuxCreated: boolean;
   constructor(
     database: Database,
     name: String,
@@ -85,8 +86,10 @@ export class SessionFactory
         : new SessionPool(database, poolOptions);
     this.pool_.on('error', this.emit.bind(database, 'error'));
     this.pool_.open();
+    this.isMuxCreated = false;
     // multiplexed session should only get created if the env variable is enabled
     if (process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true') {
+      this.isMuxCreated = true;
       this.multiplexedSession_ = new MultiplexedSession(database);
       this.multiplexedSession_.on('error', this.emit.bind(database, 'error'));
       this.multiplexedSession_.createSession();
@@ -103,10 +106,10 @@ export class SessionFactory
    */
 
   getSession(callback: GetSessionCallback): void {
-    const sessionHandler =
-      process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true'
-        ? this.multiplexedSession_
-        : this.pool_;
+    const sessionHandler = this.isMuxCreated
+      ? this.multiplexedSession_
+      : this.pool_;
+
     sessionHandler!.getSession((err, session) => callback(err, session));
   }
 
@@ -136,7 +139,7 @@ export class SessionFactory
    * @throws {Error} Throws an error if the session is invalid or cannot be released.
    */
   release(session: Session): void {
-    if (process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'false') {
+    if (!this.isMuxCreated) {
       this.pool_.release(session);
     }
   }
diff --git a/test/session-factory.ts b/test/session-factory.ts
index d172d9528..028cdff33 100644
--- a/test/session-factory.ts
+++ b/test/session-factory.ts
@@ -78,44 +78,50 @@ describe('SessionFactory', () => {
   });
 
   describe('instantiation', () => {
-    describe('SessionPool', () => {
-      before(() => {
-        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
-      });
+    before(() => {
+      process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+    });
 
-      it('should create a SessionPool object', () => {
-        assert(sessionFactory.pool_ instanceof SessionPool);
-      });
+    it('should create a SessionPool object', () => {
+      assert(sessionFactory.pool_ instanceof SessionPool);
+    });
 
-      it('should accept a custom Pool class', () => {
-        function FakePool() {}
-        FakePool.prototype.on = util.noop;
-        FakePool.prototype.open = util.noop;
+    it('should accept a custom Pool class', () => {
+      function FakePool() {}
+      FakePool.prototype.on = util.noop;
+      FakePool.prototype.open = util.noop;
+
+      const sessionFactory = new SessionFactory(
+        DATABASE,
+        NAME,
+        FakePool as {} as db.SessionPoolConstructor
+      );
+      assert(sessionFactory.pool_ instanceof FakePool);
+    });
 
-        const sessionFactory = new SessionFactory(
-          DATABASE,
-          NAME,
-          FakePool as {} as db.SessionPoolConstructor
-        );
-        assert(sessionFactory.pool_ instanceof FakePool);
-      });
+    it('should open the pool', () => {
+      const openStub = sandbox
+        .stub(SessionPool.prototype, 'open')
+        .callsFake(() => {});
 
-      it('should open the pool', () => {
-        const openStub = sandbox
-          .stub(SessionPool.prototype, 'open')
-          .callsFake(() => {});
+      new SessionFactory(DATABASE, NAME, POOL_OPTIONS);
 
-        new SessionFactory(DATABASE, NAME, POOL_OPTIONS);
+      assert.strictEqual(openStub.callCount, 1);
+    });
 
-        assert.strictEqual(openStub.callCount, 1);
-      });
+    it('should set the isMuxCreated to be false if env is disabled', () => {
+      assert.strictEqual(sessionFactory.isMuxCreated, false);
     });
 
-    describe('MultiplexedSession', () => {
+    describe('when env GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS is enabled', () => {
       before(() => {
         process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'true';
       });
 
+      it('should set the isMuxCreated to be true if env is enabled', () => {
+        assert.strictEqual(sessionFactory.isMuxCreated, true);
+      });
+
       it('should create a MultiplexedSession object', () => {
         assert(
           sessionFactory.multiplexedSession_ instanceof MultiplexedSession
diff --git a/test/spanner.ts b/test/spanner.ts
index aca33a9a3..15fb5aaca 100644
--- a/test/spanner.ts
+++ b/test/spanner.ts
@@ -5081,18 +5081,16 @@ describe('Spanner with mock server', () => {
     after(() => {
       process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
     });
-    it('should execute table mutations without leaking sessions', () => {
+
+    it('should not throw error when enabling env GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS after client initialization', done => {
       const database = newTestDatabase();
+      // enable env after database creation
       process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'true';
       const sessionFactory = database.sessionFactory_ as SessionFactory;
-      try {
-        sessionFactory.getSession(() => {});
-      } catch (e) {
-        assert.strictEqual(
-          (e as TypeError).message,
-          "Cannot read properties of undefined (reading 'getSession')"
-        );
-      }
+      sessionFactory.getSession((err, _) => {
+        assert.ifError(err);
+        done();
+      });
     });
   });
 });