feat(spanner): support for Multiplexed Session Partitioned Ops (#2252)

* feat(spanner): support for Multiplexed Session Partitioned Ops * feat(spanner): add support for snapshot isolation (#2245) This PR contains code changes to add support for option IsolationLevel at the client level and at the transaction level. supported methods are(RW and Blind Write): ``` - writeAtLeastOnce - runTransactionAsync - runTransaction - getTransaction - async getTransaction(from transaction runner class) ``` * refactor * refactor

Author: alkatrivedi

.kokoro/presubmit/node14/system-test-multiplexed-session.cfg

@@ -1964,7 +1964,7 @@ describe('Database', () => {
  },
  };
 
- let fakePool: FakeSessionPool;
+ let fakeSessionFactory: FakeSessionFactory;
  let fakeSession: FakeSession;
  let fakePartitionedDml = new FakeTransaction(
  {} as google.spanner.v1.TransactionOptions.PartitionedDml
@@ -1974,14 +1974,17 @@ describe('Database', () => {
  let beginStub;
 
  beforeEach(() => {
- fakePool = database.pool_;
+ fakeSessionFactory = database.sessionFactory_;
  fakeSession = new FakeSession();
  fakePartitionedDml = new FakeTransaction(
  {} as google.spanner.v1.TransactionOptions.PartitionedDml
  );
 
  getSessionStub = (
- sandbox.stub(fakePool, 'getSession') as sinon.SinonStub
+ sandbox.stub(
+ fakeSessionFactory,
+ 'getSessionForPartitionedOps'
+ ) as sinon.SinonStub
  ).callsFake(callback => {
  callback(null, fakeSession);
  });
@@ -2077,7 +2080,7 @@ describe('Database', () => {
  beginStub.callsFake(callback => callback(fakeError));
 
  const releaseStub = (
- sandbox.stub(fakePool, 'release') as sinon.SinonStub
+ sandbox.stub(fakeSessionFactory, 'release') as sinon.SinonStub
  ).withArgs(fakeSession);
 
  database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
@@ -2124,7 +2127,7 @@ describe('Database', () => {
 
  it('session released on transaction end', done => {
  const releaseStub = (
- sandbox.stub(fakePool, 'release') as sinon.SinonStub
+ sandbox.stub(fakeSessionFactory, 'release') as sinon.SinonStub
  ).withArgs(fakeSession);
 
  database.runPartitionedUpdate(QUERY, async () => {

.kokoro/trampoline_v2.sh

@@ -2945,7 +2945,7 @@ class Database extends common.GrpcServiceObject {
  ...this._traceConfig,
  };
  return startTrace('Database.runPartitionedUpdate', traceConfig, span => {
- this.pool_.getSession((err, session) => {
+ this.sessionFactory_.getSessionForPartitionedOps((err, session) => {
  if (err) {
  setSpanError(span, err);
  span.end();
@@ -2976,21 +2976,21 @@ class Database extends common.GrpcServiceObject {
  }
  transaction.begin(err => {
  if (err) {
- this.pool_.release(session!);
+ this.sessionFactory_.release(session!);
  callback!(err, 0);
  return;
  }
 
  transaction.runUpdate(query, (err, updateCount) => {
  if (err) {
  if (err.code !== grpc.status.ABORTED) {
- this.pool_.release(session!);
+ this.sessionFactory_.release(session!);
  callback!(err, 0);
  return;
  }
  this._runPartitionedUpdate(session, query, callback);
  } else {
- this.pool_.release(session!);
+ this.sessionFactory_.release(session!);
  callback!(null, updateCount);
  return;
  }

observability-test/database.ts

@@ -56,6 +56,14 @@ export interface SessionFactoryInterface {
  */
  getSession(callback: GetSessionCallback): void;
 
+ /**
+ * When called returns a session for paritioned dml.
+ *
+ * @name SessionFactoryInterface#getSessionForPartitionedOps
+ * @param {GetSessionCallback} callback The callback function.
+ */
+ getSessionForPartitionedOps(callback: GetSessionCallback): void;
+
  /**
  * When called returns the pool object.
  *
@@ -97,6 +105,7 @@ export class SessionFactory
  multiplexedSession_: MultiplexedSessionInterface;
  pool_: SessionPoolInterface;
  isMultiplexed: boolean;
+ isMultiplexedPartitionOps: boolean;
  constructor(
  database: Database,
  name: String,
@@ -117,6 +126,11 @@ export class SessionFactory
  process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true'
  ? (this.isMultiplexed = true)
  : (this.isMultiplexed = false);
+ // set the isMultiplexedPartitionedOps property to true if multiplexed session is enabled for paritioned ops, otherwise set the property to false
+ this.isMultiplexedPartitionOps =
+ process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true' &&
+ process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS ===
+ 'true';
  // Multiplexed sessions should only be created if its enabled.
  if (this.isMultiplexed) {
  this.multiplexedSession_.on('error', this.emit.bind(database, 'error'));
@@ -143,6 +157,23 @@ export class SessionFactory
  );
  }
 
+ /**
+ * Retrieves a session for partitioned operations, selecting the appropriate session type
+ * based on whether multiplexed sessions are enabled.
+ *
+ * If multiplexed sessions are enabled for partitioned ops this methods delegates the request to `getSession()`, which returns
+ * either a multiplexed session or a regular session based on the configuration.
+ *
+ * If the multiplexed sessions are disabled, a session is retrieved from the regular session pool.
+ *
+ * @param {GetSessionCallback} callback The callback function.
+ */
+ getSessionForPartitionedOps(callback: GetSessionCallback): void {
+ this.isMultiplexedPartitionOps
+ ? this.getSession(callback)
+ : this.pool_.getSession(callback);
+ }
+
  /**
  * Returns the regular session pool object.
  *
@@ -154,9 +185,11 @@ export class SessionFactory
  }
 
  /**
- * Releases a session back to the session pool.
+ * Releases a regular session back to the session pool.
+ *
+ * This methods does not release a multiplexed session.
  *
- * This method returns a session to the pool after it is no longer needed.
+ * It returns a session to the pool after it is no longer needed.
  * It is a no-op for multiplexed sessions.
  *
  * @param {Session} session - The session to be released. This should be an instance of `Session` that was
@@ -165,7 +198,7 @@ export class SessionFactory
  * @throws {Error} If the session is invalid or cannot be released.
  */
  release(session: Session): void {
- if (!this.isMultiplexed) {
+ if (!session.metadata?.multiplexed) {
  this.pool_.release(session);
  }
  }