Skip to content

Commit

Permalink
fix: Add explicit parameters limit in PostgreSQL and Redshift drivers
Browse files Browse the repository at this point in the history
  • Loading branch information
mcheshkov committed Oct 4, 2024
1 parent 7ac7a5b commit 35ca5a8
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 1 deletion.
17 changes: 17 additions & 0 deletions packages/cubejs-postgres-driver/src/PostgresDriver.ts
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,8 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
values: unknown[],
{ highWaterMark }: StreamOptions
): Promise<StreamTableDataWithTypes> {
PostgresDriver.checkValuesLimit(values);

const conn = await this.pool.connect();

try {
Expand Down Expand Up @@ -324,7 +326,22 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
}
}

protected static checkValuesLimit(values?: unknown[]) {
// PostgreSQL protocol allows sending up to 65535 params in a single bind message
// See https://github.com/postgres/postgres/blob/REL_16_0/src/backend/tcop/postgres.c#L1698-L1708
// See https://github.com/postgres/postgres/blob/REL_16_0/src/backend/libpq/pqformat.c#L428-L431
// But 'pg' module does not check for params count, and ends up sending incorrect bind message
// See https://github.com/brianc/node-postgres/blob/92cb640fd316972e323ced6256b2acd89b1b58e0/packages/pg-protocol/src/serializer.ts#L155
// See https://github.com/brianc/node-postgres/blob/92cb640fd316972e323ced6256b2acd89b1b58e0/packages/pg-protocol/src/buffer-writer.ts#L32-L37
const length = (values?.length ?? 0);
if (length >= 65536) {
throw new Error(`PostgreSQL protocol does not support more than 65535 parameters, but ${length} passed`);
}
}

protected async queryResponse(query: string, values: unknown[]) {
PostgresDriver.checkValuesLimit(values);

const conn = await this.pool.connect();

try {
Expand Down
26 changes: 26 additions & 0 deletions packages/cubejs-postgres-driver/test/PostgresDriver.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,10 @@ import { PostgresDriver } from '../src';

const streamToArray = require('stream-to-array');

function largeParams(): Array<string> {
return new Array(65536).fill('foo');
}

describe('PostgresDriver', () => {
let container: StartedTestContainer;
let driver: PostgresDriver;
Expand Down Expand Up @@ -56,6 +60,14 @@ describe('PostgresDriver', () => {
]);
});

test('too many params', async () => {
await expect(
driver.query(`SELECT 'foo'::TEXT;`, largeParams())
)
.rejects
.toThrow('PostgreSQL protocol does not support more than 65535 parameters, but 65536 passed');
});

test('stream', async () => {
await driver.uploadTable(
'test.streaming_test',
Expand Down Expand Up @@ -116,6 +128,20 @@ describe('PostgresDriver', () => {
}
});

test('stream (too many params)', async () => {
try {
await driver.stream('select * from test.streaming_test', largeParams(), {
highWaterMark: 1000,
});

throw new Error('stream must throw an exception');
} catch (e: any) {
expect(e.message).toEqual(
'PostgreSQL protocol does not support more than 65535 parameters, but 65536 passed'
);
}
});

// Note: This test MUST be the last in the list.
test('release', async () => {
expect(async () => {
Expand Down
34 changes: 33 additions & 1 deletion packages/cubejs-redshift-driver/src/RedshiftDriver.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,13 @@

import { getEnv } from '@cubejs-backend/shared';
import { PostgresDriver, PostgresDriverConfiguration } from '@cubejs-backend/postgres-driver';
import { DownloadTableCSVData, DriverCapabilities, UnloadOptions } from '@cubejs-backend/base-driver';
import {
DownloadTableCSVData,
DriverCapabilities,
StreamOptions,
StreamTableDataWithTypes,
UnloadOptions
} from '@cubejs-backend/base-driver';
import crypto from 'crypto';

interface RedshiftDriverExportRequiredAWS {
Expand Down Expand Up @@ -91,6 +97,32 @@ export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration>
};
}

protected static checkValuesLimit(values?: unknown[]) {
// Redshift server is not exactly compatible with PostgreSQL protocol
// And breaks after 32767 parameter values with `there is no parameter $-32768`
// This is a bug/misbehaviour on server side, nothing we can do besides generate a more meaningful error
const length = (values?.length ?? 0);
if (length >= 32768) {
throw new Error(`Redshift server does not support more than 32767 parameters, but ${length} passed`);
}
}

public override async stream(
query: string,
values: unknown[],
options: StreamOptions
): Promise<StreamTableDataWithTypes> {
RedshiftDriver.checkValuesLimit(values);

return super.stream(query, values, options);
}

protected override async queryResponse(query: string, values: unknown[]) {
RedshiftDriver.checkValuesLimit(values);

return super.queryResponse(query, values);
}

protected getExportBucket(
dataSource: string,
): RedshiftDriverExportAWS | undefined {
Expand Down

0 comments on commit 35ca5a8

Please sign in to comment.