コード例 #1
0
ファイル: store.go プロジェクト: fanyang01/crawler
func New(path string, opt *bolt.Options, e codec.Codec) (bs *BoltStore, err error) {
	if e == nil {
		e = codec.JSON
	}
	bs = &BoltStore{
		filter: bloom.NewFilter(-1, 0.0001),
		codec:  e,
	}
	if bs.DB, err = bolt.Open(path, 0644, opt); err != nil {
		return nil, err
	}
	err = bs.DB.Update(func(tx *bolt.Tx) error {
		if _, err = tx.CreateBucketIfNotExists(bkURL); err != nil {
			return err
		}
		b, err := tx.CreateBucketIfNotExists(bkCount)
		if err != nil {
			return err
		}
		for _, k := range [][]byte{
			keyVisitCount, keyURLCount, keyErrorCount, keyFinishCount,
		} {
			if _, err := bkPutNX(b, k, util.I64tob(0)); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		bs = nil
	}
	return
}
コード例 #2
0
ファイル: store.go プロジェクト: fanyang01/crawler
func New(conn string) (s *SQLStore, err error) {
	db, err := sqlx.Open("postgres", conn)
	if err != nil {
		return
	}
	s = &SQLStore{
		DB:     db,
		filter: bloom.NewFilter(-1, 0.0001),
	}
	tx, err := db.Begin()
	if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			tx.Rollback()
		} else {
			err = tx.Commit()
		}
	}()

	if _, err = tx.Exec(URLSchema); err != nil {
		return nil, err
	}
	if _, err = tx.Exec(CountSchema); err != nil {
		return nil, err
	}
	var cnt int
	if err = tx.QueryRow(
		`SELECT count(*) FROM count`,
	).Scan(&cnt); err != nil {
		return
	} else if cnt == 0 {
		_, err = tx.Exec(
			`INSERT INTO count(url_count, finish_count, error_count, visit_count)
			 VALUES (0, 0, 0, 0)`,
		)
	}
	return
}