github.com/derat/nup@v0.0.0-20230418113745-15592ba7c620/cmd/nup/dump/command.go (about)

     1  // Copyright 2020 Daniel Erat.
     2  // All rights reserved.
     3  
     4  package dump
     5  
     6  import (
     7  	"bufio"
     8  	"context"
     9  	"encoding/json"
    10  	"flag"
    11  	"fmt"
    12  	"log"
    13  	"net/http"
    14  	"os"
    15  	"strings"
    16  
    17  	"github.com/derat/nup/cmd/nup/client"
    18  	"github.com/derat/nup/server/db"
    19  	"github.com/google/subcommands"
    20  )
    21  
    22  const (
    23  	progressInterval = 100
    24  
    25  	// TODO: Tune these numbers.
    26  	defaultSongBatchSize = 400
    27  	defaultPlayBatchSize = 800
    28  	chanSize             = 50
    29  )
    30  
    31  type Command struct {
    32  	Cfg *client.Config
    33  
    34  	songBatchSize int // batch size for Song entities
    35  	playBatchSize int // batch size for Play entities
    36  }
    37  
    38  func (*Command) Name() string     { return "dump" }
    39  func (*Command) Synopsis() string { return "dump songs from the server" }
    40  func (*Command) Usage() string {
    41  	return `dump <flags>:
    42  	Dump JSON-marshaled song data from the server to stdout.
    43  
    44  `
    45  }
    46  
    47  func (cmd *Command) SetFlags(f *flag.FlagSet) {
    48  	f.IntVar(&cmd.songBatchSize, "song-batch-size", defaultSongBatchSize, "Size for each batch of entities")
    49  	f.IntVar(&cmd.playBatchSize, "play-batch-size", defaultPlayBatchSize, "Size for each batch of entities")
    50  }
    51  
    52  func (cmd *Command) Execute(ctx context.Context, _ *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
    53  	songChan := make(chan *db.Song, chanSize)
    54  	go getSongs(cmd.Cfg, cmd.songBatchSize, songChan)
    55  
    56  	playChan := make(chan *db.PlayDump, chanSize)
    57  	go getPlays(cmd.Cfg, cmd.playBatchSize, playChan)
    58  
    59  	e := json.NewEncoder(os.Stdout)
    60  
    61  	numSongs := 0
    62  	pd := <-playChan
    63  	for {
    64  		s := <-songChan
    65  		if s == nil {
    66  			break
    67  		}
    68  
    69  		for pd != nil && pd.SongID == s.SongID {
    70  			s.Plays = append(s.Plays, pd.Play)
    71  			pd = <-playChan
    72  		}
    73  
    74  		if err := e.Encode(s); err != nil {
    75  			fmt.Fprintln(os.Stderr, "Failed to encode song:", err)
    76  			return subcommands.ExitFailure
    77  		}
    78  
    79  		numSongs++
    80  		if numSongs%progressInterval == 0 {
    81  			log.Printf("Wrote %d songs", numSongs)
    82  		}
    83  	}
    84  	log.Printf("Wrote %d songs", numSongs)
    85  
    86  	if pd != nil {
    87  		fmt.Fprintf(os.Stderr, "Got orphaned play for song %v: %v\n", pd.SongID, pd.Play)
    88  		return subcommands.ExitFailure
    89  	}
    90  	return subcommands.ExitSuccess
    91  }
    92  
    93  func getEntities(cfg *client.Config, entityType string, extraArgs []string, batchSize int, f func([]byte)) {
    94  	u := cfg.GetURL("/export")
    95  	var cursor string
    96  	for {
    97  		u.RawQuery = fmt.Sprintf("type=%s&max=%d", entityType, batchSize)
    98  		if len(extraArgs) > 0 {
    99  			u.RawQuery += "&" + strings.Join(extraArgs, "&")
   100  		}
   101  		if len(cursor) > 0 {
   102  			u.RawQuery += "&cursor=" + cursor
   103  		}
   104  
   105  		req, err := http.NewRequest("GET", u.String(), nil)
   106  		if err != nil {
   107  			log.Fatal("Failed to create request: ", err)
   108  		}
   109  		req.SetBasicAuth(cfg.Username, cfg.Password)
   110  
   111  		resp, err := http.DefaultClient.Do(req)
   112  		if err != nil {
   113  			log.Fatalf("Failed to fetch %v: %v", u.String(), err)
   114  		}
   115  		defer resp.Body.Close()
   116  		if resp.StatusCode != http.StatusOK {
   117  			log.Fatal("Got non-OK status: ", resp.Status)
   118  		}
   119  
   120  		cursor = ""
   121  		scanner := bufio.NewScanner(resp.Body)
   122  		for scanner.Scan() {
   123  			if err := json.Unmarshal(scanner.Bytes(), &cursor); err != nil {
   124  				f(scanner.Bytes())
   125  			}
   126  		}
   127  		if err = scanner.Err(); err != nil {
   128  			log.Fatal("Got error while reading from server: ", err)
   129  		}
   130  
   131  		if len(cursor) == 0 {
   132  			break
   133  		}
   134  	}
   135  }
   136  
   137  func getSongs(cfg *client.Config, batchSize int, ch chan *db.Song) {
   138  	getEntities(cfg, "song", nil, batchSize, func(b []byte) {
   139  		var s db.Song
   140  		if err := json.Unmarshal(b, &s); err == nil {
   141  			ch <- &s
   142  		} else {
   143  			log.Fatalf("Got unexpected line from server: %v", string(b))
   144  		}
   145  	})
   146  	ch <- nil
   147  }
   148  
   149  func getPlays(cfg *client.Config, batchSize int, ch chan *db.PlayDump) {
   150  	getEntities(cfg, "play", nil, batchSize, func(b []byte) {
   151  		var pd db.PlayDump
   152  		if err := json.Unmarshal(b, &pd); err == nil {
   153  			ch <- &pd
   154  		} else {
   155  			log.Fatalf("Got unexpected line from server: %v", string(b))
   156  		}
   157  	})
   158  	ch <- nil
   159  }