Golang DB.Batch方法代码示例

func generateID(db *bolt.DB) ([]byte, error) {
	var id []byte
	err := db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("idgen"))
		if b == nil {
			log.Fatal(fmt.Errorf("Empty idgen"))
		}
		cid := b.Get([]byte("id"))
		if cid == nil {
			// If no ID, create ID starting at 1
			err := b.Put([]byte("id"), []byte("1"))
			if err != nil {
				return err
			}
			cid = b.Get([]byte("id"))
		} else {
			// Increment ID by 1 if it already exists
			cidInt, err := strconv.Atoi(string(cid))
			if err != nil {
				return err
			}
			cidInt = cidInt + 1
			cid = []byte(strconv.Itoa(cidInt))
			b.Put([]byte("id"), cid)
		}
		id = cid
		return nil
	})
	if err != nil {
		return nil, err
	}
	return id, nil
}

func sync(db *bolt.DB) {
	tc := transmission.New("http://127.0.0.1:9091", "", "")
	if _, err := tc.GetTorrents(); err != nil {
		fmt.Println("Could not connect to Transmission RPC API: " + err.Error())
		return
	}

	var (
		shows   []*Show
		results []*search.Result
	)

	waiting := 0
	fin := make(chan bool)
	rec := make(chan *search.Result)

	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("Shows"))
		c := b.Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			show, err := ShowFromBytes(v)
			if err == nil {
				shows = append(shows, show)
			}
		}

		return nil
	})

	for _, show := range shows {
		if show.Paused {
			continue
		}

		waiting += 1
		go syncShow(show, rec, fin)
	}

	for waiting > 0 {
		select {
		case result := <-rec:
			results = append(results, result)

		case <-fin:
			waiting -= 1
		}
	}

	for _, result := range results {
		fmt.Printf("Found torrent: '%s'\n", result.Name)

		if _, err := tc.AddTorrentByFilename(result.MagnetURL, ""); err != nil {
			fmt.Println("Error adding torrent: " + err.Error())
			fmt.Println("Stopping...")
			return
		}
	}

	fmt.Printf("Found %d torrent(s)\n", len(results))

	err := db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("Shows"))

		for _, show := range shows {
			if err := b.Put([]byte(show.Name), show.Bytes()); err != nil {
				return err
			}
		}

		return nil
	})

	if err != nil {
		fmt.Printf("Error occurred while saving shows: %s", err.Error())
	}
}

func scanCities(
	db *bolt.DB, filename string, minPopulation int,
) (int, error) {
	file, err := os.Open(filename)
	if err != nil {
		return 0, err
	}
	defer file.Close()

	reader := bufio.NewReader(file)
	scanner := bufio.NewScanner(reader)

	citiesCount := 0

	err = db.Batch(func(tx *bolt.Tx) error {
		citiesBucket := tx.Bucket(ds.CitiesBucketName)
		cityNamesBucket := tx.Bucket(ds.CityNamesBucketName)

		for scanner.Scan() {
			cityData := strings.Split(scanner.Text(), "\t")
			cityBytes, err := prepareCityBytes(cityData)
			if err != nil {
				return err
			}

			population, _ := strconv.ParseInt(cityData[14], 0, 64)
			if population > int64(minPopulation) {
				citiesBucket.Put([]byte(cityData[0]), cityBytes)

				addCityToIndex(
					cityNamesBucket, cityData[0], cityData[1], "", uint32(population),
				)

				citiesCount++
			}
		}

		return err
	})

	return citiesCount, err
}

func scanCountries(db *bolt.DB, filename string) (int, error) {
	file, err := os.Open(filename)
	if err != nil {
		return 0, err
	}
	defer file.Close()

	reader := bufio.NewReader(file)
	scanner := bufio.NewScanner(reader)

	countriesCount := 0

	err = db.Batch(func(tx *bolt.Tx) error {
		countriesBucket := tx.Bucket(ds.CountriesBucketName)

		for scanner.Scan() {
			countryString := scanner.Text()
			if strings.HasPrefix(countryString, "#") {
				continue
			}

			countryData := strings.Split(countryString, "\t")
			countryBytes, err := prepareCountryBytes(countryData)
			if err != nil {
				return err
			}

			if id := countryData[16]; id != "" {
				countriesBucket.Put([]byte(id), countryBytes)
				countriesCount++
			}
		}

		return err
	})

	return countriesCount, err
}