OLD | NEW |
1 // Copyright 2015 The Chromium Authors. All rights reserved. | 1 // Copyright 2015 The Chromium Authors. All rights reserved. |
2 // Use of this source code is governed by a BSD-style license that can be | 2 // Use of this source code is governed by a BSD-style license that can be |
3 // found in the LICENSE file. | 3 // found in the LICENSE file. |
4 | 4 |
5 package datastore | 5 package datastore |
6 | 6 |
7 import ( | 7 import ( |
8 "github.com/luci/gae/service/info" | 8 "github.com/luci/gae/service/info" |
9 "golang.org/x/net/context" | 9 "golang.org/x/net/context" |
10 ) | 10 ) |
11 | 11 |
12 type key int | 12 type key int |
13 | 13 |
14 var ( | 14 var ( |
15 rawDatastoreKey key | 15 rawDatastoreKey key |
16 rawDatastoreFilterKey key = 1 | 16 rawDatastoreFilterKey key = 1 |
17 ) | 17 ) |
18 | 18 |
19 // RawFactory is the function signature for factory methods compatible with | 19 // RawFactory is the function signature for factory methods compatible with |
20 // SetRawFactory. | 20 // SetRawFactory. wantTxn is true if the Factory should return the datastore in |
21 type RawFactory func(context.Context) RawInterface | 21 // the current transaction, and false if the Factory should return the |
| 22 // non-transactional (root) datastore. |
| 23 type RawFactory func(c context.Context, wantTxn bool) RawInterface |
22 | 24 |
23 // RawFilter is the function signature for a RawFilter implementation. It | 25 // RawFilter is the function signature for a RawFilter implementation. It |
24 // gets the current RDS implementation, and returns a new RDS implementation | 26 // gets the current RDS implementation, and returns a new RDS implementation |
25 // backed by the one passed in. | 27 // backed by the one passed in. |
26 type RawFilter func(context.Context, RawInterface) RawInterface | 28 type RawFilter func(context.Context, RawInterface) RawInterface |
27 | 29 |
28 // getUnfiltered gets gets the RawInterface implementation from context without | 30 // getUnfiltered gets gets the RawInterface implementation from context without |
29 // any of the filters applied. | 31 // any of the filters applied. |
30 func getUnfiltered(c context.Context) RawInterface { | 32 func getUnfiltered(c context.Context, wantTxn bool) RawInterface { |
31 if f, ok := c.Value(rawDatastoreKey).(RawFactory); ok && f != nil { | 33 if f, ok := c.Value(rawDatastoreKey).(RawFactory); ok && f != nil { |
32 » » return f(c) | 34 » » return f(c, wantTxn) |
33 } | 35 } |
34 return nil | 36 return nil |
35 } | 37 } |
36 | 38 |
37 // GetRaw gets the RawInterface implementation from context. | 39 // getFiltered gets the datastore (transactional or not), and applies all of |
38 func GetRaw(c context.Context) RawInterface { | 40 // the currently installed filters to it. |
39 » ret := getUnfiltered(c) | 41 func getFiltered(c context.Context, wantTxn bool) RawInterface { |
| 42 » ret := getUnfiltered(c, wantTxn) |
40 if ret == nil { | 43 if ret == nil { |
41 return nil | 44 return nil |
42 } | 45 } |
43 for _, f := range getCurFilters(c) { | 46 for _, f := range getCurFilters(c) { |
44 ret = f(c, ret) | 47 ret = f(c, ret) |
45 } | 48 } |
46 return applyCheckFilter(c, ret) | 49 return applyCheckFilter(c, ret) |
47 } | 50 } |
48 | 51 |
| 52 // GetRaw gets the RawInterface implementation from context. |
| 53 func GetRaw(c context.Context) RawInterface { |
| 54 return getFiltered(c, true) |
| 55 } |
| 56 |
| 57 // GetRawNoTxn gets the RawInterface implementation from context. If there's a |
| 58 // currently active transaction, this will return a non-transactional connection |
| 59 // to the datastore, otherwise this is the same as GetRaw. |
| 60 func GetRawNoTxn(c context.Context) RawInterface { |
| 61 return getFiltered(c, false) |
| 62 } |
| 63 |
49 // Get gets the Interface implementation from context. | 64 // Get gets the Interface implementation from context. |
50 func Get(c context.Context) Interface { | 65 func Get(c context.Context) Interface { |
51 inf := info.Get(c) | 66 inf := info.Get(c) |
52 return &datastoreImpl{ | 67 return &datastoreImpl{ |
53 GetRaw(c), | 68 GetRaw(c), |
54 inf.FullyQualifiedAppID(), | 69 inf.FullyQualifiedAppID(), |
55 inf.GetNamespace(), | 70 inf.GetNamespace(), |
56 } | 71 } |
57 } | 72 } |
58 | 73 |
| 74 // GetNoTxn gets the Interface implementation from context. If there's a |
| 75 // currently active transaction, this will return a non-transactional connection |
| 76 // to the datastore, otherwise this is the same as GetRaw. |
| 77 // Get gets the Interface implementation from context. |
| 78 func GetNoTxn(c context.Context) Interface { |
| 79 inf := info.Get(c) |
| 80 return &datastoreImpl{ |
| 81 GetRawNoTxn(c), |
| 82 inf.FullyQualifiedAppID(), |
| 83 inf.GetNamespace(), |
| 84 } |
| 85 } |
| 86 |
59 // SetRawFactory sets the function to produce Datastore instances, as returned b
y | 87 // SetRawFactory sets the function to produce Datastore instances, as returned b
y |
60 // the GetRaw method. | 88 // the GetRaw method. |
61 func SetRawFactory(c context.Context, rdsf RawFactory) context.Context { | 89 func SetRawFactory(c context.Context, rdsf RawFactory) context.Context { |
62 return context.WithValue(c, rawDatastoreKey, rdsf) | 90 return context.WithValue(c, rawDatastoreKey, rdsf) |
63 } | 91 } |
64 | 92 |
65 // SetRaw sets the current Datastore object in the context. Useful for testing | 93 // SetRaw sets the current Datastore object in the context. Useful for testing |
66 // with a quick mock. This is just a shorthand SetRawFactory invocation to set | 94 // with a quick mock. This is just a shorthand SetRawFactory invocation to set |
67 // a factory which always returns the same object. | 95 // a factory which always returns the same object. |
68 func SetRaw(c context.Context, rds RawInterface) context.Context { | 96 func SetRaw(c context.Context, rds RawInterface) context.Context { |
69 » return SetRawFactory(c, func(context.Context) RawInterface { return rds
}) | 97 » return SetRawFactory(c, func(context.Context, bool) RawInterface { retur
n rds }) |
70 } | 98 } |
71 | 99 |
72 func getCurFilters(c context.Context) []RawFilter { | 100 func getCurFilters(c context.Context) []RawFilter { |
73 curFiltsI := c.Value(rawDatastoreFilterKey) | 101 curFiltsI := c.Value(rawDatastoreFilterKey) |
74 if curFiltsI != nil { | 102 if curFiltsI != nil { |
75 return curFiltsI.([]RawFilter) | 103 return curFiltsI.([]RawFilter) |
76 } | 104 } |
77 return nil | 105 return nil |
78 } | 106 } |
79 | 107 |
80 // AddRawFilters adds RawInterface filters to the context. | 108 // AddRawFilters adds RawInterface filters to the context. |
81 func AddRawFilters(c context.Context, filts ...RawFilter) context.Context { | 109 func AddRawFilters(c context.Context, filts ...RawFilter) context.Context { |
82 if len(filts) == 0 { | 110 if len(filts) == 0 { |
83 return c | 111 return c |
84 } | 112 } |
85 cur := getCurFilters(c) | 113 cur := getCurFilters(c) |
86 newFilts := make([]RawFilter, 0, len(cur)+len(filts)) | 114 newFilts := make([]RawFilter, 0, len(cur)+len(filts)) |
87 newFilts = append(newFilts, getCurFilters(c)...) | 115 newFilts = append(newFilts, getCurFilters(c)...) |
88 newFilts = append(newFilts, filts...) | 116 newFilts = append(newFilts, filts...) |
89 return context.WithValue(c, rawDatastoreFilterKey, newFilts) | 117 return context.WithValue(c, rawDatastoreFilterKey, newFilts) |
90 } | 118 } |
OLD | NEW |