diff --git a/handwritten/spanner/src/channel-factory.ts b/handwritten/spanner/src/channel-factory.ts new file mode 100644 index 00000000000..662ed3087e7 --- /dev/null +++ b/handwritten/spanner/src/channel-factory.ts @@ -0,0 +1,93 @@ +/*! + * Copyright 2026 Google LLC. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import {grpc} from 'google-gax'; +import grpcGcpModule = require('grpc-gcp'); + +const grpcGcp = grpcGcpModule(grpc); + +export interface ChannelFactoryWithWithoutAffinity { + _withoutAffinity?: unknown; +} + +/** + * Creates a delegate for GcpChannelFactory that overrides getAffinityConfig + * to return undefined. This causes grpc-gcp to skip affinity lookups and select + * a channel using its native stream load balancer (getActiveStreamsCount). + */ +export function createChannelFactoryWithoutAffinity( + channelFactory: object, +): object { + const channelFactoryWithoutAffinity = Object.create(channelFactory); + ( + channelFactoryWithoutAffinity as {getAffinityConfig?: () => undefined} + ).getAffinityConfig = () => undefined; + return channelFactoryWithoutAffinity; +} + +/** + * Custom channel factory override that pre-allocates a static delegate + * without affinity lookup on the factory instance. This avoids any object + * or closure allocations per request. + */ +export function spannerChannelFactoryOverride( + address: string, + credentials: grpc.ChannelCredentials, + options: object, +) { + const channelFactory = grpcGcp.gcpChannelFactoryOverride( + address, + credentials, + options, + ); + if (channelFactory) { + (channelFactory as ChannelFactoryWithWithoutAffinity)._withoutAffinity = + createChannelFactoryWithoutAffinity(channelFactory); + } + return channelFactory; +} + +interface SingleUseTransactionArgument { + transaction?: { + singleUse?: unknown; + single_use?: unknown; + }; +} + +/** + * Intercepts calls before dispatch. For single-use transactions (e.g. single queries), + * routes through the pre-allocated delegate to distribute across channels in the pool. + * For read/write and multi-use transactions, uses the standard channel factory so + * requests adhere to session-to-channel affinity. + */ +export function spannerCallInvocationTransformer( + callProperties: grpc.CallProperties, +): grpc.CallProperties { + if (!callProperties) { + return callProperties; + } + const argument = callProperties.argument as + SingleUseTransactionArgument | undefined; + if (argument?.transaction?.singleUse || argument?.transaction?.single_use) { + const channelFactory = + callProperties.channel as ChannelFactoryWithWithoutAffinity; + if (channelFactory?._withoutAffinity) { + callProperties.channel = + channelFactory._withoutAffinity as typeof callProperties.channel; + } + } + return grpcGcp.gcpCallInvocationTransformer(callProperties); +} diff --git a/handwritten/spanner/src/index.ts b/handwritten/spanner/src/index.ts index 17287fba98b..37dc7f9dae6 100644 --- a/handwritten/spanner/src/index.ts +++ b/handwritten/spanner/src/index.ts @@ -107,6 +107,11 @@ import {MetricsTracer} from './metrics/metrics-tracer'; // eslint-disable-next-line @typescript-eslint/no-var-requires const gcpApiConfig = require('./spanner_grpc_config.json'); +import { + spannerCallInvocationTransformer, + spannerChannelFactoryOverride, +} from './channel-factory'; + export type IOperation = instanceAdmin.longrunning.IOperation; export type GetInstancesOptions = PagedOptionsWithFilter; @@ -424,8 +429,8 @@ class Spanner extends GrpcService { // Add grpc keep alive setting 'grpc.keepalive_time_ms': 120000, // Enable grpc-gcp support - 'grpc.callInvocationTransformer': grpcGcp.gcpCallInvocationTransformer, - 'grpc.channelFactoryOverride': grpcGcp.gcpChannelFactoryOverride, + 'grpc.callInvocationTransformer': spannerCallInvocationTransformer, + 'grpc.channelFactoryOverride': spannerChannelFactoryOverride, 'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig), grpc, }, diff --git a/handwritten/spanner/test/channel-factory.ts b/handwritten/spanner/test/channel-factory.ts new file mode 100644 index 00000000000..d69b76fb12b --- /dev/null +++ b/handwritten/spanner/test/channel-factory.ts @@ -0,0 +1,241 @@ +/*! + * Copyright 2026 Google LLC. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import * as sinon from 'sinon'; +import * as proxyquire from 'proxyquire'; +import {grpc} from 'google-gax'; +import { + createChannelFactoryWithoutAffinity, + spannerCallInvocationTransformer, + spannerChannelFactoryOverride, + ChannelFactoryWithWithoutAffinity, +} from '../src/channel-factory'; + +describe('ChannelFactory and Transformer', () => { + let sandbox: sinon.SinonSandbox; + + beforeEach(() => { + sandbox = sinon.createSandbox(); + }); + + afterEach(() => { + sandbox.restore(); + }); + + describe('createChannelFactoryWithoutAffinity', () => { + it('should create a prototype delegate that returns undefined for getAffinityConfig', () => { + const originalFactory: { + options: {foo: string}; + getAffinityConfig(path?: string): {command: string} | undefined; + } = { + options: {foo: 'bar'}, + getAffinityConfig() { + return {command: 'BOUND'}; + }, + }; + + const delegate = createChannelFactoryWithoutAffinity( + originalFactory, + ) as typeof originalFactory; + + assert.strictEqual(delegate.getAffinityConfig('/test.Method'), undefined); + assert.strictEqual( + originalFactory.getAffinityConfig('/test.Method')?.command, + 'BOUND', + ); + assert.strictEqual(delegate.options, originalFactory.options); + assert.strictEqual(Object.getPrototypeOf(delegate), originalFactory); + }); + }); + + describe('spannerChannelFactoryOverride', () => { + it('should attach _withoutAffinity delegate to the created channel factory', () => { + const channelFactory = spannerChannelFactoryOverride( + 'localhost:443', + grpc.credentials.createInsecure(), + {}, + ) as ChannelFactoryWithWithoutAffinity & { + getAffinityConfig: (path: string) => unknown; + }; + + assert.ok(channelFactory); + assert.ok(channelFactory._withoutAffinity); + const withoutAffinity = channelFactory._withoutAffinity as { + getAffinityConfig: (path: string) => unknown; + }; + assert.strictEqual(typeof withoutAffinity.getAffinityConfig, 'function'); + assert.strictEqual( + withoutAffinity.getAffinityConfig('/test.Method'), + undefined, + ); + }); + + it('should return nullish channel factory when gcpChannelFactoryOverride returns nullish', () => { + const proxiedModule = proxyquire('../src/channel-factory', { + 'grpc-gcp': () => ({ + gcpChannelFactoryOverride: () => null, + gcpCallInvocationTransformer: (props: unknown) => props, + }), + }); + + const channelFactory = proxiedModule.spannerChannelFactoryOverride( + 'localhost:443', + grpc.credentials.createInsecure(), + {}, + ); + + assert.strictEqual(channelFactory, null); + }); + }); + + describe('spannerCallInvocationTransformer', () => { + function createMockCallProperties(argument?: unknown, channel?: unknown) { + return { + argument, + channel: channel ?? { + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }, + } as unknown as grpc.CallProperties; + } + + it('should route through _withoutAffinity for camelCase singleUse transaction', () => { + const mockChannelFactory = { + name: 'originalChannel', + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }; + const callProps = createMockCallProperties( + {transaction: {singleUse: {readOnly: {}}}}, + mockChannelFactory, + ); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual( + transformed.channel, + mockChannelFactory._withoutAffinity, + ); + }); + + it('should route through _withoutAffinity for snake_case single_use transaction', () => { + const mockChannelFactory = { + name: 'originalChannel', + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }; + const callProps = createMockCallProperties( + {transaction: {single_use: {read_only: {}}}}, + mockChannelFactory, + ); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual( + transformed.channel, + mockChannelFactory._withoutAffinity, + ); + }); + + it('should not change channel for multi-use transaction with id', () => { + const mockChannelFactory = { + name: 'originalChannel', + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }; + const callProps = createMockCallProperties( + {transaction: {id: Buffer.from('tx-123')}}, + mockChannelFactory, + ); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual(transformed.channel, mockChannelFactory); + }); + + it('should not change channel for read-write transaction with begin', () => { + const mockChannelFactory = { + name: 'originalChannel', + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }; + const callProps = createMockCallProperties( + {transaction: {begin: {readWrite: {}}}}, + mockChannelFactory, + ); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual(transformed.channel, mockChannelFactory); + }); + + it('should not change channel when argument is undefined', () => { + const mockChannelFactory = { + name: 'originalChannel', + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }; + const callProps = createMockCallProperties(undefined, mockChannelFactory); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual(transformed.channel, mockChannelFactory); + }); + + it('should not change channel when argument is null', () => { + const mockChannelFactory = { + name: 'originalChannel', + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }; + const callProps = createMockCallProperties(null, mockChannelFactory); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual(transformed.channel, mockChannelFactory); + }); + + it('should not change channel when argument has no transaction property', () => { + const mockChannelFactory = { + name: 'originalChannel', + _withoutAffinity: {name: 'withoutAffinityChannel'}, + }; + const callProps = createMockCallProperties( + {name: 'some-session'}, + mockChannelFactory, + ); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual(transformed.channel, mockChannelFactory); + }); + + it('should fall back gracefully when _withoutAffinity is not present on channel', () => { + const mockChannelFactory = { + name: 'originalChannel', + }; + const callProps = createMockCallProperties( + {transaction: {singleUse: {readOnly: {}}}}, + mockChannelFactory, + ); + + const transformed = spannerCallInvocationTransformer(callProps); + assert.strictEqual(transformed.channel, mockChannelFactory); + }); + + it('should return callProperties directly when callProperties is undefined or null', () => { + assert.strictEqual( + spannerCallInvocationTransformer( + undefined as unknown as grpc.CallProperties, + ), + undefined, + ); + assert.strictEqual( + spannerCallInvocationTransformer( + null as unknown as grpc.CallProperties, + ), + null, + ); + }); + }); +}); diff --git a/handwritten/spanner/test/index.ts b/handwritten/spanner/test/index.ts index 6652d25fec6..a16cfe9eb3e 100644 --- a/handwritten/spanner/test/index.ts +++ b/handwritten/spanner/test/index.ts @@ -39,6 +39,10 @@ import { } from '../src'; import {Duplex} from 'stream'; import {CLOUD_RESOURCE_HEADER, AFE_SERVER_TIMING_HEADER} from '../src/common'; +import { + spannerCallInvocationTransformer, + spannerChannelFactoryOverride, +} from '../src/channel-factory'; import {MetricsTracerFactory} from '../src/metrics/metrics-tracer-factory'; import IsolationLevel = protos.google.spanner.v1.TransactionOptions.IsolationLevel; import ReadLockMode = protos.google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode; @@ -176,7 +180,6 @@ class FakeInstanceConfig { } describe('Spanner', () => { - // tslint:disable-next-line variable-name let Spanner: typeof spnr.Spanner; let spanner: spnr.Spanner; let sandbox: sinon.SinonSandbox; @@ -185,6 +188,8 @@ describe('Spanner', () => { projectId: 'project-id', }; + let EXPECTED_OPTIONS; + before(() => { Spanner = proxyquire('../src', { './common-grpc/service': { @@ -203,6 +208,19 @@ describe('Spanner', () => { './instance-config.js': {InstanceConfig: FakeInstanceConfig}, './v1': fakeV1, }).Spanner; + + EXPECTED_OPTIONS = Object.assign({}, OPTIONS, { + libName: 'gccl', + libVersion: require('../../package.json').version, + scopes: [], + grpc, + 'grpc.keepalive_time_ms': 120000, + 'grpc.callInvocationTransformer': spannerCallInvocationTransformer, + 'grpc.channelFactoryOverride': spannerChannelFactoryOverride, + 'grpc.gcpApiConfig': { + calledWith_: apiConfig, + }, + }); }); beforeEach(async () => { @@ -224,20 +242,6 @@ describe('Spanner', () => { afterEach(() => sandbox.restore()); describe('instantiation', () => { - const EXPECTED_OPTIONS = Object.assign({}, OPTIONS, { - libName: 'gccl', - libVersion: require('../../package.json').version, - scopes: [], - grpc, - 'grpc.keepalive_time_ms': 120000, - 'grpc.callInvocationTransformer': - fakeGrpcGcp().gcpCallInvocationTransformer, - 'grpc.channelFactoryOverride': fakeGrpcGcp().gcpChannelFactoryOverride, - 'grpc.gcpApiConfig': { - calledWith_: apiConfig, - }, - }); - it('should localize a cached gapic client map', () => { assert(spanner.clients_ instanceof Map); assert.strictEqual(spanner.clients_.size, 0); diff --git a/handwritten/spanner/test/spanner.ts b/handwritten/spanner/test/spanner.ts index 0e8baaeb9db..a55e959915c 100644 --- a/handwritten/spanner/test/spanner.ts +++ b/handwritten/spanner/test/spanner.ts @@ -257,6 +257,44 @@ class XGoogRequestHeaderInterceptor { } describe('Spanner with mock server', () => { let sandbox: sinon.SinonSandbox; + + interface FakeChannelRef { + channelId: number; + getActiveStreamsCount(): number; + } + + interface FakeGcpChannelFactory { + channelRefs: FakeChannelRef[]; + maxSize: number; + addChannel(): FakeChannelRef; + getChannelRef(affinityKey?: string): FakeChannelRef; + } + + interface SpannerGaxClient { + spannerStub: Promise>; + } + + async function getGcpChannelFactory( + spannerInstance: Spanner, + ): Promise { + const gaxClient = spannerInstance.clients_.get( + 'SpannerClient', + ) as unknown as SpannerGaxClient; + assert.ok(gaxClient, 'SpannerClient should be created'); + const stub = await gaxClient.spannerStub; + const symbols = Object.getOwnPropertySymbols(stub); + const channelFactory = symbols + .map(s => stub[s]) + .find( + (v): v is FakeGcpChannelFactory => + typeof v === 'object' && + v !== null && + v.constructor?.name === 'GcpChannelFactory', + ); + assert.ok(channelFactory, 'GcpChannelFactory should exist on stub'); + return channelFactory!; + } + const selectSql = 'SELECT NUM, NAME FROM NUMBERS'; const select1 = 'SELECT 1'; const invalidSql = 'SELECT * FROM FOO'; @@ -2046,6 +2084,297 @@ describe('Spanner with mock server', () => { }); }); + it('should distribute single-use queries across multiple gRPC channels when using multiplexed session', async () => { + const database = newTestDatabase(); + // Warm up the database and client so SpannerClient is instantiated and multiplexed session is created + await database.run(selectSql); + + const channelFactory = await getGcpChannelFactory(spanner); + + // Ensure channel pool has multiple channels available (up to maxSize = 10) + while (channelFactory.channelRefs.length < channelFactory.maxSize) { + channelFactory.addChannel(); + } + assert.strictEqual(channelFactory.channelRefs.length, 10); + + // Record channel selection for queries + const selectedChannelIndices: number[] = []; + const originalGetChannelRef = + channelFactory.getChannelRef.bind(channelFactory); + const getChannelRefStub = sandbox + .stub(channelFactory, 'getChannelRef') + .callsFake(affinityKey => { + const channelReference = originalGetChannelRef(affinityKey); + selectedChannelIndices.push(channelReference.channelId); + return channelReference; + }); + + try { + // Execute multiple concurrent single-use queries + const queryCount = 5; + const promises: Promise[] = []; + for (let i = 0; i < queryCount; i++) { + promises.push(database.run(selectSql)); + } + await Promise.all(promises); + + const uniqueChannels = new Set(selectedChannelIndices); + assert.strictEqual( + uniqueChannels.size > 1, + true, + `Expected ${queryCount} queries to be distributed across multiple channels, but all were sent to channelId ${Array.from(uniqueChannels)[0]}`, + ); + } finally { + getChannelRefStub.restore(); + } + }); + + it('should scale up the gRPC channel pool when single-use queries exceed watermark', async () => { + const testSpanner = new Spanner({ + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + }); + + try { + const testInstance = testSpanner.instance('instance'); + const database = testInstance.database('scale-up-db'); + // Warm up the database and client so SpannerClient is instantiated and multiplexed session is created + await database.run(selectSql); + + const channelFactory = await getGcpChannelFactory(testSpanner); + // Initially, the channel pool has 1 channel + assert.strictEqual(channelFactory.channelRefs.length, 1); + + // Freeze the mock server so active streams accumulate on the channel + spannerMock.freeze(); + const queryPromises: Promise[] = []; + try { + // Send 26 concurrent queries to exceed maxConcurrentStreamsLowWatermark (25) + for (let i = 0; i < 26; i++) { + queryPromises.push(database.run(selectSql)); + } + + // Busy-wait for requests to be dispatched and for the channel pool to scale up + const deadline = Date.now() + 5000; + while ( + channelFactory.channelRefs.length <= 1 && + Date.now() < deadline + ) { + await new Promise(resolve => setImmediate(resolve)); + } + + // Verify that the pool scaled up from 1 to 2 channels + assert.strictEqual( + channelFactory.channelRefs.length > 1, + true, + `Expected channel pool to scale up beyond 1 channel, but got ${channelFactory.channelRefs.length}`, + ); + } finally { + spannerMock.unfreeze(); + await Promise.all(queryPromises); + } + } finally { + await testSpanner.close(); + } + }); + + it('should distribute single-use streaming queries across multiple gRPC channels when using multiplexed session', async () => { + const database = newTestDatabase(); + await database.run(selectSql); + + const channelFactory = await getGcpChannelFactory(spanner); + while (channelFactory.channelRefs.length < channelFactory.maxSize) { + channelFactory.addChannel(); + } + assert.strictEqual(channelFactory.channelRefs.length, 10); + + const selectedChannelIndices: number[] = []; + const originalGetChannelRef = + channelFactory.getChannelRef.bind(channelFactory); + const getChannelRefStub = sandbox + .stub(channelFactory, 'getChannelRef') + .callsFake(affinityKey => { + const channelReference = originalGetChannelRef(affinityKey); + selectedChannelIndices.push(channelReference.channelId); + return channelReference; + }); + + try { + const queryCount = 5; + const streamPromises: Promise[] = []; + for (let i = 0; i < queryCount; i++) { + streamPromises.push( + new Promise((resolve, reject) => { + database + .runStream(selectSql) + .on('data', () => {}) + .on('error', reject) + .on('end', resolve); + }), + ); + } + await Promise.all(streamPromises); + + const uniqueChannels = new Set(selectedChannelIndices); + assert.strictEqual( + uniqueChannels.size > 1, + true, + `Expected ${queryCount} streaming queries to be distributed across multiple channels, but all were sent to channelId ${Array.from(uniqueChannels)[0]}`, + ); + } finally { + getChannelRefStub.restore(); + } + }); + + it('should keep multi-use snapshot queries pinned to the same channel when using multiplexed session', async () => { + const database = newTestDatabase(); + await database.run(selectSql); + + const channelFactory = await getGcpChannelFactory(spanner); + + const selectedChannelIndices: number[] = []; + const originalGetChannelRef = + channelFactory.getChannelRef.bind(channelFactory); + const getChannelRefStub = sandbox + .stub(channelFactory, 'getChannelRef') + .callsFake(affinityKey => { + const channelReference = originalGetChannelRef(affinityKey); + selectedChannelIndices.push(channelReference.channelId); + return channelReference; + }); + + try { + const [snapshot] = await database.getSnapshot(); + await snapshot.run(selectSql); + await snapshot.run(selectSql); + snapshot.end(); + + const uniqueChannels = new Set(selectedChannelIndices); + assert.strictEqual( + uniqueChannels.size, + 1, + 'Expected all queries in multi-use snapshot to stick to the same channel', + ); + } finally { + getChannelRefStub.restore(); + } + }); + + it('should keep multi-statement read-write transaction pinned to the same channel when using multiplexed session', async () => { + const database = newTestDatabase(); + await database.run(selectSql); + + const channelFactory = await getGcpChannelFactory(spanner); + + const selectedChannelIndices: number[] = []; + const originalGetChannelRef = + channelFactory.getChannelRef.bind(channelFactory); + const getChannelRefStub = sandbox + .stub(channelFactory, 'getChannelRef') + .callsFake(affinityKey => { + const channelReference = originalGetChannelRef(affinityKey); + selectedChannelIndices.push(channelReference.channelId); + return channelReference; + }); + + try { + await database.runTransactionAsync(async transaction => { + await transaction.run(selectSql); + await transaction.run(selectSql); + await transaction.commit(); + }); + + const uniqueChannels = new Set(selectedChannelIndices); + assert.strictEqual( + uniqueChannels.size, + 1, + 'Expected all statements in read-write transaction to stick to the same channel', + ); + } finally { + getChannelRefStub.restore(); + } + }); + + it('should keep table.read pinned to the same channel as snapshot.begin', async () => { + const fields = [ + protobuf.StructType.Field.create({ + name: 'C1', + type: protobuf.Type.create({code: protobuf.TypeCode.STRING}), + }), + ]; + const metadata = new protobuf.ResultSetMetadata({ + rowType: new protobuf.StructType({ + fields, + }), + }); + const results: PartialResultSet[] = [ + PartialResultSet.create({ + metadata, + values: [{stringValue: 'V1'}], + }), + ]; + const request = { + table: 'TestTable', + columns: ['C1'], + keySet: { + keys: [], + all: true, + ranges: [], + }, + }; + spannerMock.putReadRequestResult( + request, + mock.ReadRequestResult.resultSet(results), + ); + + const database = newTestDatabase(); + await database.run(selectSql); + + const channelFactory = await getGcpChannelFactory(spanner); + const selectedChannelIndices: number[] = []; + const originalGetChannelRef = + channelFactory.getChannelRef.bind(channelFactory); + const getChannelRefStub = sandbox + .stub(channelFactory, 'getChannelRef') + .callsFake(affinityKey => { + const channelReference = originalGetChannelRef(affinityKey); + selectedChannelIndices.push(channelReference.channelId); + return channelReference; + }); + + try { + const table = database.table('TestTable'); + const [rows] = await table.read({columns: ['C1']}); + assert.strictEqual(rows.length, 1); + + const uniqueChannels = new Set(selectedChannelIndices); + assert.strictEqual( + uniqueChannels.size, + 1, + 'Expected table.read (which uses database.getSnapshot) to stick to the same channel', + ); + } finally { + getChannelRefStub.restore(); + } + }); + + it('should propagate errors on single-use queries cleanly through transformed channel pipeline', async () => { + const database = newTestDatabase(); + const invalidSql = 'SELECT * FROM NonExistentTable'; + const error = new Error('Table not found: NonExistentTable'); + (error as grpc.ServiceError).code = grpc.status.NOT_FOUND; + spannerMock.putStatementResult( + invalidSql, + mock.StatementResult.error(error as grpc.ServiceError), + ); + + await assert.rejects( + database.run(invalidSql), + /Table not found: NonExistentTable/, + ); + }); + it('should execute the transaction(database.getSnapshot) successfully using multiplexed session', done => { const database = newTestDatabase(); const pool = (database.sessionFactory_ as SessionFactory) @@ -2183,6 +2512,46 @@ describe('Spanner with mock server', () => { }); }); + it('should distribute single-use queries across multiple gRPC channels when using regular sessions', async () => { + const database = newTestDatabase(); + await database.run(selectSql); + + const channelFactory = await getGcpChannelFactory(spanner); + while (channelFactory.channelRefs.length < channelFactory.maxSize) { + channelFactory.addChannel(); + } + assert.strictEqual(channelFactory.channelRefs.length, 10); + + const selectedChannelIndices: number[] = []; + const originalGetChannelRef = + channelFactory.getChannelRef.bind(channelFactory); + const getChannelRefStub = sandbox + .stub(channelFactory, 'getChannelRef') + .callsFake(affinityKey => { + const channelReference = originalGetChannelRef(affinityKey); + selectedChannelIndices.push(channelReference.channelId); + return channelReference; + }); + + try { + const queryCount = 5; + const promises: Promise[] = []; + for (let i = 0; i < queryCount; i++) { + promises.push(database.run(selectSql)); + } + await Promise.all(promises); + + const uniqueChannels = new Set(selectedChannelIndices); + assert.strictEqual( + uniqueChannels.size > 1, + true, + `Expected ${queryCount} queries on regular sessions to distribute across channels, but got ${Array.from(uniqueChannels)}`, + ); + } finally { + getChannelRefStub.restore(); + } + }); + it('should execute the transaction(database.getSnapshot) successfully using regular session', done => { const database = newTestDatabase({min: 1, max: 1}); const pool = (database.sessionFactory_ as SessionFactory)