storage/sqlitePieceCompletion.go

   1 // +build cgo
   2
   3 package storage
   4
   5 import (
   6         "database/sql"
   7         "path/filepath"
   8
   9         _ "github.com/mattn/go-sqlite3"
  10
  11         "github.com/anacrolix/torrent/metainfo"
  12 )
  13
  14 type sqlitePieceCompletion struct {
  15         db *sql.DB
  16 }
  17
  18 var _ PieceCompletion = (*sqlitePieceCompletion)(nil)
  19
  20 func NewSqlitePieceCompletion(dir string) (ret *sqlitePieceCompletion, err error) {
  21         p := filepath.Join(dir, ".torrent.db")
  22         db, err := sql.Open("sqlite3", p)
  23         if err != nil {
  24                 return
  25         }
  26         db.SetMaxOpenConns(1)
  27         db.Exec(`PRAGMA journal_mode=WAL`)
  28         db.Exec(`PRAGMA synchronous=1`)
  29         _, err = db.Exec(`create table if not exists piece_completion(infohash, "index", complete, unique(infohash, "index"))`)
  30         if err != nil {
  31                 db.Close()
  32                 return
  33         }
  34         ret = &sqlitePieceCompletion{db}
  35         return
  36 }
  37
  38 func (me *sqlitePieceCompletion) Get(pk metainfo.PieceKey) (c Completion, err error) {
  39         row := me.db.QueryRow(`select complete from piece_completion where infohash=? and "index"=?`, pk.InfoHash.HexString(), pk.Index)
  40         err = row.Scan(&c.Complete)
  41         if err == sql.ErrNoRows {
  42                 err = nil
  43         } else if err == nil {
  44                 c.Ok = true
  45         }
  46         return
  47 }
  48
  49 func (me *sqlitePieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
  50         _, err := me.db.Exec(`insert or replace into piece_completion(infohash, "index", complete) values(?, ?, ?)`, pk.InfoHash.HexString(), pk.Index, b)
  51         return err
  52 }
  53
  54 func (me *sqlitePieceCompletion) Close() error {
  55         return me.db.Close()
  56 }