feat: allow binding a Session instance to a Request (#3518)

Adds a `sessionId?: string` field to `Request`. If this is set,
`BasicCrawler` will fetch this given session when processing the request
(and will fail if the session is not present in the `SessionPool`).

This PR also adds a `sessionId` parameter to the `enqueueLinks` (and
other?) helper functions. If the request-adding function is called with
this parameter, all enqueued requests will be bound to this session.

Closes #3446

Author: barjin

packages/basic-crawler/src/internals/basic-crawler.ts

@@ -43,6 +43,7 @@ import {
     KeyValueStore,
     LogLevel,
     mergeCookies,
+    MissingSessionError,
     NavigationSkippedError,
     NonRetryableError,
     purgeDefaultStorages,
@@ -1104,6 +1105,16 @@ export class BasicCrawler<
     private async resolveSession({ request }: { request: Request }) {
         const session = await this._timeoutAndRetry(
             async () => {
+                if (request.sessionId) {
+                    const existingSession = await this.sessionPool!.getSession(request.sessionId);
+
+                    if (!existingSession) {
+                        throw new ContextPipelineInitializationError(new MissingSessionError(request.sessionId));
+                    }
+
+                    return existingSession;
+                }
+
                 return await this.sessionPool!.newSession({
                     proxyInfo: await this.proxyConfiguration?.newProxyInfo({
                         request: request ?? undefined,

packages/basic-crawler/test/batch-add-requests.test.ts

@@ -61,3 +61,50 @@ describe('BasicCrawler#addRequests with big batch sizes', () => {
         expect(result.addedRequests).toHaveLength(2000);
     });
 });
+
+describe('BasicCrawler - request.sessionId', () => {
+    const localStorageEmulator = new MemoryStorageEmulator();
+
+    beforeEach(async () => {
+        await localStorageEmulator.init();
+    });
+
+    afterAll(async () => {
+        await localStorageEmulator.destroy();
+    });
+
+    test('uses the session matching request.sessionId from the session pool', async () => {
+        let resolvedSessionId: string | undefined;
+
+        const crawler = new BasicCrawler({
+            requestHandler({ session }) {
+                resolvedSessionId = session.id;
+            },
+        });
+
+        const session = await crawler.sessionPool.newSession({ id: 'my-session' });
+
+        await crawler.run([{ url: 'http://localhost', sessionId: session.id }]);
+
+        expect(resolvedSessionId).toBe('my-session');
+    });
+
+    test('throws when request.sessionId is not found in the session pool', async () => {
+        const errors: Error[] = [];
+
+        const crawler = new BasicCrawler({
+            maxRequestRetries: 0,
+            requestHandler() {},
+            failedRequestHandler(_ctx, error) {
+                errors.push(error);
+            },
+        });
+
+        await crawler.run([{ url: 'http://localhost', sessionId: 'nonexistent' }]);
+
+        expect(errors).toHaveLength(1);
+        expect(errors[0].message).toContain(
+            "The current SessionPool instance doesn't contain the requested sessionId: nonexistent",
+        );
+    });
+});

packages/core/src/enqueue_links/enqueue_links.ts

@@ -55,6 +55,9 @@ export interface EnqueueLinksOptions extends RequestQueueOperationOptions {
      */
     label?: string;
 
+    /** Sets {@apilink Request.sessionId} for newly enqueued requests. */
+    sessionId?: string;
+
     /**
      * If set to `true`, tells the crawler to skip navigation and process the request directly.
      * @default false
@@ -304,6 +307,7 @@ export async function enqueueLinks(
             onSkippedRequest: ow.optional.function,
             forefront: ow.optional.boolean,
             skipNavigation: ow.optional.boolean,
+            sessionId: ow.optional.string,
             limit: ow.optional.number,
             selector: ow.optional.string,
             baseUrl: ow.optional.string,

packages/core/src/enqueue_links/shared.ts

@@ -217,7 +217,10 @@ export function filterRequestOptionsByPatterns(
  */
 export function createRequestOptions(
     sources: readonly (string | Record<string, unknown>)[],
-    options: Pick<EnqueueLinksOptions, 'label' | 'userData' | 'baseUrl' | 'skipNavigation' | 'strategy'> = {},
+    options: Pick<
+        EnqueueLinksOptions,
+        'label' | 'userData' | 'baseUrl' | 'skipNavigation' | 'sessionId' | 'strategy'
+    > = {},
 ): RequestOptions[] {
     return sources
         .map((src) =>
@@ -247,6 +250,10 @@ export function createRequestOptions(
                 requestOptions.skipNavigation = true;
             }
 
+            if (options.sessionId) {
+                requestOptions.sessionId = options.sessionId;
+            }
+
             return requestOptions;
         });
 }

packages/core/src/errors.ts

@@ -36,6 +36,15 @@ export class SessionError extends RetryRequestError {
     }
 }
 
+/**
+ * Thrown when a requested session is not found in the referenced SessionPool.
+ */
+export class MissingSessionError extends Error {
+    constructor(sessionId: string) {
+        super(`The current SessionPool instance doesn't contain the requested sessionId: ${sessionId}`);
+    }
+}
+
 export class ContextPipelineInterruptedError extends Error {
     constructor(message?: string) {
         super(`Request handling was interrupted during context initialization ${message ? ` - ${message}` : ''}`);

packages/core/src/request.ts

@@ -24,6 +24,7 @@ const requestOptionalPredicates = {
     noRetry: ow.optional.boolean,
     retryCount: ow.optional.number,
     sessionRotationCount: ow.optional.number,
+    sessionId: ow.optional.string,
     maxRetries: ow.optional.number,
     errorMessages: ow.optional.array.ofType(ow.string),
     headers: ow.optional.object,
@@ -170,6 +171,7 @@ class CrawleeRequest<UserData extends Dictionary = Dictionary> {
             noRetry = false,
             retryCount = 0,
             sessionRotationCount = 0,
+            sessionId,
             maxRetries,
             errorMessages = [],
             headers = {},
@@ -185,6 +187,7 @@ class CrawleeRequest<UserData extends Dictionary = Dictionary> {
             loadedUrl?: string;
             retryCount?: number;
             sessionRotationCount?: number;
+            sessionId?: string;
             errorMessages?: string[];
             handledAt?: string | Date;
         };
@@ -256,6 +259,7 @@ class CrawleeRequest<UserData extends Dictionary = Dictionary> {
         if (skipNavigation != null) this.skipNavigation = skipNavigation;
         if (maxRetries != null) this.maxRetries = maxRetries;
         if (crawlDepth != null) this.userData.__crawlee.crawlDepth ??= crawlDepth;
+        if (sessionId) this.sessionId = sessionId;
 
         // If it's already set, don't override it (for instance when fetching from storage)
         if (enqueueStrategy) {
@@ -332,6 +336,16 @@ class CrawleeRequest<UserData extends Dictionary = Dictionary> {
         }
     }
 
+    /** ID of a session to use for this request. When set, the crawler will fetch this session from the session pool instead of creating a new one. */
+    get sessionId(): string | undefined {
+        return this.userData.__crawlee?.sessionId;
+    }
+
+    set sessionId(value: string | undefined) {
+        (this.userData as Dictionary).__crawlee ??= {};
+        this.userData.__crawlee.sessionId = value;
+    }
+
     /** shortcut for getting `request.userData.label` */
     get label(): string | undefined {
         return this.userData.label;
@@ -557,6 +571,12 @@ export interface RequestOptions<UserData extends Dictionary = Dictionary> {
      */
     noRetry?: boolean;
 
+    /**
+     * ID of a session from the crawler's `SessionPool` to use for this request.
+     * When set, the crawler will fetch this session from the pool instead of creating a new one.
+     */
+    sessionId?: string;
+
     /**
      * If set to `true` then the crawler processing this request evaluates
      * the `requestHandler` immediately without prior browser navigation.

packages/core/test/enqueue_links/userData.test.ts

@@ -105,4 +105,36 @@ describe("enqueueLinks() - userData shouldn't be changed and outer label must ta
         expect(enqueued[0].url).toBe('https://example.com/first');
         expect(enqueued[0].label).toBe('first');
     });
+
+    test('sets sessionId on all enqueued requests', async () => {
+        const { enqueued, requestQueue } = createRequestQueueMock();
+
+        await cheerioCrawlerEnqueueLinks({
+            options: {
+                sessionId: 'my-session',
+            },
+            $,
+            requestQueue,
+            originalRequestUrl: 'https://example.com',
+        });
+
+        expect(enqueued).toHaveLength(2);
+        expect(enqueued[0].userData?.__crawlee?.sessionId).toBe('my-session');
+        expect(enqueued[1].userData?.__crawlee?.sessionId).toBe('my-session');
+    });
+
+    test('does not set sessionId when option is not provided', async () => {
+        const { enqueued, requestQueue } = createRequestQueueMock();
+
+        await cheerioCrawlerEnqueueLinks({
+            options: {},
+            $,
+            requestQueue,
+            originalRequestUrl: 'https://example.com',
+        });
+
+        expect(enqueued).toHaveLength(2);
+        expect(enqueued[0].userData?.__crawlee?.sessionId).toBeUndefined();
+        expect(enqueued[1].userData?.__crawlee?.sessionId).toBeUndefined();
+    });
 });