feat(utils): add discoverValidSitemaps utility (#3339)

Related to https://github.com/apify/apify-sdk-js/issues/486. I'm
[developing generic sitemap
scraper](apify/actor-scraper#205) and it's going
to share a big utility function (main chunk of logic) with wcc -
`discoverValidSitemaps`. I've asked @barjin if I could factor it out and
he told this util could fit into crawlee. It's mainly copied from wcc,
but to keep the dependencies unchanged, it's using got-scraping to check
for url existence instead of impit (I think it doesn't matter for
sitemaps), and `urlExists` is inlined (until we don't add http client to
these utils in v4 as @barjin told me). It's also turned into an async
generator. Let me know if you see a better place for this util.

Author: foxt451

packages/utils/src/internals/iterables.ts

@@ -216,3 +216,56 @@ export function peekableAsyncIterable<T>(iterable: AsyncIterable<T> | Iterable<T
         },
     };
 }
+
+// Source - https://stackoverflow.com/a/71288323
+/**
+ * Merges multiple async iterables into a single async iterable, yielding values concurrently.
+ *
+ * **Example usage:**
+ * ```ts
+ * const asyncIterable1 = async function* () {
+ *   yield 1; yield 3; yield 5;
+ * };
+ *
+ * const asyncIterable2 = async function* () {
+ *   yield 2; yield 4; yield 6;
+ * };
+ *
+ * for await (const value of mergeAsyncIterables(asyncIterable1(), asyncIterable2())) {
+ *   console.log(value);
+ * }
+ * ```
+ */
+export async function* mergeAsyncIterables<T>(...iterables: AsyncIterable<T>[]): AsyncIterable<T> {
+    const asyncIterators = iterables.map((iterable) => iterable[Symbol.asyncIterator]());
+    const results = [];
+    let count = asyncIterators.length;
+    const never: Promise<never> = new Promise(() => {});
+    async function getNext(asyncIterator: AsyncIterator<T>, index: number) {
+        const result = await asyncIterator.next();
+        return {
+            index,
+            result,
+        };
+    }
+    const nextPromises = asyncIterators.map(getNext);
+    try {
+        while (count) {
+            const { index, result } = await Promise.race(nextPromises);
+            if (result.done) {
+                nextPromises[index] = never;
+                results[index] = result.value;
+                count--;
+            } else {
+                nextPromises[index] = getNext(asyncIterators[index], index);
+                yield result.value;
+            }
+        }
+    } finally {
+        for (const [index, iterator] of asyncIterators.entries()) {
+            // no await here - see https://github.com/tc39/proposal-async-iteration/issues/126
+            if (nextPromises[index] !== never && iterator.return != null) void iterator.return();
+        }
+    }
+    return results;
+}

packages/utils/src/internals/sitemap.ts

@@ -11,6 +11,9 @@ import MIMEType from 'whatwg-mimetype';
 
 import log from '@apify/log';
 
+import { mergeAsyncIterables } from './iterables';
+import { RobotsFile } from './robots';
+
 interface SitemapUrlData {
     loc: string;
     lastmod?: Date;
@@ -436,3 +439,102 @@ export class Sitemap {
         return new Sitemap(urls);
     }
 }
+
+/**
+ * Given a list of URLs, discover related sitemap files for these domains by checking the `robots.txt` file,
+ * the default `sitemap.xml` & `sitemap.txt` files and the URLs themselves.
+ * @param `urls` The list of URLs to discover sitemaps for.
+ * @param `options` Options for sitemap discovery
+ * @returns An async iterable with the discovered sitemap URLs.
+ */
+export async function* discoverValidSitemaps(
+    urls: string[],
+    options: {
+        /**
+         * Proxy URL to be used for network requests.
+         */
+        proxyUrl?: string;
+    } = {},
+): AsyncIterable<string> {
+    const { proxyUrl } = options;
+    const { gotScraping } = await import('got-scraping');
+    const sitemapUrls = new Set<string>();
+
+    const addSitemapUrl = (url: string): string | undefined => {
+        const sizeBefore = sitemapUrls.size;
+
+        sitemapUrls.add(url);
+
+        if (sitemapUrls.size > sizeBefore) {
+            return url;
+        }
+
+        return undefined;
+    };
+
+    const urlExists = (url: string) =>
+        gotScraping({
+            proxyUrl,
+            url,
+            method: 'HEAD',
+        }).then((response) => response.statusCode >= 200 && response.statusCode < 400);
+
+    const discoverSitemapsForDomainUrls = async function* (hostname: string, domainUrls: string[]) {
+        if (!hostname) {
+            return;
+        }
+
+        try {
+            const robotsFile = await RobotsFile.find(domainUrls[0], proxyUrl);
+
+            for (const sitemapUrl of robotsFile.getSitemaps()) {
+                if (addSitemapUrl(sitemapUrl)) {
+                    yield sitemapUrl;
+                }
+            }
+        } catch (err) {
+            log.warning(`Failed to fetch robots.txt file for ${hostname}`, { error: err });
+        }
+
+        const sitemapUrl = domainUrls.find((url) => /sitemap\.(?:xml|txt)(?:\.gz)?$/i.test(url));
+
+        if (sitemapUrl !== undefined) {
+            if (addSitemapUrl(sitemapUrl)) {
+                yield sitemapUrl;
+            }
+        } else {
+            const firstUrl = new URL(domainUrls[0]);
+            firstUrl.pathname = '/sitemap.xml';
+            if (await urlExists(firstUrl.toString())) {
+                if (addSitemapUrl(firstUrl.toString())) {
+                    yield firstUrl.toString();
+                }
+            }
+
+            firstUrl.pathname = '/sitemap.txt';
+            if (await urlExists(firstUrl.toString())) {
+                if (addSitemapUrl(firstUrl.toString())) {
+                    yield firstUrl.toString();
+                }
+            }
+        }
+    };
+
+    const groupedUrls = urls.reduce(
+        (acc, url) => {
+            const hostname = new URL(url)?.hostname ?? '';
+            acc[hostname] ??= [];
+            acc[hostname].push(url);
+            return acc;
+        },
+        {} as Record<string, string[]>,
+    );
+
+    const iterables = Object.entries(groupedUrls).map(([hostname, domainUrls]) =>
+        discoverSitemapsForDomainUrls(hostname, domainUrls),
+    );
+
+    for await (const url of mergeAsyncIterables(...iterables)) {
+        yield url;
+    }
+}

packages/utils/test/sitemap.test.ts

@@ -4,7 +4,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 import log from '@apify/log';
 
 import type { SitemapUrl } from '../src/internals/sitemap';
-import { parseSitemap, Sitemap } from '../src/internals/sitemap';
+import { discoverValidSitemaps, parseSitemap, Sitemap } from '../src/internals/sitemap';
 
 describe('Sitemap', () => {
     beforeEach(() => {
@@ -449,3 +449,114 @@ describe('Sitemap', () => {
         );
     });
 });
+
+describe('discoverValidSitemaps', () => {
+    beforeEach(() => {
+        nock.disableNetConnect();
+    });
+
+    afterEach(() => {
+        nock.cleanAll();
+        nock.enableNetConnect();
+    });
+
+    it('extracts sitemap from robots.txt', async () => {
+        nock('http://sitemap-discovery.com')
+            .get('/robots.txt')
+            .reply(200, 'Sitemap: http://sitemap-discovery.com/some-sitemap.xml')
+            .head('/some-sitemap.xml')
+            .reply(200, '')
+            .head('/sitemap.xml')
+            .reply(404, '')
+            .head('/sitemap.txt')
+            .reply(404, '');
+
+        const urls = [];
+        for await (const url of discoverValidSitemaps(['http://sitemap-discovery.com'])) {
+            urls.push(url);
+        }
+
+        expect(urls).toEqual(['http://sitemap-discovery.com/some-sitemap.xml']);
+    });
+
+    it('extracts sitemap from well-known paths if robots.txt is missing', async () => {
+        nock('http://sitemap-discovery.com')
+            .get('/robots.txt')
+            .reply(404)
+            .head('/sitemap.xml')
+            .reply(200, '')
+            .head('/sitemap.txt')
+            .reply(404, '');
+
+        const urls = [];
+        for await (const url of discoverValidSitemaps(['http://sitemap-discovery.com'])) {
+            urls.push(url);
+        }
+
+        expect(urls).toEqual(['http://sitemap-discovery.com/sitemap.xml']);
+    });
+
+    it('extracts sitemap from well-known paths if robots.txt is missing (txt)', async () => {
+        nock('http://sitemap-discovery.com')
+            .get('/robots.txt')
+            .reply(404)
+            .head('/sitemap.xml')
+            .reply(404, '')
+            .head('/sitemap.txt')
+            .reply(200, '');
+
+        const urls = [];
+        for await (const url of discoverValidSitemaps(['http://sitemap-discovery.com'])) {
+            urls.push(url);
+        }
+
+        expect(urls).toEqual(['http://sitemap-discovery.com/sitemap.txt']);
+    });
+
+    it('extracts sitemap from input url', async () => {
+        nock('http://sitemap-discovery.com').get('/robots.txt').reply(404);
+
+        const urls = [];
+        for await (const url of discoverValidSitemaps(['http://sitemap-discovery.com/sitemap.xml'])) {
+            urls.push(url);
+        }
+
+        expect(urls).toEqual(['http://sitemap-discovery.com/sitemap.xml']);
+    });
+
+    it('extracts sitemaps from multiple domains with mixed order', async () => {
+        nock('http://domain-a.com')
+            .get('/robots.txt')
+            .delay(10)
+            .reply(404)
+            .head('/sitemap.xml')
+            .delay(30)
+            .reply(200, '')
+            .head('/sitemap.txt')
+            .delay(50)
+            .reply(200, '');
+
+        nock('http://domain-b.com')
+            .get('/robots.txt')
+            .delay(20)
+            .reply(404)
+            .head('/sitemap.xml')
+            .delay(40)
+            .reply(200, '')
+            .head('/sitemap.txt')
+            .delay(60)
+            .reply(200, '');
+
+        const urls = [];
+        for await (const url of discoverValidSitemaps(['http://domain-a.com', 'http://domain-b.com'])) {
+            urls.push(url);
+        }
+
+        expect(urls).toEqual([
+            'http://domain-a.com/sitemap.xml',
+            'http://domain-b.com/sitemap.xml',
+            'http://domain-a.com/sitemap.txt',
+            'http://domain-b.com/sitemap.txt',
+        ]);
+    });
+});