1010package mvcc
1111
1212import (
13+ "crypto/rand"
14+ "fmt"
15+ "net/url"
16+ "strings"
1317 "sync"
18+ "testing"
1419
1520 "github.com/ncruces/go-sqlite3/vfs"
21+ "github.com/ncruces/wbt"
1622)
1723
1824func init () {
@@ -26,42 +32,103 @@ var (
2632)
2733
2834// Create creates a shared memory database,
29- // using data as its initial contents.
30- func Create (name string , data string ) {
35+ // using a snapshot as its initial contents.
36+ func Create (name string , snapshot Snapshot ) {
3137 memoryMtx .Lock ()
3238 defer memoryMtx .Unlock ()
3339
34- db : = & mvccDB {
40+ memoryDBs [ name ] = & mvccDB {
3541 refs : 1 ,
3642 name : name ,
43+ data : snapshot .Tree ,
3744 }
38- memoryDBs [name ] = db
39- if len (data ) == 0 {
40- return
41- }
45+ }
46+
47+ // Delete deletes a shared memory database.
48+ func Delete (name string ) {
49+ name = getName (name )
50+
51+ memoryMtx .Lock ()
52+ defer memoryMtx .Unlock ()
53+ delete (memoryDBs , name )
54+ }
55+
56+ // Snapshot represents a database snapshot.
57+ type Snapshot struct {
58+ * wbt.Tree [int64 , string ]
59+ }
60+
61+ // NewSnapshot creates a snapshot from data.
62+ func NewSnapshot (data string ) Snapshot {
63+ var tree * wbt.Tree [int64 , string ]
64+
4265 // Convert data from WAL/2 to rollback journal.
4366 if len (data ) >= 20 && (false ||
4467 data [18 ] == 2 && data [19 ] == 2 ||
4568 data [18 ] == 3 && data [19 ] == 3 ) {
46- db . data = db . data .
69+ tree = tree .
4770 Put (0 , data [:18 ]).
4871 Put (18 , "\001 \001 " ).
4972 Put (20 , data [20 :])
50- } else {
51- db . data = db . data .Put (0 , data )
73+ } else if len ( data ) > 0 {
74+ tree = tree .Put (0 , data )
5275 }
76+
77+ return Snapshot {tree }
5378}
5479
55- // Delete deletes a shared memory database.
56- func Delete (name string ) {
80+ // TakeSnapshot takes a snapshot of a database.
81+ // Name may be a URI filename.
82+ func TakeSnapshot (name string ) Snapshot {
83+ name = getName (name )
84+
5785 memoryMtx .Lock ()
5886 defer memoryMtx .Unlock ()
59- delete (memoryDBs , name )
87+ db := memoryDBs [name ]
88+ if db == nil {
89+ return Snapshot {}
90+ }
91+
92+ db .mtx .Lock ()
93+ defer db .mtx .Unlock ()
94+ return Snapshot {db .data }
6095}
6196
62- // Snapshot stores a snapshot of database src into dst.
63- func Snapshot (dst , src string ) {
64- memoryMtx .Lock ()
65- defer memoryMtx .Unlock ()
66- memoryDBs [dst ] = memoryDBs [src ].fork ()
97+ // TestDB creates a shared database from a snapshot for the test to use.
98+ // The database is automatically deleted when the test and all its subtests complete.
99+ // Returns a URI filename appropriate to call Open with.
100+ // Each subsequent call to TestDB returns a unique database.
101+ func TestDB (tb testing.TB , snapshot Snapshot , params ... url.Values ) string {
102+ tb .Helper ()
103+
104+ name := fmt .Sprintf ("%s_%s" , tb .Name (), rand .Text ())
105+ tb .Cleanup (func () { Delete (name ) })
106+ Create (name , snapshot )
107+
108+ p := url.Values {"vfs" : {"mvcc" }}
109+ for _ , v := range params {
110+ for k , v := range v {
111+ for _ , v := range v {
112+ p .Add (k , v )
113+ }
114+ }
115+ }
116+
117+ return (& url.URL {
118+ Scheme : "file" ,
119+ OmitHost : true ,
120+ Path : "/" + name ,
121+ RawQuery : p .Encode (),
122+ }).String ()
123+ }
124+
125+ func getName (dsn string ) string {
126+ u , err := url .Parse (dsn )
127+ if err == nil &&
128+ u .Scheme == "file" &&
129+ strings .HasPrefix (u .Path , "/" ) &&
130+ u .Query ().Get ("vfs" ) == "mvcc" {
131+ return u .Path [1 :]
132+ }
133+ return dsn
67134}
0 commit comments