Merge pull request #33 from planetlabs/filter

Optional filter for crawler

Author: tschaub

crawler/crawler.go

@@ -104,6 +104,7 @@ type Crawler struct {
 	visitor     Visitor
 	recursion   RecursionType
 	concurrency int
+	filter      func(string) bool
 	queue       workgroup.Queue[*Task]
 }
 
@@ -116,6 +117,12 @@ type Options struct {
 	// a single resource.  Use Children to only visit linked item/child resources.
 	Recursion RecursionType
 
+	// Optional function to limit which resources to crawl.  If provided, the function
+	// will be called with the URL or absolute path to a resource before it is crawled.
+	// If the function returns false, the resource will not be read and the visitor will
+	// not be called.
+	Filter func(string) bool
+
 	Queue workgroup.Queue[*Task]
 }
 
@@ -129,6 +136,9 @@ func (c *Crawler) apply(options *Options) {
 	if options.Queue != nil {
 		c.queue = options.Queue
 	}
+	if options.Filter != nil {
+		c.filter = options.Filter
+	}
 }
 
 // DefaultOptions used when creating a new crawler.
@@ -211,6 +221,9 @@ func (c *Crawler) load(loc *normurl.Locator, value interface{}) error {
 }
 
 func (c *Crawler) crawl(worker *workgroup.Worker[*Task], t *Task) error {
+	if c.filter != nil && !c.filter(t.Url) {
+		return nil
+	}
 	switch t.Type {
 	case resourceTask:
 		return c.crawlResource(worker, t.Url)

crawler/crawler_test.go

@@ -35,6 +35,81 @@ func TestCrawler(t *testing.T) {
 	assert.NoError(t, err)
 
 	assert.Equal(t, uint64(3), count)
+
+	wd, wdErr := os.Getwd()
+	require.NoError(t, wdErr)
+
+	_, visitedCatalog := visited.Load(filepath.Join(wd, "testdata/v1.0.0/catalog-with-collection-of-items.json"))
+	assert.True(t, visitedCatalog)
+
+	_, visitedCollection := visited.Load(filepath.Join(wd, "testdata/v1.0.0/collection-with-items.json"))
+	assert.True(t, visitedCollection)
+
+	_, visitedItem := visited.Load(filepath.Join(wd, "testdata/v1.0.0/item-in-collection.json"))
+	assert.True(t, visitedItem)
+}
+
+func TestCrawlerFilterItem(t *testing.T) {
+	count := uint64(0)
+	visited := &sync.Map{}
+
+	visitor := func(location string, resource crawler.Resource) error {
+		atomic.AddUint64(&count, 1)
+		_, loaded := visited.LoadOrStore(location, true)
+		if loaded {
+			return fmt.Errorf("already visited %s", location)
+		}
+		return nil
+	}
+	c := crawler.New(visitor, &crawler.Options{
+		Filter: func(location string) bool {
+			return !strings.HasSuffix(location, "/item-in-collection.json")
+		},
+	})
+
+	err := c.Crawl(context.Background(), "testdata/v1.0.0/catalog-with-collection-of-items.json")
+	assert.NoError(t, err)
+
+	assert.Equal(t, uint64(2), count)
+
+	wd, wdErr := os.Getwd()
+	require.NoError(t, wdErr)
+
+	_, visitedCatalog := visited.Load(filepath.Join(wd, "testdata/v1.0.0/catalog-with-collection-of-items.json"))
+	assert.True(t, visitedCatalog)
+
+	_, visitedCollection := visited.Load(filepath.Join(wd, "testdata/v1.0.0/collection-with-items.json"))
+	assert.True(t, visitedCollection)
+}
+
+func TestCrawlerFilterCollection(t *testing.T) {
+	count := uint64(0)
+	visited := &sync.Map{}
+
+	visitor := func(location string, resource crawler.Resource) error {
+		atomic.AddUint64(&count, 1)
+		_, loaded := visited.LoadOrStore(location, true)
+		if loaded {
+			return fmt.Errorf("already visited %s", location)
+		}
+		return nil
+	}
+	c := crawler.New(visitor, &crawler.Options{
+		Filter: func(location string) bool {
+			return !strings.HasSuffix(location, "/collection-with-items.json")
+		},
+	})
+
+	err := c.Crawl(context.Background(), "testdata/v1.0.0/catalog-with-collection-of-items.json")
+	assert.NoError(t, err)
+
+	assert.Equal(t, uint64(1), count)
+
+	wd, wdErr := os.Getwd()
+	require.NoError(t, wdErr)
+
+	_, visitedCatalog := visited.Load(filepath.Join(wd, "testdata/v1.0.0/catalog-with-collection-of-items.json"))
+	assert.True(t, visitedCatalog)
 }
 
 func TestCrawlerHTTP(t *testing.T) {