fix(key-value-stores): store the value piped into set-value

src/commands/key-value-stores/set-value.ts

@@ -1,5 +1,6 @@
 import type { ApifyApiError } from 'apify-client';
 
+import { cachedStdinInput } from '../../entrypoints/_shared.js';
 import { ApifyCommand } from '../../lib/command-framework/apify-command.js';
 import { Args } from '../../lib/command-framework/args.js';
 import { Flags } from '../../lib/command-framework/flags.js';
@@ -52,6 +53,17 @@ export class KeyValueStoresSetValueCommand extends ApifyCommand<typeof KeyValueS
 		const { storeId, itemKey, value } = this.args;
 		const { contentType } = this.flags;
 
+		// NOTE: If the value is passed as an argument and on stdin at the same time, the argument wins.
+		// Stdin was already read at startup (see `cachedStdinInput`), so the raw `process.stdin` stream is drained.
+		const recordValue = value || cachedStdinInput;
+
+		if (!recordValue) {
+			error({ message: 'No value was provided. Pass it as an argument or pipe it into the command.' });
+			process.exitCode = 1;
+
+			return;
+		}
+
 		const apifyClient = await getLoggedClientOrThrow();
 		const maybeStore = await tryToGetKeyValueStore(apifyClient, storeId);
 
@@ -67,7 +79,7 @@ export class KeyValueStoresSetValueCommand extends ApifyCommand<typeof KeyValueS
 
 		try {
 			// TODO: again, the types need to be fixed -w-
-			await client.setRecord({ key: itemKey, value: (value || process.stdin) as string, contentType });
+			await client.setRecord({ key: itemKey, value: recordValue as string, contentType });
 
 			success({
 				message: `Value with key "${itemKey}" set in the key-value store.`,

test/e2e/__helpers__/run-cli.ts

@@ -9,7 +9,7 @@ const DistActor = fileURLToPath(new URL('./dist/actor.js', ProjectRoot));
 
 export interface RunCliOptions {
 	cwd?: string | URL;
-	stdin?: string;
+	stdin?: string | Uint8Array;
 	env?: Record<string, string>;
 	timeout?: number;
 }

test/e2e/commands/key-value-stores/set-value-stdin.test.ts

@@ -0,0 +1,116 @@
+import { createHash } from 'node:crypto';
+import { once } from 'node:events';
+import type { Server } from 'node:http';
+import { createServer } from 'node:http';
+import type { AddressInfo } from 'node:net';
+import { brotliDecompressSync, gunzipSync, inflateSync } from 'node:zlib';
+
+import { runCli } from '../../__helpers__/run-cli.js';
+
+// A value that is neither valid UTF-8 nor small: every byte value, repeated to ~100 KB (the command line cannot carry it).
+const binaryValue = Buffer.from(Array.from({ length: 100_000 }, (_, i) => i % 256));
+
+const sha256 = (data: Uint8Array) => createHash('sha256').update(data).digest('hex');
+
+// The API client compresses larger request bodies; the mock stores what the server would store.
+const decodeBody = (raw: Buffer, encoding: string | undefined) => {
+	switch (encoding) {
+		case 'gzip':
+			return gunzipSync(raw);
+		case 'br':
+			return brotliDecompressSync(raw);
+		case 'deflate':
+			return inflateSync(raw);
+		default:
+			return raw;
+	}
+};
+
+describe('[e2e] key-value-stores set-value reads the value from stdin (#1474)', () => {
+	// Fake API: answers the store lookup and records what is PUT.
+	describe('apify key-value-stores set-value', () => {
+		let server: Server;
+		let apiUrl: string;
+		let puts: { url: string; contentType: string | undefined; body: Buffer }[];
+
+		beforeAll(async () => {
+			server = createServer((req, res) => {
+				const chunks: Buffer[] = [];
+				req.on('data', (chunk: Buffer) => chunks.push(chunk));
+				req.on('end', () => {
+					const raw = Buffer.concat(chunks);
+					const reply = (data: unknown, status = 200) => {
+						res.writeHead(status, { 'content-type': 'application/json' });
+						res.end(JSON.stringify({ data }));
+					};
+
+					if (req.method === 'PUT') {
+						puts.push({
+							url: req.url!,
+							contentType: req.headers['content-type'],
+							body: decodeBody(raw, req.headers['content-encoding']),
+						});
+						return reply({}, 201);
+					}
+
+					if (req.url === '/v2/users/me') return reply({ id: 'user1', username: 'someone' });
+
+					// The store lookup by ID or name.
+					return reply({ id: 'store1', name: 'store1' });
+				});
+			}).listen(0, '127.0.0.1');
+			await once(server, 'listening');
+			apiUrl = `http://127.0.0.1:${(server.address() as AddressInfo).port}`;
+		});
+
+		afterAll(async () => {
+			server.close();
+			await once(server, 'close');
+		});
+
+		beforeEach(() => {
+			puts = [];
+		});
+
+		const runSetValue = (args: string[], stdin?: string | Uint8Array) =>
+			runCli('apify', ['key-value-stores', 'set-value', 'store1', ...args], {
+				stdin,
+				env: { APIFY_TOKEN: 'dummy-token', APIFY_CLIENT_BASE_URL: apiUrl },
+			});
+
+		it('stores the JSON piped on stdin', async () => {
+			const result = await runSetValue(['stdin-test', '--content-type', 'application/json'], '{"a":1}');
+
+			expect(result.exitCode, `stderr: ${result.stderr}`).toBe(0);
+			expect(puts).toHaveLength(1);
+			expect(puts[0].url).toContain('/records/stdin-test');
+			expect(puts[0].contentType).toContain('application/json');
+			expect(puts[0].body.toString('utf8')).toBe('{"a":1}');
+		});
+
+		it('stores a ~100 KB binary value piped on stdin byte for byte', async () => {
+			const result = await runSetValue(['big', '--content-type', 'application/octet-stream'], binaryValue);
+
+			expect(result.exitCode, `stderr: ${result.stderr}`).toBe(0);
+			expect(puts).toHaveLength(1);
+			expect(puts[0].body.length).toBe(binaryValue.length);
+			expect(sha256(puts[0].body)).toBe(sha256(binaryValue));
+		});
+
+		it('prefers the value argument over stdin', async () => {
+			const result = await runSetValue(['key', '{"from":"argument"}'], 'from stdin');
+
+			expect(result.exitCode, `stderr: ${result.stderr}`).toBe(0);
+			expect(puts).toHaveLength(1);
+			expect(puts[0].body.toString('utf8')).toBe('{"from":"argument"}');
+		});
+
+		it('fails without writing a record when there is no value', async () => {
+			const result = await runSetValue(['key']);
+
+			expect(result.exitCode).toBe(1);
+			expect(result.stderr).toContain('No value was provided');
+			expect(puts).toHaveLength(0);
+		});
+	});
+});

test/local/commands/key-value-stores/set-value.test.ts

@@ -0,0 +1,98 @@
+import type { ApifyClient } from 'apify-client';
+
+import { KeyValueStoresSetValueCommand } from '../../../../src/commands/key-value-stores/set-value.js';
+import { testRunCommand } from '../../../../src/lib/command-framework/apify-command.js';
+import { useAuthSetup } from '../../../__setup__/hooks/useAuthSetup.js';
+import { useConsoleSpy } from '../../../__setup__/hooks/useConsoleSpy.js';
+
+const { mockGetLoggedClientOrThrow, stdin } = vitest.hoisted(() => ({
+	mockGetLoggedClientOrThrow: vitest.fn(),
+	// What the CLI read from stdin at startup (`cachedStdinInput`); the raw `process.stdin` stream is drained by then.
+	stdin: { value: undefined as Buffer | undefined },
+}));
+
+vitest.mock('../../../../src/entrypoints/_shared.js', () => ({
+	get cachedStdinInput() {
+		return stdin.value;
+	},
+}));
+
+vitest.mock('../../../../src/lib/utils.js', async (importOriginal) => {
+	const original = await importOriginal<typeof import('../../../../src/lib/utils.js')>();
+	return {
+		...original,
+		getLoggedClientOrThrow: mockGetLoggedClientOrThrow,
+	};
+});
+
+useAuthSetup();
+
+const { logMessages } = useConsoleSpy();
+
+const setRecord = vitest.fn();
+
+const fakeClient = () =>
+	({
+		keyValueStore: (id: string) => ({
+			get: async () => ({ id }),
+			setRecord,
+		}),
+	}) as unknown as ApifyClient;
+
+beforeEach(() => {
+	mockGetLoggedClientOrThrow.mockResolvedValue(fakeClient());
+	setRecord.mockResolvedValue(undefined);
+});
+
+afterEach(() => {
+	mockGetLoggedClientOrThrow.mockReset();
+	setRecord.mockReset();
+	stdin.value = undefined;
+	process.exitCode = undefined;
+});
+
+describe('apify key-value-stores set-value', () => {
+	it('stores the data piped on stdin byte for byte when no value argument is given', async () => {
+		// Not valid UTF-8, so any text round trip would corrupt it.
+		stdin.value = Buffer.from([0x00, 0xff, 0xfe, 0x80, 0x7b, 0x22, 0x61, 0x22, 0x3a, 0x31, 0x7d, 0x0a]);
+
+		await testRunCommand(KeyValueStoresSetValueCommand, {
+			args_storeId: 'store1',
+			args_itemKey: 'stdin-test',
+			flags_contentType: 'application/octet-stream',
+		});
+
+		expect(setRecord).toHaveBeenCalledTimes(1);
+		expect(setRecord).toHaveBeenCalledWith({
+			key: 'stdin-test',
+			value: stdin.value,
+			contentType: 'application/octet-stream',
+		});
+		expect(process.exitCode).toBeFalsy();
+	});
+
+	it('prefers the value argument over stdin', async () => {
+		stdin.value = Buffer.from('from stdin');
+
+		await testRunCommand(KeyValueStoresSetValueCommand, {
+			args_storeId: 'store1',
+			args_itemKey: 'key',
+			args_value: '{"a":1}',
+			flags_contentType: 'application/json',
+		});
+
+		expect(setRecord).toHaveBeenCalledWith({ key: 'key', value: '{"a":1}', contentType: 'application/json' });
+	});
+
+	it('fails without calling the API when there is neither a value argument nor stdin data', async () => {
+		await testRunCommand(KeyValueStoresSetValueCommand, {
+			args_storeId: 'store1',
+			args_itemKey: 'key',
+			flags_contentType: 'application/json',
+		});
+
+		expect(setRecord).not.toHaveBeenCalled();
+		expect(process.exitCode).toBe(1);
+		expect(logMessages.error.join('\n')).toContain('No value was provided');
+	});
+});