From 9b5d2c9d8fe4b44e322ddcdb51f08e78f696fe25 Mon Sep 17 00:00:00 2001 From: Giorgos Komninos Date: Tue, 21 Jul 2026 18:57:31 +0300 Subject: [PATCH] Add resumable CLI scraping Allow interrupted CLI file scrapes to continue without starting over or duplicating already-written places. Large scrape jobs can fail or be stopped after producing useful partial output. Previously, rerunning the same command would recreate the run from scratch and risk duplicating results or overwriting output. The new resume mode lets users keep existing CSV/JSONL output, append missing places and skip inputs that have already completed. This keeps the default behavior unchanged while giving long-running CLI jobs a practical recovery path. --- README.md | 18 ++ gmaps/completion.go | 6 + gmaps/completion_test.go | 91 +++++++++ gmaps/job.go | 11 + runner/filerunner/filerunner.go | 153 +++++++++++++- runner/filerunner/filerunner_test.go | 189 +++++++++++++++++ runner/jobs.go | 92 ++++++++- runner/jobs_test.go | 69 +++++++ runner/resume/progress.go | 98 +++++++++ runner/resume/progress_test.go | 78 +++++++ runner/resume/results.go | 291 +++++++++++++++++++++++++++ runner/resume/results_test.go | 110 ++++++++++ runner/resume/state.go | 156 ++++++++++++++ runner/resume/state_test.go | 55 +++++ runner/resume/writer.go | 206 +++++++++++++++++++ runner/resume/writer_test.go | 195 ++++++++++++++++++ runner/runner.go | 2 + 17 files changed, 1814 insertions(+), 6 deletions(-) create mode 100644 gmaps/completion.go create mode 100644 gmaps/completion_test.go create mode 100644 runner/filerunner/filerunner_test.go create mode 100644 runner/resume/progress.go create mode 100644 runner/resume/progress_test.go create mode 100644 runner/resume/results.go create mode 100644 runner/resume/results_test.go create mode 100644 runner/resume/state.go create mode 100644 runner/resume/state_test.go create mode 100644 runner/resume/writer.go create mode 100644 runner/resume/writer_test.go diff --git a/README.md b/README.md index bcbe808..0875e00 100644 --- a/README.md +++ b/README.md @@ -273,6 +273,7 @@ Useful options: | Extract emails from business websites | `-email` | | Write JSON instead of CSV | `-json -results /out/results.json` | | Collect extra reviews | `-extra-reviews -json -results /out/results.json` | +| Resume an interrupted file scrape | `-resume -results /out/results.csv` | | Increase concurrency | `-c 4`, `-c 8`, or `-c 16` | | Run multiple pages per browser | `-pages-per-browser 4` | | Limit browser processes | `-browser-pool-size 2` | @@ -493,6 +494,7 @@ Core Options: -input string Path to input file with queries (one per line) -results string Output file path (default: stdout) -json Output JSON instead of CSV + -resume Resume a CLI file scrape by appending missing places -depth int Max scroll depth in results (default: 10) -c int Concurrency level (default: half of CPU cores) @@ -540,6 +542,22 @@ Notes: Run `./google-maps-scraper -h` for the complete list. +### Resuming Interrupted CLI Runs + +Use `-resume` to continue a CLI file scrape after a crash or manual stop: + +```bash +./google-maps-scraper \ + -resume \ + -input queries.txt \ + -results results.csv \ + -depth 10 +``` + +Resume mode reads the existing CSV or JSONL results file, appends new results, and skips places that were already written. It also writes completed input queries to `.resume.json`; future resume runs use that sidecar file to skip fully completed queries. + +`-resume` is only supported with regular file output. It requires `-results` to be a file path and does not support `stdout`, `-writer`, `-leadsdb-api-key`, or `-fast-mode`. Resume runs should use the same input and scrape options as the original run. + ### Using Proxies For larger scraping jobs, proxies help avoid rate limiting. Here's how to configure them: diff --git a/gmaps/completion.go b/gmaps/completion.go new file mode 100644 index 0000000..744932c --- /dev/null +++ b/gmaps/completion.go @@ -0,0 +1,6 @@ +package gmaps + +// CompletionTracker receives per-input job completion signals. +type CompletionTracker interface { + SeedDiscovered(inputID string, placesFound int) error +} diff --git a/gmaps/completion_test.go b/gmaps/completion_test.go new file mode 100644 index 0000000..4c71cb3 --- /dev/null +++ b/gmaps/completion_test.go @@ -0,0 +1,91 @@ +package gmaps_test + +import ( + "context" + "errors" + "os" + "strings" + "testing" + + "github.com/PuerkitoBio/goquery" + "github.com/gosom/google-maps-scraper/gmaps" + "github.com/gosom/scrapemate" + "github.com/stretchr/testify/require" +) + +type recordingCompletionTracker struct { + seedInputID string + seedPlacesFound int +} + +func (t *recordingCompletionTracker) SeedDiscovered(inputID string, placesFound int) error { + t.seedInputID = inputID + t.seedPlacesFound = placesFound + + return nil +} + +func TestGmapJobProcessNotifiesSeedDiscovery(t *testing.T) { + t.Parallel() + + tracker := &recordingCompletionTracker{} + job := gmaps.NewGmapJob("input-1", "en", "coffee", 1, false, "", 0, gmaps.WithGmapCompletionTracker(tracker)) + doc, err := goquery.NewDocumentFromReader(strings.NewReader(` + +
+
+
+
+ + `)) + require.NoError(t, err) + + _, next, err := job.Process(context.Background(), &scrapemate.Response{Document: doc}) + + require.NoError(t, err) + require.Len(t, next, 2) + require.Equal(t, "input-1", tracker.seedInputID) + require.Equal(t, 2, tracker.seedPlacesFound) +} + +func TestPlaceJobProcessDoesNotCompleteInputOnTerminalError(t *testing.T) { + t.Parallel() + + job := gmaps.NewPlaceJob("input-1", "en", "https://www.google.com/maps/place/a", false, false) + + _, next, err := job.Process(context.Background(), &scrapemate.Response{Error: errors.New("fetch failed")}) + + require.Error(t, err) + require.Empty(t, next) +} + +func TestPlaceJobProcessDoesNotCompleteInputBeforeResultIsWritten(t *testing.T) { + t.Parallel() + + job := gmaps.NewPlaceJob("input-1", "en", "https://www.google.com/maps/place/a", false, false) + raw, err := os.ReadFile("../testdata/raw.json") + require.NoError(t, err) + + result, next, err := job.Process(context.Background(), &scrapemate.Response{ + Meta: map[string]any{ + "json": raw, + }, + }) + + require.NoError(t, err) + require.Empty(t, next) + require.NotNil(t, result) +} + +func TestEmailExtractJobProcessDoesNotCompleteInputBeforeResultIsWritten(t *testing.T) { + t.Parallel() + + entry := &gmaps.Entry{ID: "input-1", WebSite: "https://example.com"} + job := gmaps.NewEmailJob("place-1", entry) + + result, next, err := job.Process(context.Background(), &scrapemate.Response{Error: errors.New("fetch failed")}) + + require.NoError(t, err) + require.Empty(t, next) + require.Equal(t, entry, result) +} diff --git a/gmaps/job.go b/gmaps/job.go index 051dc8c..96f2a68 100644 --- a/gmaps/job.go +++ b/gmaps/job.go @@ -34,6 +34,7 @@ type GmapJob struct { ExitMonitor exiter.Exiter ExtractExtraReviews bool WriterManagedCompletion bool + CompletionTracker CompletionTracker } func NewGmapJob( @@ -112,6 +113,12 @@ func WithWriterManagedCompletion() GmapJobOptions { } } +func WithGmapCompletionTracker(tracker CompletionTracker) GmapJobOptions { + return func(j *GmapJob) { + j.CompletionTracker = tracker + } +} + func (j *GmapJob) UseInResults() bool { return false } @@ -186,6 +193,10 @@ func (j *GmapJob) Process(ctx context.Context, resp *scrapemate.Response) (any, j.ExitMonitor.IncrSeedCompleted(1) } + if j.CompletionTracker != nil { + _ = j.CompletionTracker.SeedDiscovered(j.ID, len(next)) + } + log.Info(fmt.Sprintf("%d places found", len(next))) return nil, next, nil diff --git a/runner/filerunner/filerunner.go b/runner/filerunner/filerunner.go index 6e9e4eb..9cf5d03 100644 --- a/runner/filerunner/filerunner.go +++ b/runner/filerunner/filerunner.go @@ -14,6 +14,7 @@ import ( "github.com/gosom/google-maps-scraper/grid" "github.com/gosom/google-maps-scraper/leadsdb" "github.com/gosom/google-maps-scraper/runner" + "github.com/gosom/google-maps-scraper/runner/resume" "github.com/gosom/google-maps-scraper/tlmt" "github.com/gosom/scrapemate" "github.com/gosom/scrapemate/adapters/writers/csvwriter" @@ -27,6 +28,11 @@ type fileRunner struct { writers []scrapemate.ResultWriter app *scrapemateapp.ScrapemateApp outfile *os.File + + resumeIDs *resume.IdentitySet + resumeDedup *resume.IdentitySet + resumeState *resume.State + resumeProgress *resume.ProgressTracker } func New(cfg *runner.Config) (runner.Runner, error) { @@ -34,6 +40,28 @@ func New(cfg *runner.Config) (runner.Runner, error) { return nil, fmt.Errorf("%w: %d", runner.ErrInvalidRunMode, cfg.RunMode) } + if cfg.Resume { + if cfg.FastMode { + return nil, fmt.Errorf("-resume does not support fast mode") + } + + if cfg.ResultsFile == "stdout" { + return nil, fmt.Errorf("-resume requires -results to be a file path") + } + + if err := validateResumeFiles(cfg.ResultsFile); err != nil { + return nil, err + } + + if cfg.CustomWriter != "" { + return nil, fmt.Errorf("-resume does not support custom writers") + } + + if cfg.LeadsDBAPIKey != "" { + return nil, fmt.Errorf("-resume does not support LeadsDB output") + } + } + ans := &fileRunner{ cfg: cfg, } @@ -75,7 +103,13 @@ func (r *fileRunner) Run(ctx context.Context) (err error) { }() dedup := deduper.New() + + if r.cfg.Resume { + dedup = r.resumeDedup + } + exitMonitor := exiter.New() + seedOpts := r.seedJobOptions() if r.cfg.GridBBox != "" { if r.cfg.FastMode { @@ -101,6 +135,7 @@ func (r *fileRunner) Run(ctx context.Context) (err error) { dedup, exitMonitor, r.cfg.ExtraReviews, + seedOpts..., ) } else { seedJobs, err = runner.CreateSeedJobs( @@ -115,6 +150,7 @@ func (r *fileRunner) Run(ctx context.Context) (err error) { dedup, exitMonitor, r.cfg.ExtraReviews, + seedOpts..., ) } @@ -154,6 +190,18 @@ func (r *fileRunner) Close(context.Context) error { return nil } +func (r *fileRunner) seedJobOptions() []runner.SeedJobOption { + if !r.cfg.Resume { + return nil + } + + return []runner.SeedJobOption{ + runner.WithDeterministicSeedIDs(), + runner.WithCompletedInputSkipper(r.resumeState.IsInputCompleted), + runner.WithCompletionTracker(r.resumeProgress), + } +} + func (r *fileRunner) setInput() error { switch r.cfg.InputFile { case "stdin": @@ -189,13 +237,19 @@ func (r *fileRunner) setWriters() error { case r.cfg.LeadsDBAPIKey != "": r.writers = append(r.writers, leadsdb.New(r.cfg.LeadsDBAPIKey)) default: + if !r.cfg.Resume { + if err := removeResumeState(r.cfg.ResultsFile); err != nil { + return err + } + } + var resultsWriter io.Writer switch r.cfg.ResultsFile { case "stdout": resultsWriter = os.Stdout default: - f, err := os.Create(r.cfg.ResultsFile) + f, err := r.openResultsFile() if err != nil { return err } @@ -205,11 +259,37 @@ func (r *fileRunner) setWriters() error { resultsWriter = r.outfile } - csvWriter := csvwriter.NewCsvWriter(csv.NewWriter(resultsWriter)) + switch { + case r.cfg.Resume: + if err := r.initResume(); err != nil { + return err + } - if r.cfg.JSON { + if r.cfg.JSON { + r.writers = append(r.writers, resume.NewJSONLAppendWriter( + resultsWriter, + r.resumeIDs, + r.resumeProgress, + r.outfile.Sync, + )) + } else { + writeHeader, headerErr := r.shouldWriteCSVHeader() + if headerErr != nil { + return headerErr + } + + r.writers = append(r.writers, resume.NewCSVAppendWriter( + csv.NewWriter(resultsWriter), + writeHeader, + r.resumeIDs, + r.resumeProgress, + r.outfile.Sync, + )) + } + case r.cfg.JSON: r.writers = append(r.writers, jsonwriter.NewJSONWriter(resultsWriter)) - } else { + default: + csvWriter := csvwriter.NewCsvWriter(csv.NewWriter(resultsWriter)) r.writers = append(r.writers, csvWriter) } } @@ -217,6 +297,71 @@ func (r *fileRunner) setWriters() error { return nil } +func validateResumeFiles(resultsPath string) error { + if _, err := os.Stat(resultsPath); err == nil { + return nil + } else if !os.IsNotExist(err) { + return err + } + + if _, err := os.Stat(resume.DefaultStatePath(resultsPath)); err == nil { + return fmt.Errorf("resume state exists but results file is missing") + } else if !os.IsNotExist(err) { + return err + } + + return nil +} + +func removeResumeState(resultsPath string) error { + err := os.Remove(resume.DefaultStatePath(resultsPath)) + if os.IsNotExist(err) { + return nil + } + + return err +} + +func (r *fileRunner) openResultsFile() (*os.File, error) { + if r.cfg.Resume { + return os.OpenFile(r.cfg.ResultsFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) + } + + return os.Create(r.cfg.ResultsFile) +} + +func (r *fileRunner) initResume() error { + ids, err := resume.LoadResultIdentities(r.cfg.ResultsFile, r.cfg.JSON) + if err != nil { + return err + } + + state, err := resume.LoadState(resume.DefaultStatePath(r.cfg.ResultsFile)) + if err != nil { + return err + } + + r.resumeIDs = ids + r.resumeDedup = ids.Clone() + r.resumeState = state + r.resumeProgress = resume.NewProgressTracker(state) + + return nil +} + +func (r *fileRunner) shouldWriteCSVHeader() (bool, error) { + info, err := os.Stat(r.cfg.ResultsFile) + if err != nil { + if os.IsNotExist(err) { + return true, nil + } + + return false, err + } + + return info.Size() == 0, nil +} + func (r *fileRunner) setApp() error { opts := []func(*scrapemateapp.Config) error{ // scrapemateapp.WithCache("leveldb", "cache"), diff --git a/runner/filerunner/filerunner_test.go b/runner/filerunner/filerunner_test.go new file mode 100644 index 0000000..19b90de --- /dev/null +++ b/runner/filerunner/filerunner_test.go @@ -0,0 +1,189 @@ +package filerunner_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/gosom/google-maps-scraper/runner" + "github.com/gosom/google-maps-scraper/runner/filerunner" + "github.com/stretchr/testify/require" +) + +func TestNewRejectsResumeWithStdoutResults(t *testing.T) { + t.Parallel() + + cfg := &runner.Config{ + RunMode: runner.RunModeFile, + InputFile: "testdata/input.txt", + ResultsFile: "stdout", + Resume: true, + } + + _, err := filerunner.New(cfg) + + require.ErrorContains(t, err, "-resume requires -results to be a file path") +} + +func TestNewRejectsResumeWithCustomWriter(t *testing.T) { + t.Parallel() + + cfg := &runner.Config{ + RunMode: runner.RunModeFile, + InputFile: "testdata/input.txt", + ResultsFile: "results.csv", + CustomWriter: "/tmp:Writer", + Resume: true, + } + + _, err := filerunner.New(cfg) + + require.ErrorContains(t, err, "-resume does not support custom writers") +} + +func TestNewRejectsResumeWithLeadsDB(t *testing.T) { + t.Parallel() + + cfg := &runner.Config{ + RunMode: runner.RunModeFile, + InputFile: "testdata/input.txt", + ResultsFile: "results.csv", + LeadsDBAPIKey: "key", + Resume: true, + } + + _, err := filerunner.New(cfg) + + require.ErrorContains(t, err, "-resume does not support LeadsDB output") +} + +func TestNewRejectsResumeWithFastMode(t *testing.T) { + t.Parallel() + + cfg := &runner.Config{ + RunMode: runner.RunModeFile, + InputFile: "testdata/input.txt", + ResultsFile: "results.csv", + Resume: true, + FastMode: true, + } + + _, err := filerunner.New(cfg) + + require.ErrorContains(t, err, "-resume does not support fast mode") +} + +func TestNewResumeOpensResultsFileForAppend(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + inputPath := filepath.Join(dir, "input.txt") + resultsPath := filepath.Join(dir, "results.jsonl") + + require.NoError(t, os.WriteFile(inputPath, []byte("coffee\n"), 0o600)) + + existing := `{"input_id":"q1","link":"https://maps/place/existing"}` + "\n" + + require.NoError(t, os.WriteFile(resultsPath, []byte(existing), 0o600)) + + cfg := testConfig(inputPath, resultsPath) + cfg.Resume = true + cfg.JSON = true + + r, err := filerunner.New(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, r.Close(t.Context())) }) + + require.NoError(t, r.Close(t.Context())) + + f, err := os.OpenFile(resultsPath, os.O_WRONLY|os.O_APPEND, 0o600) + require.NoError(t, err) + _, err = f.WriteString("new\n") + require.NoError(t, err) + require.NoError(t, f.Close()) + + got, err := os.ReadFile(resultsPath) + require.NoError(t, err) + require.Equal(t, existing+"new\n", string(got)) +} + +func TestNewWithoutResumeTruncatesResultsFile(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + inputPath := filepath.Join(dir, "input.txt") + resultsPath := filepath.Join(dir, "results.jsonl") + + require.NoError(t, os.WriteFile(inputPath, []byte("coffee\n"), 0o600)) + require.NoError(t, os.WriteFile(resultsPath, []byte("existing\n"), 0o600)) + + r, err := filerunner.New(testConfig(inputPath, resultsPath)) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, r.Close(t.Context())) }) + + require.NoError(t, r.Close(t.Context())) + + got, err := os.ReadFile(resultsPath) + require.NoError(t, err) + require.Empty(t, got) +} + +func TestNewResumeRejectsSidecarWithoutResultsFile(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + inputPath := filepath.Join(dir, "input.txt") + resultsPath := filepath.Join(dir, "results.jsonl") + + require.NoError(t, os.WriteFile(inputPath, []byte("coffee\n"), 0o600)) + require.NoError(t, os.WriteFile(resultsPath+".resume.json", []byte(`{ + "version": 1, + "completed_inputs": ["q1"] +}`), 0o600)) + + cfg := testConfig(inputPath, resultsPath) + cfg.Resume = true + cfg.JSON = true + + _, err := filerunner.New(cfg) + + require.ErrorContains(t, err, "resume state exists but results file is missing") + + _, statErr := os.Stat(resultsPath) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func TestNewWithoutResumeInvalidatesExistingSidecar(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + inputPath := filepath.Join(dir, "input.txt") + resultsPath := filepath.Join(dir, "results.jsonl") + statePath := resultsPath + ".resume.json" + + require.NoError(t, os.WriteFile(inputPath, []byte("coffee\n"), 0o600)) + require.NoError(t, os.WriteFile(resultsPath, []byte("existing\n"), 0o600)) + require.NoError(t, os.WriteFile(statePath, []byte(`{"version":1,"completed_inputs":["q1"]}`), 0o600)) + + r, err := filerunner.New(testConfig(inputPath, resultsPath)) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, r.Close(t.Context())) }) + + _, statErr := os.Stat(statePath) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func testConfig(inputPath, resultsPath string) *runner.Config { + return &runner.Config{ + RunMode: runner.RunModeFile, + InputFile: inputPath, + ResultsFile: resultsPath, + Concurrency: 1, + MaxDepth: 1, + LangCode: "en", + Zoom: 15, + Radius: 10000, + ExitOnInactivityDuration: 0, + MaxPagesPerBrowser: 1, + } +} diff --git a/runner/jobs.go b/runner/jobs.go index 9ceeb5b..8cf340b 100644 --- a/runner/jobs.go +++ b/runner/jobs.go @@ -2,6 +2,7 @@ package runner import ( "bufio" + "crypto/sha256" "fmt" "io" "os" @@ -18,6 +19,36 @@ import ( "github.com/gosom/scrapemate" ) +type seedJobConfig struct { + completedInputSkipper func(string) bool + completionTracker gmaps.CompletionTracker + deterministicIDs bool +} + +// SeedJobOption configures seed job creation. +type SeedJobOption func(*seedJobConfig) + +// WithCompletedInputSkipper skips seed jobs whose input IDs are complete. +func WithCompletedInputSkipper(skipper func(string) bool) SeedJobOption { + return func(cfg *seedJobConfig) { + cfg.completedInputSkipper = skipper + } +} + +// WithCompletionTracker attaches a per-input completion tracker to created jobs. +func WithCompletionTracker(tracker gmaps.CompletionTracker) SeedJobOption { + return func(cfg *seedJobConfig) { + cfg.completionTracker = tracker + } +} + +// WithDeterministicSeedIDs derives stable IDs for input lines without custom IDs. +func WithDeterministicSeedIDs() SeedJobOption { + return func(cfg *seedJobConfig) { + cfg.deterministicIDs = true + } +} + func CreateSeedJobs( fastmode bool, langCode string, @@ -30,7 +61,10 @@ func CreateSeedJobs( dedup deduper.Deduper, exitMonitor exiter.Exiter, extraReviews bool, + opts ...SeedJobOption, ) (jobs []scrapemate.IJob, err error) { + createCfg := newSeedJobConfig(opts...) + var lat, lon float64 if fastmode { @@ -85,6 +119,14 @@ func CreateSeedJobs( query := q.text id := q.id + if createCfg.deterministicIDs && id == "" { + id = deterministicSeedID(query) + } + + if createCfg.completedInputSkipper != nil && createCfg.completedInputSkipper(id) { + continue + } + var job scrapemate.IJob if !fastmode { @@ -102,6 +144,10 @@ func CreateSeedJobs( opts = append(opts, gmaps.WithExtraReviews()) } + if createCfg.completionTracker != nil { + opts = append(opts, gmaps.WithGmapCompletionTracker(createCfg.completionTracker)) + } + job = gmaps.NewGmapJob(id, langCode, query, maxDepth, email, geoCoordinates, zoom, opts...) } else { jparams := gmaps.MapSearchParams{ @@ -123,7 +169,13 @@ func CreateSeedJobs( opts = append(opts, gmaps.WithSearchJobExitMonitor(exitMonitor)) } - job = gmaps.NewSearchJob(&jparams, opts...) + searchJob := gmaps.NewSearchJob(&jparams, opts...) + + if id != "" { + searchJob.ID = id + } + + job = searchJob } jobs = append(jobs, job) @@ -149,7 +201,10 @@ func CreateGridSeedJobs( dedup deduper.Deduper, exitMonitor exiter.Exiter, extraReviews bool, + opts ...SeedJobOption, ) ([]scrapemate.IJob, error) { + createCfg := newSeedJobConfig(opts...) + if zoom < 1 || zoom > 21 { return nil, fmt.Errorf("invalid zoom level: %d", zoom) } @@ -177,10 +232,20 @@ func CreateGridSeedJobs( for _, cell := range cells { // Each cell gets a unique ID derived from the query ID (or a new UUID). cellID := uuid.NewV4().String() - if queryID != "" { + + switch { + case createCfg.deterministicIDs && queryID != "": + cellID = deterministicSeedID(queryID, cell.GeoCoordinates()) + case createCfg.deterministicIDs: + cellID = deterministicSeedID(queryText, cell.GeoCoordinates()) + case queryID != "": cellID = fmt.Sprintf("%s-%s", queryID, cellID) } + if createCfg.completedInputSkipper != nil && createCfg.completedInputSkipper(cellID) { + continue + } + opts := []gmaps.GmapJobOptions{} if dedup != nil { @@ -195,6 +260,10 @@ func CreateGridSeedJobs( opts = append(opts, gmaps.WithExtraReviews()) } + if createCfg.completionTracker != nil { + opts = append(opts, gmaps.WithGmapCompletionTracker(createCfg.completionTracker)) + } + job := gmaps.NewGmapJob( cellID, langCode, @@ -213,6 +282,25 @@ func CreateGridSeedJobs( return jobs, nil } +func newSeedJobConfig(opts ...SeedJobOption) seedJobConfig { + var cfg seedJobConfig + for _, opt := range opts { + opt(&cfg) + } + + return cfg +} + +func deterministicSeedID(parts ...string) string { + hash := sha256.New() + for _, part := range parts { + _, _ = hash.Write([]byte(part)) + _, _ = hash.Write([]byte{0}) + } + + return fmt.Sprintf("resume:%x", hash.Sum(nil)) +} + // query holds a parsed input line. type query struct { text string diff --git a/runner/jobs_test.go b/runner/jobs_test.go index 36ac65b..5cb171a 100644 --- a/runner/jobs_test.go +++ b/runner/jobs_test.go @@ -4,6 +4,7 @@ import ( "strings" "testing" + "github.com/gosom/google-maps-scraper/gmaps" "github.com/gosom/google-maps-scraper/grid" "github.com/gosom/google-maps-scraper/runner" ) @@ -82,3 +83,71 @@ func TestCreateGridSeedJobsRejectsEmptyQueryBeforeCustomID(t *testing.T) { t.Fatalf("expected empty query text error, got %v", err) } } + +func TestCreateSeedJobsSkipsCompletedInputs(t *testing.T) { + t.Parallel() + + jobs, err := runner.CreateSeedJobs( + false, + "en", + strings.NewReader("coffee\ntea\n"), + 10, + false, + "", + 15, + 10000, + nil, + nil, + false, + runner.WithDeterministicSeedIDs(), + runner.WithCompletedInputSkipper(func(inputID string) bool { + return strings.HasPrefix(inputID, "resume:") + }), + ) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(jobs) != 0 { + t.Fatalf("expected completed deterministic inputs to be skipped, got %d jobs", len(jobs)) + } +} + +func TestCreateSeedJobsAddsCompletionTracker(t *testing.T) { + t.Parallel() + + tracker := &recordingSeedTracker{} + + jobs, err := runner.CreateSeedJobs( + false, + "en", + strings.NewReader("coffee\n"), + 10, + false, + "", + 15, + 10000, + nil, + nil, + false, + runner.WithCompletionTracker(tracker), + ) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + job, ok := jobs[0].(*gmaps.GmapJob) + if !ok { + t.Fatalf("expected *gmaps.GmapJob, got %T", jobs[0]) + } + + if job.CompletionTracker != tracker { + t.Fatalf("completion tracker was not attached") + } +} + +type recordingSeedTracker struct{} + +func (t *recordingSeedTracker) SeedDiscovered(string, int) error { + return nil +} diff --git a/runner/resume/progress.go b/runner/resume/progress.go new file mode 100644 index 0000000..2c6a1f7 --- /dev/null +++ b/runner/resume/progress.go @@ -0,0 +1,98 @@ +package resume + +import "sync" + +// InputMarker persists fully completed input IDs. +type InputMarker interface { + MarkInputCompleted(inputID string) error +} + +type inputProgress struct { + seedDone bool + placesFound int + placesCompleted int + marked bool +} + +// ProgressTracker tracks per-input discovery and child job completion. +type ProgressTracker struct { + mu sync.Mutex + marker InputMarker + progress map[string]*inputProgress +} + +// NewProgressTracker creates a per-input completion tracker. +func NewProgressTracker(marker InputMarker) *ProgressTracker { + return &ProgressTracker{ + marker: marker, + progress: make(map[string]*inputProgress), + } +} + +// SeedDiscovered records that an input's search job finished discovery. +func (t *ProgressTracker) SeedDiscovered(inputID string, placesFound int) error { + if inputID == "" { + return nil + } + + t.mu.Lock() + defer t.mu.Unlock() + + progress := t.getLocked(inputID) + progress.seedDone = true + progress.placesFound += placesFound + + return t.maybeMarkLocked(inputID, progress) +} + +// ResultPersisted records that one output-producing child was durably handled. +func (t *ProgressTracker) ResultPersisted(inputID string, syncOutput func() error) error { + if inputID == "" { + return nil + } + + t.mu.Lock() + defer t.mu.Unlock() + + progress := t.getLocked(inputID) + progress.placesCompleted++ + + return t.maybeMarkLocked(inputID, progress, syncOutput) +} + +func (t *ProgressTracker) getLocked(inputID string) *inputProgress { + progress, ok := t.progress[inputID] + if ok { + return progress + } + + progress = &inputProgress{} + t.progress[inputID] = progress + + return progress +} + +func (t *ProgressTracker) maybeMarkLocked(inputID string, progress *inputProgress, syncOutput ...func() error) error { + if progress.marked || !progress.seedDone || progress.placesCompleted < progress.placesFound { + return nil + } + + if len(syncOutput) > 0 && syncOutput[0] != nil { + if err := syncOutput[0](); err != nil { + return err + } + } + + if t.marker == nil { + progress.marked = true + return nil + } + + if err := t.marker.MarkInputCompleted(inputID); err != nil { + return err + } + + progress.marked = true + + return nil +} diff --git a/runner/resume/progress_test.go b/runner/resume/progress_test.go new file mode 100644 index 0000000..3a7bc4b --- /dev/null +++ b/runner/resume/progress_test.go @@ -0,0 +1,78 @@ +package resume_test + +import ( + "testing" + + "github.com/gosom/google-maps-scraper/runner/resume" + "github.com/stretchr/testify/require" +) + +type recordingMarker struct { + completed map[string]int + events []string +} + +func newRecordingMarker() *recordingMarker { + return &recordingMarker{completed: make(map[string]int)} +} + +func (m *recordingMarker) MarkInputCompleted(inputID string) error { + m.events = append(m.events, "marked") + m.completed[inputID]++ + + return nil +} + +func TestProgressMarksZeroResultInputComplete(t *testing.T) { + t.Parallel() + + marker := newRecordingMarker() + tracker := resume.NewProgressTracker(marker) + + require.NoError(t, tracker.SeedDiscovered("q1", 0)) + + require.Equal(t, 1, marker.completed["q1"]) +} + +func TestProgressWaitsForPlaces(t *testing.T) { + t.Parallel() + + marker := newRecordingMarker() + tracker := resume.NewProgressTracker(marker) + + require.NoError(t, tracker.SeedDiscovered("q1", 2)) + require.Zero(t, marker.completed["q1"]) + require.NoError(t, tracker.ResultPersisted("q1", nil)) + require.Zero(t, marker.completed["q1"]) + require.NoError(t, tracker.ResultPersisted("q1", nil)) + require.Equal(t, 1, marker.completed["q1"]) +} + +func TestProgressMarksInputOnlyOnce(t *testing.T) { + t.Parallel() + + marker := newRecordingMarker() + tracker := resume.NewProgressTracker(marker) + + require.NoError(t, tracker.SeedDiscovered("q1", 1)) + require.NoError(t, tracker.ResultPersisted("q1", nil)) + require.NoError(t, tracker.ResultPersisted("q1", nil)) + + require.Equal(t, 1, marker.completed["q1"]) +} + +func TestProgressSyncsOutputBeforeMarkingInputComplete(t *testing.T) { + t.Parallel() + + marker := newRecordingMarker() + tracker := resume.NewProgressTracker(marker) + require.NoError(t, tracker.SeedDiscovered("q1", 1)) + + err := tracker.ResultPersisted("q1", func() error { + marker.events = append(marker.events, "synced") + return nil + }) + + require.NoError(t, err) + require.Equal(t, []string{"synced", "marked"}, marker.events) +} diff --git a/runner/resume/results.go b/runner/resume/results.go new file mode 100644 index 0000000..fac59cd --- /dev/null +++ b/runner/resume/results.go @@ -0,0 +1,291 @@ +package resume + +import ( + "context" + "encoding/csv" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "strings" + "sync" + + "github.com/gosom/google-maps-scraper/gmaps" +) + +// IdentitySet tracks stable place identities already emitted to a result file. +type IdentitySet struct { + mu sync.RWMutex + keys map[string]struct{} +} + +// NewIdentitySet creates an empty result identity set. +func NewIdentitySet() *IdentitySet { + return &IdentitySet{ + keys: make(map[string]struct{}), + } +} + +// Add records a stable identity key. Empty keys are ignored. +func (s *IdentitySet) Add(key string) { + key = strings.TrimSpace(key) + if key == "" { + return + } + + s.mu.Lock() + defer s.mu.Unlock() + + s.keys[key] = struct{}{} +} + +// Has reports whether key has already been recorded. +func (s *IdentitySet) Has(key string) bool { + key = strings.TrimSpace(key) + if key == "" { + return false + } + + s.mu.RLock() + defer s.mu.RUnlock() + + _, ok := s.keys[key] + + return ok +} + +// Clone returns an independent copy of the identity set. +func (s *IdentitySet) Clone() *IdentitySet { + clone := NewIdentitySet() + + s.mu.RLock() + defer s.mu.RUnlock() + + for key := range s.keys { + clone.keys[key] = struct{}{} + } + + return clone +} + +// HasEntry reports whether any stable identity for entry is already recorded. +func (s *IdentitySet) HasEntry(entry *gmaps.Entry) bool { + if entry == nil { + return false + } + + s.mu.RLock() + defer s.mu.RUnlock() + + for _, key := range entryIdentities(entry) { + if _, ok := s.keys[key]; ok { + return true + } + } + + return false +} + +// AddIfNotExists implements deduper.Deduper for discovered place URLs. +func (s *IdentitySet) AddIfNotExists(_ context.Context, key string) bool { + key = strings.TrimSpace(key) + if key == "" { + return false + } + + s.mu.Lock() + defer s.mu.Unlock() + + if _, ok := s.keys[key]; ok { + return false + } + + s.keys[key] = struct{}{} + + return true +} + +// AddEntry records all stable identities for entry. +func (s *IdentitySet) AddEntry(entry *gmaps.Entry) { + if entry == nil { + return + } + + for _, key := range entryIdentities(entry) { + s.Add(key) + } +} + +func entryIdentities(entry *gmaps.Entry) []string { + values := []string{entry.PlaceID, entry.Cid, entry.DataID, entry.Link} + identities := make([]string, 0, len(values)) + + for _, value := range values { + if value = strings.TrimSpace(value); value != "" { + identities = append(identities, value) + } + } + + return identities +} + +// EntryIdentity returns the preferred stable identity for entry. +func EntryIdentity(entry *gmaps.Entry) string { + if entry == nil { + return "" + } + + return firstNonEmpty(entry.PlaceID, entry.Cid, entry.DataID, entry.Link) +} + +// LoadResultIdentities reads existing CSV or JSONL results into an IdentitySet. +func LoadResultIdentities(path string, jsonl bool) (*IdentitySet, error) { + ids, err := loadResultIdentities(path, jsonl) + if err == nil || !hasIncompleteTrailingRecord(path) { + return ids, err + } + + if err := truncateIncompleteTrailingRecord(path); err != nil { + return nil, fmt.Errorf("repair incomplete result record: %w", err) + } + + return loadResultIdentities(path, jsonl) +} + +func loadResultIdentities(path string, jsonl bool) (*IdentitySet, error) { + ids := NewIdentitySet() + + f, err := os.Open(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return ids, nil + } + + return nil, err + } + defer f.Close() + + info, err := f.Stat() + if err != nil { + return nil, err + } + + if info.Size() == 0 { + return ids, nil + } + + if jsonl { + return ids, loadJSONLIdentities(f, ids) + } + + return ids, loadCSVIdentities(f, ids) +} + +func hasIncompleteTrailingRecord(path string) bool { + data, err := os.ReadFile(path) + if err != nil || len(data) == 0 { + return false + } + + return data[len(data)-1] != '\n' +} + +func truncateIncompleteTrailingRecord(path string) error { + data, err := os.ReadFile(path) + if err != nil { + return err + } + + lastNewline := strings.LastIndexByte(string(data), '\n') + if lastNewline < 0 { + lastNewline = 0 + } else { + lastNewline++ + } + + file, err := os.OpenFile(path, os.O_WRONLY, 0) + if err != nil { + return err + } + + if err := file.Truncate(int64(lastNewline)); err != nil { + _ = file.Close() + return err + } + + return file.Close() +} + +func loadJSONLIdentities(r io.Reader, ids *IdentitySet) error { + dec := json.NewDecoder(r) + + for { + var entry gmaps.Entry + if err := dec.Decode(&entry); err != nil { + if errors.Is(err, io.EOF) { + break + } + + return fmt.Errorf("parse JSONL result: %w", err) + } + + ids.AddEntry(&entry) + } + + return nil +} + +func loadCSVIdentities(r io.Reader, ids *IdentitySet) error { + reader := csv.NewReader(r) + + headers, err := reader.Read() + if err != nil { + if errors.Is(err, io.EOF) { + return nil + } + + return fmt.Errorf("read CSV header: %w", err) + } + + indexes := map[string]int{} + for i, header := range headers { + indexes[strings.TrimSpace(header)] = i + } + + for { + row, err := reader.Read() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + + return fmt.Errorf("read CSV row: %w", err) + } + + addCSVValues(ids, row, indexes, "place_id", "cid", "data_id", "link") + } + + return nil +} + +func addCSVValues(ids *IdentitySet, row []string, indexes map[string]int, names ...string) { + for _, name := range names { + index, ok := indexes[name] + if !ok || index >= len(row) { + continue + } + + ids.Add(row[index]) + } +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value = strings.TrimSpace(value); value != "" { + return value + } + } + + return "" +} diff --git a/runner/resume/results_test.go b/runner/resume/results_test.go new file mode 100644 index 0000000..3a38dd7 --- /dev/null +++ b/runner/resume/results_test.go @@ -0,0 +1,110 @@ +package resume_test + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/gosom/google-maps-scraper/gmaps" + "github.com/gosom/google-maps-scraper/runner/resume" + "github.com/stretchr/testify/require" +) + +func TestLoadCSVIdentities(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "results.csv") + body := "input_id,link,title,cid,data_id,place_id\n" + + "q1,https://maps/place/a,A,cid-a,data-a,place-a\n" + + "q2,https://maps/place/b,B,,data-b,\n" + require.NoError(t, os.WriteFile(path, []byte(body), 0o600)) + + ids, err := resume.LoadResultIdentities(path, false) + + require.NoError(t, err) + require.True(t, ids.Has("place-a")) + require.True(t, ids.Has("cid-a")) + require.True(t, ids.Has("data-b")) +} + +func TestLoadJSONLIdentities(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "results.jsonl") + body := `{"input_id":"q1","link":"https://maps/place/a","place_id":"place-a"}` + "\n" + + `{"input_id":"q2","link":"https://maps/place/b","cid":"cid-b"}` + "\n" + require.NoError(t, os.WriteFile(path, []byte(body), 0o600)) + + ids, err := resume.LoadResultIdentities(path, true) + + require.NoError(t, err) + require.True(t, ids.Has("place-a")) + require.True(t, ids.Has("cid-b")) +} + +func TestLoadJSONLIdentitiesSupportsLargeEntries(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "results.jsonl") + body := `{"input_id":"q1","link":"https://maps/place/a","place_id":"place-a","description":"` + + strings.Repeat("x", 70*1024) + `"}` + "\n" + require.NoError(t, os.WriteFile(path, []byte(body), 0o600)) + + ids, err := resume.LoadResultIdentities(path, true) + + require.NoError(t, err) + require.True(t, ids.Has("place-a")) +} + +func TestLoadResultIdentitiesMissingFile(t *testing.T) { + t.Parallel() + + ids, err := resume.LoadResultIdentities(filepath.Join(t.TempDir(), "missing.csv"), false) + + require.NoError(t, err) + require.False(t, ids.Has("anything")) +} + +func TestIdentitySetCloneIsIndependent(t *testing.T) { + t.Parallel() + + emitted := resume.NewIdentitySet() + emitted.Add("https://maps/place/existing") + discovered := emitted.Clone() + + discovered.Add("https://maps/place/new") + + require.True(t, discovered.Has("https://maps/place/new")) + require.False(t, emitted.Has("https://maps/place/new")) +} + +func TestIdentitySetMatchesAnyEntryIdentity(t *testing.T) { + t.Parallel() + + ids := resume.NewIdentitySet() + ids.Add("cid-a") + + require.True(t, ids.HasEntry(&gmaps.Entry{ + PlaceID: "place-a", + Cid: "cid-a", + })) +} + +func TestLoadJSONLIdentitiesRepairsIncompleteTrailingRecord(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "results.jsonl") + body := `{"place_id":"place-a"}` + "\n" + `{"place_id":"place-b"` + require.NoError(t, os.WriteFile(path, []byte(body), 0o600)) + + ids, err := resume.LoadResultIdentities(path, true) + + require.NoError(t, err) + require.True(t, ids.Has("place-a")) + require.False(t, ids.Has("place-b")) + + repaired, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, `{"place_id":"place-a"}`+"\n", string(repaired)) +} diff --git a/runner/resume/state.go b/runner/resume/state.go new file mode 100644 index 0000000..54b83d8 --- /dev/null +++ b/runner/resume/state.go @@ -0,0 +1,156 @@ +package resume + +import ( + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "slices" + "sync" +) + +const stateVersion = 1 + +type stateFile struct { + Version int `json:"version"` + CompletedInput []string `json:"completed_inputs"` +} + +// State tracks input queries that fully completed in a resume-capable run. +type State struct { + path string + mu sync.Mutex + completed map[string]struct{} +} + +// DefaultStatePath returns the sidecar path for a results file. +func DefaultStatePath(resultsPath string) string { + return resultsPath + ".resume.json" +} + +// LoadState reads a resume sidecar. Missing files produce an empty state. +func LoadState(path string) (*State, error) { + state := &State{ + path: path, + completed: make(map[string]struct{}), + } + + data, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return state, nil + } + + return nil, err + } + + if len(data) == 0 { + return state, nil + } + + var file stateFile + if err := json.Unmarshal(data, &file); err != nil { + return nil, err + } + + if file.Version != stateVersion { + return nil, fmt.Errorf("unsupported resume state version: %d", file.Version) + } + + for _, inputID := range file.CompletedInput { + if inputID != "" { + state.completed[inputID] = struct{}{} + } + } + + return state, nil +} + +// IsInputCompleted reports whether inputID has been marked complete. +func (s *State) IsInputCompleted(inputID string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + _, ok := s.completed[inputID] + + return ok +} + +// MarkInputCompleted records and persists a completed input ID. +func (s *State) MarkInputCompleted(inputID string) error { + if inputID == "" { + return nil + } + + s.mu.Lock() + defer s.mu.Unlock() + + if _, ok := s.completed[inputID]; ok { + return nil + } + + completed := make(map[string]struct{}, len(s.completed)+1) + for completedInputID := range s.completed { + completed[completedInputID] = struct{}{} + } + + completed[inputID] = struct{}{} + + if err := s.saveLocked(completed); err != nil { + return err + } + + s.completed = completed + + return nil +} + +func (s *State) saveLocked(completedSet map[string]struct{}) error { + completed := make([]string, 0, len(completedSet)) + for inputID := range completedSet { + completed = append(completed, inputID) + } + + slices.Sort(completed) + + data, err := json.MarshalIndent(stateFile{ + Version: stateVersion, + CompletedInput: completed, + }, "", " ") + if err != nil { + return err + } + + data = append(data, '\n') + + tmp, err := os.CreateTemp(filepath.Dir(s.path), ".resume-*.tmp") + if err != nil { + return err + } + + tmpPath := tmp.Name() + + defer func() { _ = os.Remove(tmpPath) }() + + if err := tmp.Chmod(0o600); err != nil { + _ = tmp.Close() + return err + } + + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return err + } + + if err := tmp.Sync(); err != nil { + _ = tmp.Close() + return err + } + + if err := tmp.Close(); err != nil { + return err + } + + return os.Rename(tmpPath, s.path) +} diff --git a/runner/resume/state_test.go b/runner/resume/state_test.go new file mode 100644 index 0000000..bcae346 --- /dev/null +++ b/runner/resume/state_test.go @@ -0,0 +1,55 @@ +package resume_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/gosom/google-maps-scraper/runner/resume" + "github.com/stretchr/testify/require" +) + +func TestStateLoadMissingFile(t *testing.T) { + t.Parallel() + + state, err := resume.LoadState(filepath.Join(t.TempDir(), "missing.resume.json")) + + require.NoError(t, err) + require.False(t, state.IsInputCompleted("q1")) +} + +func TestStateMarkInputCompletedPersists(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "results.csv.resume.json") + state, err := resume.LoadState(path) + require.NoError(t, err) + + require.NoError(t, state.MarkInputCompleted("q1")) + + reloaded, err := resume.LoadState(path) + require.NoError(t, err) + require.True(t, reloaded.IsInputCompleted("q1")) +} + +func TestDefaultStatePath(t *testing.T) { + t.Parallel() + + require.Equal(t, "results.csv.resume.json", resume.DefaultStatePath("results.csv")) +} + +func TestStateDoesNotMarkInputCompletedWhenPersistenceFails(t *testing.T) { + t.Parallel() + + dir := filepath.Join(t.TempDir(), "state") + require.NoError(t, os.Mkdir(dir, 0o700)) + path := filepath.Join(dir, "results.csv.resume.json") + state, err := resume.LoadState(path) + require.NoError(t, err) + require.NoError(t, os.Remove(dir)) + + err = state.MarkInputCompleted("q1") + + require.Error(t, err) + require.False(t, state.IsInputCompleted("q1")) +} diff --git a/runner/resume/writer.go b/runner/resume/writer.go new file mode 100644 index 0000000..4a1460e --- /dev/null +++ b/runner/resume/writer.go @@ -0,0 +1,206 @@ +package resume + +import ( + "context" + "encoding/csv" + "encoding/json" + "fmt" + "io" + "reflect" + + "github.com/gosom/google-maps-scraper/gmaps" + "github.com/gosom/scrapemate" +) + +var ( + _ scrapemate.ResultWriter = (*csvAppendWriter)(nil) + _ scrapemate.ResultWriter = (*jsonlAppendWriter)(nil) +) + +type csvAppendWriter struct { + w *csv.Writer + writeHeader bool + headerDone bool + ids *IdentitySet + tracker ResultTracker + syncOutput func() error +} + +// ResultTracker acknowledges results after they have been persisted. +type ResultTracker interface { + ResultPersisted(inputID string, syncOutput func() error) error +} + +// NewCSVAppendWriter creates a CSV writer suitable for resume append mode. +func NewCSVAppendWriter( + w *csv.Writer, + writeHeader bool, + ids *IdentitySet, + tracker ResultTracker, + syncOutput func() error, +) scrapemate.ResultWriter { + return &csvAppendWriter{ + w: w, + writeHeader: writeHeader, + ids: ids, + tracker: tracker, + syncOutput: syncOutput, + } +} + +func (w *csvAppendWriter) Run(_ context.Context, in <-chan scrapemate.Result) error { + for result := range in { + entries, err := entriesFromResult(result.Data) + if err != nil { + return err + } + + if len(entries) == 0 { + continue + } + + if err := w.writeHeaderIfNeeded(entries[0]); err != nil { + return err + } + + for _, entry := range entries { + if w.ids != nil && w.ids.HasEntry(entry) { + continue + } + + if err := w.w.Write(entry.CsvRow()); err != nil { + return err + } + + if w.ids != nil { + w.ids.AddEntry(entry) + } + } + + w.w.Flush() + + if err := w.w.Error(); err != nil { + return err + } + + for _, entry := range entries { + if err := acknowledgeResult(w.tracker, entry, w.syncOutput); err != nil { + return err + } + } + } + + return w.w.Error() +} + +func (w *csvAppendWriter) writeHeaderIfNeeded(entry *gmaps.Entry) error { + if !w.writeHeader || w.headerDone { + return nil + } + + if err := w.w.Write(entry.CsvHeaders()); err != nil { + return err + } + + w.w.Flush() + + if err := w.w.Error(); err != nil { + return err + } + + w.headerDone = true + + return nil +} + +type jsonlAppendWriter struct { + enc *json.Encoder + ids *IdentitySet + tracker ResultTracker + syncOutput func() error +} + +// NewJSONLAppendWriter creates a JSONL writer suitable for resume append mode. +func NewJSONLAppendWriter( + w io.Writer, + ids *IdentitySet, + tracker ResultTracker, + syncOutput func() error, +) scrapemate.ResultWriter { + return &jsonlAppendWriter{ + enc: json.NewEncoder(w), + ids: ids, + tracker: tracker, + syncOutput: syncOutput, + } +} + +func (w *jsonlAppendWriter) Run(_ context.Context, in <-chan scrapemate.Result) error { + for result := range in { + entries, err := entriesFromResult(result.Data) + if err != nil { + return err + } + + for _, entry := range entries { + if w.ids == nil || !w.ids.HasEntry(entry) { + if err := w.enc.Encode(entry); err != nil { + return err + } + + if w.ids != nil { + w.ids.AddEntry(entry) + } + } + + if err := acknowledgeResult(w.tracker, entry, w.syncOutput); err != nil { + return err + } + } + } + + return nil +} + +func acknowledgeResult(tracker ResultTracker, entry *gmaps.Entry, syncOutput func() error) error { + if tracker == nil || entry == nil { + return nil + } + + if err := tracker.ResultPersisted(entry.ID, syncOutput); err != nil { + return fmt.Errorf("acknowledge persisted result: %w", err) + } + + return nil +} + +func entriesFromResult(data any) ([]*gmaps.Entry, error) { + if data == nil { + return nil, nil + } + + if entry, ok := data.(*gmaps.Entry); ok { + return []*gmaps.Entry{entry}, nil + } + + value := reflect.ValueOf(data) + if value.Kind() != reflect.Slice { + return nil, fmt.Errorf("unexpected resume writer data type: %T", data) + } + + length := value.Len() + entries := make([]*gmaps.Entry, 0, length) + + for i := 0; i < length; i++ { + item := value.Index(i).Interface() + entry, ok := item.(*gmaps.Entry) + + if !ok { + return nil, fmt.Errorf("unexpected resume writer slice item type: %T", item) + } + + entries = append(entries, entry) + } + + return entries, nil +} diff --git a/runner/resume/writer_test.go b/runner/resume/writer_test.go new file mode 100644 index 0000000..9251485 --- /dev/null +++ b/runner/resume/writer_test.go @@ -0,0 +1,195 @@ +package resume_test + +import ( + "bytes" + "context" + "encoding/csv" + "errors" + "strings" + "testing" + + "github.com/gosom/google-maps-scraper/gmaps" + "github.com/gosom/google-maps-scraper/runner/resume" + "github.com/gosom/scrapemate" + "github.com/stretchr/testify/require" +) + +func TestCSVAppendWriterWritesHeaderForNewFile(t *testing.T) { + t.Parallel() + + var out bytes.Buffer + + ids := resume.NewIdentitySet() + writer := resume.NewCSVAppendWriter(csv.NewWriter(&out), true, ids, nil, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: &gmaps.Entry{ID: "q1", Title: "A", Link: "https://maps/place/a", PlaceID: "place-a"}} + + close(in) + + require.NoError(t, writer.Run(context.Background(), in)) + + lines := strings.Split(strings.TrimSpace(out.String()), "\n") + require.Equal(t, "input_id,link,title,category,address,open_hours,popular_times,website,phone,plus_code,review_count,review_rating,reviews_per_rating,latitude,longitude,cid,status,descriptions,reviews_link,thumbnail,timezone,price_range,data_id,street_view_url,place_id,images,reservations,order_online,menu,owner,complete_address,credit_cards_accepted,about,user_reviews,user_reviews_extended,emails", lines[0]) + require.Len(t, lines, 2) + require.True(t, ids.Has("place-a")) +} + +func TestCSVAppendWriterReturnsHeaderWriteError(t *testing.T) { + t.Parallel() + + ids := resume.NewIdentitySet() + writer := resume.NewCSVAppendWriter(csv.NewWriter(failingWriter{}), true, ids, nil, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: &gmaps.Entry{ID: "q1", Title: "A", Link: "https://maps/place/a", PlaceID: "place-a"}} + + close(in) + + require.ErrorContains(t, writer.Run(context.Background(), in), "write failed") + require.False(t, ids.Has("place-a")) +} + +type failingWriter struct{} + +func (failingWriter) Write([]byte) (int, error) { + return 0, errors.New("write failed") +} + +func TestCSVAppendWriterSkipsHeaderForExistingFile(t *testing.T) { + t.Parallel() + + var out bytes.Buffer + + ids := resume.NewIdentitySet() + writer := resume.NewCSVAppendWriter(csv.NewWriter(&out), false, ids, nil, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: &gmaps.Entry{ID: "q1", Title: "A", Link: "https://maps/place/a", PlaceID: "place-a"}} + + close(in) + + require.NoError(t, writer.Run(context.Background(), in)) + + require.NotContains(t, out.String(), "input_id,link,title") + require.True(t, ids.Has("place-a")) +} + +func TestJSONLAppendWriterWritesOneObjectPerLine(t *testing.T) { + t.Parallel() + + var out bytes.Buffer + + ids := resume.NewIdentitySet() + writer := resume.NewJSONLAppendWriter(&out, ids, nil, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: []*gmaps.Entry{ + {ID: "q1", Title: "A", Link: "https://maps/place/a", PlaceID: "place-a"}, + {ID: "q1", Title: "B", Link: "https://maps/place/b", Cid: "cid-b"}, + }} + + close(in) + + require.NoError(t, writer.Run(context.Background(), in)) + + lines := strings.Split(strings.TrimSpace(out.String()), "\n") + require.Len(t, lines, 2) + require.JSONEq(t, `{"input_id":"q1","link":"https://maps/place/a","title":"A","place_id":"place-a","longitude":0,"categories":null,"open_hours":null,"popular_times":null,"review_count":0,"review_rating":0,"reviews_per_rating":null,"latitude":0,"longtitude":0,"images":null,"reservations":null,"order_online":null,"credit_cards_accepted":null,"about":null,"user_reviews":null,"user_reviews_extended":null,"emails":null,"cid":"","category":"","address":"","web_site":"","phone":"","plus_code":"","status":"","description":"","reviews_link":"","thumbnail":"","timezone":"","price_range":"","data_id":"","street_view_url":"","owner":{"id":"","name":"","link":""},"complete_address":{"borough":"","street":"","city":"","postal_code":"","state":"","country":""},"menu":{"link":"","source":""}}`, lines[0]) + require.True(t, ids.Has("place-a")) + require.True(t, ids.Has("cid-b")) +} + +func TestJSONLAppendWriterSkipsExistingIdentity(t *testing.T) { + t.Parallel() + + var out bytes.Buffer + + ids := resume.NewIdentitySet() + ids.Add("place-a") + writer := resume.NewJSONLAppendWriter(&out, ids, nil, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: []*gmaps.Entry{ + {ID: "q1", Title: "A", Link: "https://maps/place/a", PlaceID: "place-a"}, + {ID: "q1", Title: "B", Link: "https://maps/place/b", Cid: "cid-b"}, + }} + + close(in) + + require.NoError(t, writer.Run(context.Background(), in)) + + lines := strings.Split(strings.TrimSpace(out.String()), "\n") + require.Len(t, lines, 1) + require.Contains(t, lines[0], `"cid":"cid-b"`) + require.True(t, ids.Has("cid-b")) +} + +type recordingResultTracker struct { + inputIDs []string +} + +func (t *recordingResultTracker) ResultPersisted(inputID string, syncOutput func() error) error { + if syncOutput != nil { + if err := syncOutput(); err != nil { + return err + } + } + + t.inputIDs = append(t.inputIDs, inputID) + + return nil +} + +func TestJSONLAppendWriterAcknowledgesResultAfterWrite(t *testing.T) { + t.Parallel() + + var out bytes.Buffer + + tracker := &recordingResultTracker{} + writer := resume.NewJSONLAppendWriter(&out, resume.NewIdentitySet(), tracker, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: &gmaps.Entry{ID: "q1", PlaceID: "place-a"}} + + close(in) + + require.NoError(t, writer.Run(context.Background(), in)) + require.Equal(t, []string{"q1"}, tracker.inputIDs) +} + +func TestJSONLAppendWriterDoesNotAcknowledgeFailedWrite(t *testing.T) { + t.Parallel() + + tracker := &recordingResultTracker{} + writer := resume.NewJSONLAppendWriter(failingWriter{}, resume.NewIdentitySet(), tracker, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: &gmaps.Entry{ID: "q1", PlaceID: "place-a"}} + + close(in) + + require.Error(t, writer.Run(context.Background(), in)) + require.Empty(t, tracker.inputIDs) +} + +func TestJSONLAppendWriterChecksAllExistingIdentities(t *testing.T) { + t.Parallel() + + var out bytes.Buffer + + ids := resume.NewIdentitySet() + ids.Add("cid-a") + + tracker := &recordingResultTracker{} + writer := resume.NewJSONLAppendWriter(&out, ids, tracker, nil) + in := make(chan scrapemate.Result, 1) + + in <- scrapemate.Result{Data: &gmaps.Entry{ID: "q1", PlaceID: "place-a", Cid: "cid-a"}} + + close(in) + + require.NoError(t, writer.Run(context.Background(), in)) + require.Empty(t, out.String()) + require.Equal(t, []string{"q1"}, tracker.inputIDs) +} diff --git a/runner/runner.go b/runner/runner.go index 054f0d7..fe716a7 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -85,6 +85,7 @@ type Config struct { LeadsDBAPIKey string BrowserPoolSize int MaxPagesPerBrowser int + Resume bool // Grid scraping — divide a bounding box into cells to bypass the ~120 // results-per-search limit imposed by Google Maps. @@ -141,6 +142,7 @@ func ParseConfig() *Config { flag.BoolVar(&cfg.DisablePageReuse, "disable-page-reuse", false, "disable page reuse in playwright") flag.BoolVar(&cfg.ExtraReviews, "extra-reviews", false, "enable extra reviews collection") flag.StringVar(&cfg.LeadsDBAPIKey, "leadsdb-api-key", "", "LeadsDB API key for exporting results to LeadsDB") + flag.BoolVar(&cfg.Resume, "resume", false, "resume a CLI file scrape by reading existing results and appending missing places") flag.StringVar(&cfg.GridBBox, "grid-bbox", "", "bounding box for grid scraping: 'minLat,minLon,maxLat,maxLon' (e.g. '40.30,-3.80,40.50,-3.60')") flag.Float64Var(&cfg.GridCellKm, "grid-cell", 1.0, "grid cell size in km [default: 1.0]. Use with -grid-bbox") flag.IntVar(&cfg.BrowserPoolSize, "browser-pool-size", 0, "number of browser contexts for JS mode; 0 derives from concurrency and pages-per-browser")