session.go 105 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517
  1. // mgo - MongoDB driver for Go
  2. //
  3. // Copyright (c) 2010-2012 - Gustavo Niemeyer <gustavo@niemeyer.net>
  4. //
  5. // All rights reserved.
  6. //
  7. // Redistribution and use in source and binary forms, with or without
  8. // modification, are permitted provided that the following conditions are met:
  9. //
  10. // 1. Redistributions of source code must retain the above copyright notice, this
  11. // list of conditions and the following disclaimer.
  12. // 2. Redistributions in binary form must reproduce the above copyright notice,
  13. // this list of conditions and the following disclaimer in the documentation
  14. // and/or other materials provided with the distribution.
  15. //
  16. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
  17. // ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  18. // WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  19. // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
  20. // ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  21. // (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  22. // LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  23. // ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  24. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  25. // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  26. package mgo
  27. import (
  28. "crypto/md5"
  29. "encoding/hex"
  30. "errors"
  31. "fmt"
  32. "labix.org/v2/mgo/bson"
  33. "math"
  34. "net"
  35. "net/url"
  36. "reflect"
  37. "sort"
  38. "strconv"
  39. "strings"
  40. "sync"
  41. "time"
  42. )
  43. type mode int
  44. const (
  45. Eventual mode = 0
  46. Monotonic mode = 1
  47. Strong mode = 2
  48. )
  49. // When changing the Session type, check if newSession and copySession
  50. // need to be updated too.
  51. type Session struct {
  52. m sync.RWMutex
  53. cluster_ *mongoCluster
  54. slaveSocket *mongoSocket
  55. masterSocket *mongoSocket
  56. slaveOk bool
  57. consistency mode
  58. queryConfig query
  59. safeOp *queryOp
  60. syncTimeout time.Duration
  61. sockTimeout time.Duration
  62. defaultdb string
  63. sourcedb string
  64. dialCred *Credential
  65. creds []Credential
  66. }
  67. type Database struct {
  68. Session *Session
  69. Name string
  70. }
  71. type Collection struct {
  72. Database *Database
  73. Name string // "collection"
  74. FullName string // "db.collection"
  75. }
  76. type Query struct {
  77. m sync.Mutex
  78. session *Session
  79. query // Enables default settings in session.
  80. }
  81. type query struct {
  82. op queryOp
  83. prefetch float64
  84. limit int32
  85. }
  86. type getLastError struct {
  87. CmdName int "getLastError"
  88. W interface{} "w,omitempty"
  89. WTimeout int "wtimeout,omitempty"
  90. FSync bool "fsync,omitempty"
  91. J bool "j,omitempty"
  92. }
  93. type Iter struct {
  94. m sync.Mutex
  95. gotReply sync.Cond
  96. session *Session
  97. server *mongoServer
  98. docData queue
  99. err error
  100. op getMoreOp
  101. prefetch float64
  102. limit int32
  103. docsToReceive int
  104. docsBeforeMore int
  105. timeout time.Duration
  106. timedout bool
  107. }
  108. var ErrNotFound = errors.New("not found")
  109. const defaultPrefetch = 0.25
  110. // Dial establishes a new session to the cluster identified by the given seed
  111. // server(s). The session will enable communication with all of the servers in
  112. // the cluster, so the seed servers are used only to find out about the cluster
  113. // topology.
  114. //
  115. // Dial will timeout after 10 seconds if a server isn't reached. The returned
  116. // session will timeout operations after one minute by default if servers
  117. // aren't available. To customize the timeout, see DialWithTimeout,
  118. // SetSyncTimeout, and SetSocketTimeout.
  119. //
  120. // This method is generally called just once for a given cluster. Further
  121. // sessions to the same cluster are then established using the New or Copy
  122. // methods on the obtained session. This will make them share the underlying
  123. // cluster, and manage the pool of connections appropriately.
  124. //
  125. // Once the session is not useful anymore, Close must be called to release the
  126. // resources appropriately.
  127. //
  128. // The seed servers must be provided in the following format:
  129. //
  130. // [mongodb://][user:pass@]host1[:port1][,host2[:port2],...][/database][?options]
  131. //
  132. // For example, it may be as simple as:
  133. //
  134. // localhost
  135. //
  136. // Or more involved like:
  137. //
  138. // mongodb://myuser:mypass@localhost:40001,otherhost:40001/mydb
  139. //
  140. // If the port number is not provided for a server, it defaults to 27017.
  141. //
  142. // The username and password provided in the URL will be used to authenticate
  143. // into the database named after the slash at the end of the host names, or
  144. // into the "admin" database if none is provided. The authentication information
  145. // will persist in sessions obtained through the New method as well.
  146. //
  147. // The following connection options are supported after the question mark:
  148. //
  149. // connect=direct
  150. //
  151. // Disables the automatic replica set server discovery logic, and
  152. // forces the use of servers provided only (even if secondaries).
  153. // Note that to talk to a secondary the consistency requirements
  154. // must be relaxed to Monotonic or Eventual via SetMode.
  155. //
  156. //
  157. // authSource=<db>
  158. //
  159. // Informs the database used to establish credentials and privileges
  160. // with a MongoDB server. Defaults to the database name provided via
  161. // the URL path, and "admin" if that's unset.
  162. //
  163. //
  164. // authMechanism=<mechanism>
  165. //
  166. // Defines the protocol for credential negotiation. Defaults to "MONGODB-CR",
  167. // which is the default username/password challenge-response mechanism.
  168. //
  169. //
  170. // gssapiServiceName=<name>
  171. //
  172. // Defines the service name to use when authenticating with the GSSAPI
  173. // mechanism. Defaults to "mongodb".
  174. //
  175. //
  176. // Relevant documentation:
  177. //
  178. // http://docs.mongodb.org/manual/reference/connection-string/
  179. //
  180. func Dial(url string) (*Session, error) {
  181. session, err := DialWithTimeout(url, 10*time.Second)
  182. if err == nil {
  183. session.SetSyncTimeout(1 * time.Minute)
  184. session.SetSocketTimeout(1 * time.Minute)
  185. }
  186. return session, err
  187. }
  188. // DialWithTimeout works like Dial, but uses timeout as the amount of time to
  189. // wait for a server to respond when first connecting and also on follow up
  190. // operations in the session. If timeout is zero, the call may block
  191. // forever waiting for a connection to be made.
  192. //
  193. // See SetSyncTimeout for customizing the timeout for the session.
  194. func DialWithTimeout(url string, timeout time.Duration) (*Session, error) {
  195. uinfo, err := parseURL(url)
  196. if err != nil {
  197. return nil, err
  198. }
  199. direct := false
  200. mechanism := ""
  201. service := ""
  202. source := ""
  203. for k, v := range uinfo.options {
  204. switch k {
  205. case "authSource":
  206. source = v
  207. case "authMechanism":
  208. mechanism = v
  209. case "gssapiServiceName":
  210. service = v
  211. case "connect":
  212. if v == "direct" {
  213. direct = true
  214. break
  215. }
  216. if v == "replicaSet" {
  217. break
  218. }
  219. fallthrough
  220. default:
  221. return nil, errors.New("unsupported connection URL option: " + k + "=" + v)
  222. }
  223. }
  224. info := DialInfo{
  225. Addrs: uinfo.addrs,
  226. Direct: direct,
  227. Timeout: timeout,
  228. Database: uinfo.db,
  229. Username: uinfo.user,
  230. Password: uinfo.pass,
  231. Mechanism: mechanism,
  232. Service: service,
  233. Source: source,
  234. }
  235. return DialWithInfo(&info)
  236. }
  237. // DialInfo holds options for establishing a session with a MongoDB cluster.
  238. // To use a URL, see the Dial function.
  239. type DialInfo struct {
  240. // Addrs holds the addresses for the seed servers.
  241. Addrs []string
  242. // Direct informs whether to establish connections only with the
  243. // specified seed servers, or to obtain information for the whole
  244. // cluster and establish connections with further servers too.
  245. Direct bool
  246. // Timeout is the amount of time to wait for a server to respond when
  247. // first connecting and on follow up operations in the session. If
  248. // timeout is zero, the call may block forever waiting for a connection
  249. // to be established.
  250. Timeout time.Duration
  251. // FailFast will cause connection and query attempts to fail faster when
  252. // the server is unavailable, instead of retrying until the configured
  253. // timeout period. Note that an unavailable server may silently drop
  254. // packets instead of rejecting them, in which case it's impossible to
  255. // distinguish it from a slow server, so the timeout stays relevant.
  256. FailFast bool
  257. // Database is the default database name used when the Session.DB method
  258. // is called with an empty name, and is also used during the intial
  259. // authenticatoin if Source is unset.
  260. Database string
  261. // Source is the database used to establish credentials and privileges
  262. // with a MongoDB server. Defaults to the value of Database, if that is
  263. // set, or "admin" otherwise.
  264. Source string
  265. // Service defines the service name to use when authenticating with the GSSAPI
  266. // mechanism. Defaults to "mongodb".
  267. Service string
  268. // Mechanism defines the protocol for credential negotiation.
  269. // Defaults to "MONGODB-CR".
  270. Mechanism string
  271. // Username and Password inform the credentials for the initial authentication
  272. // done on the database defined by the Source field. See Session.Login.
  273. Username string
  274. Password string
  275. // DialServer optionally specifies the dial function for establishing
  276. // connections with the MongoDB servers.
  277. DialServer func(addr *ServerAddr) (net.Conn, error)
  278. // WARNING: This field is obsolete. See DialServer above.
  279. Dial func(addr net.Addr) (net.Conn, error)
  280. }
  281. // ServerAddr represents the address for establishing a connection to an
  282. // individual MongoDB server.
  283. type ServerAddr struct {
  284. str string
  285. tcp *net.TCPAddr
  286. }
  287. // String returns the address that was provided for the server before resolution.
  288. func (addr *ServerAddr) String() string {
  289. return addr.str
  290. }
  291. // TCPAddr returns the resolved TCP address for the server.
  292. func (addr *ServerAddr) TCPAddr() *net.TCPAddr {
  293. return addr.tcp
  294. }
  295. // DialWithInfo establishes a new session to the cluster identified by info.
  296. func DialWithInfo(info *DialInfo) (*Session, error) {
  297. addrs := make([]string, len(info.Addrs))
  298. for i, addr := range info.Addrs {
  299. p := strings.LastIndexAny(addr, "]:")
  300. if p == -1 || addr[p] != ':' {
  301. // XXX This is untested. The test suite doesn't use the standard port.
  302. addr += ":27017"
  303. }
  304. addrs[i] = addr
  305. }
  306. cluster := newCluster(addrs, info.Direct, info.FailFast, dialer{info.Dial, info.DialServer})
  307. session := newSession(Eventual, cluster, info.Timeout)
  308. session.defaultdb = info.Database
  309. if session.defaultdb == "" {
  310. session.defaultdb = "test"
  311. }
  312. session.sourcedb = info.Source
  313. if session.sourcedb == "" {
  314. session.sourcedb = info.Database
  315. if session.sourcedb == "" {
  316. session.sourcedb = "admin"
  317. }
  318. }
  319. if info.Username != "" {
  320. source := session.sourcedb
  321. if info.Source == "" && (info.Mechanism == "GSSAPI" || info.Mechanism == "PLAIN") {
  322. source = "$external"
  323. }
  324. session.dialCred = &Credential{
  325. Username: info.Username,
  326. Password: info.Password,
  327. Mechanism: info.Mechanism,
  328. Service: info.Service,
  329. Source: source,
  330. }
  331. session.creds = []Credential{*session.dialCred}
  332. }
  333. cluster.Release()
  334. // People get confused when we return a session that is not actually
  335. // established to any servers yet (e.g. what if url was wrong). So,
  336. // ping the server to ensure there's someone there, and abort if it
  337. // fails.
  338. if err := session.Ping(); err != nil {
  339. session.Close()
  340. return nil, err
  341. }
  342. session.SetMode(Strong, true)
  343. return session, nil
  344. }
  345. func isOptSep(c rune) bool {
  346. return c == ';' || c == '&'
  347. }
  348. type urlInfo struct {
  349. addrs []string
  350. user string
  351. pass string
  352. db string
  353. options map[string]string
  354. }
  355. func parseURL(s string) (*urlInfo, error) {
  356. if strings.HasPrefix(s, "mongodb://") {
  357. s = s[10:]
  358. }
  359. info := &urlInfo{options: make(map[string]string)}
  360. if c := strings.Index(s, "?"); c != -1 {
  361. for _, pair := range strings.FieldsFunc(s[c+1:], isOptSep) {
  362. l := strings.SplitN(pair, "=", 2)
  363. if len(l) != 2 || l[0] == "" || l[1] == "" {
  364. return nil, errors.New("connection option must be key=value: " + pair)
  365. }
  366. info.options[l[0]] = l[1]
  367. }
  368. s = s[:c]
  369. }
  370. if c := strings.Index(s, "@"); c != -1 {
  371. pair := strings.SplitN(s[:c], ":", 2)
  372. if len(pair) > 2 || pair[0] == "" {
  373. return nil, errors.New("credentials must be provided as user:pass@host")
  374. }
  375. var err error
  376. info.user, err = url.QueryUnescape(pair[0])
  377. if err != nil {
  378. return nil, fmt.Errorf("cannot unescape username in URL: %q", pair[0])
  379. }
  380. if len(pair) > 1 {
  381. info.pass, err = url.QueryUnescape(pair[1])
  382. if err != nil {
  383. return nil, fmt.Errorf("cannot unescape password in URL")
  384. }
  385. }
  386. s = s[c+1:]
  387. }
  388. if c := strings.Index(s, "/"); c != -1 {
  389. info.db = s[c+1:]
  390. s = s[:c]
  391. }
  392. info.addrs = strings.Split(s, ",")
  393. return info, nil
  394. }
  395. func newSession(consistency mode, cluster *mongoCluster, timeout time.Duration) (session *Session) {
  396. cluster.Acquire()
  397. session = &Session{cluster_: cluster, syncTimeout: timeout, sockTimeout: timeout}
  398. debugf("New session %p on cluster %p", session, cluster)
  399. session.SetMode(consistency, true)
  400. session.SetSafe(&Safe{})
  401. session.queryConfig.prefetch = defaultPrefetch
  402. return session
  403. }
  404. func copySession(session *Session, keepCreds bool) (s *Session) {
  405. cluster := session.cluster()
  406. cluster.Acquire()
  407. if session.masterSocket != nil {
  408. session.masterSocket.Acquire()
  409. }
  410. if session.slaveSocket != nil {
  411. session.slaveSocket.Acquire()
  412. }
  413. var creds []Credential
  414. if keepCreds {
  415. creds = make([]Credential, len(session.creds))
  416. copy(creds, session.creds)
  417. } else if session.dialCred != nil {
  418. creds = []Credential{*session.dialCred}
  419. }
  420. scopy := *session
  421. scopy.m = sync.RWMutex{}
  422. scopy.creds = creds
  423. s = &scopy
  424. debugf("New session %p on cluster %p (copy from %p)", s, cluster, session)
  425. return s
  426. }
  427. // LiveServers returns a list of server addresses which are
  428. // currently known to be alive.
  429. func (s *Session) LiveServers() (addrs []string) {
  430. s.m.RLock()
  431. addrs = s.cluster().LiveServers()
  432. s.m.RUnlock()
  433. return addrs
  434. }
  435. // DB returns a value representing the named database. If name
  436. // is empty, the database name provided in the dialed URL is
  437. // used instead. If that is also empty, "test" is used as a
  438. // fallback in a way equivalent to the mongo shell.
  439. //
  440. // Creating this value is a very lightweight operation, and
  441. // involves no network communication.
  442. func (s *Session) DB(name string) *Database {
  443. if name == "" {
  444. name = s.defaultdb
  445. }
  446. return &Database{s, name}
  447. }
  448. // C returns a value representing the named collection.
  449. //
  450. // Creating this value is a very lightweight operation, and
  451. // involves no network communication.
  452. func (db *Database) C(name string) *Collection {
  453. return &Collection{db, name, db.Name + "." + name}
  454. }
  455. // With returns a copy of db that uses session s.
  456. func (db *Database) With(s *Session) *Database {
  457. newdb := *db
  458. newdb.Session = s
  459. return &newdb
  460. }
  461. // With returns a copy of c that uses session s.
  462. func (c *Collection) With(s *Session) *Collection {
  463. newdb := *c.Database
  464. newdb.Session = s
  465. newc := *c
  466. newc.Database = &newdb
  467. return &newc
  468. }
  469. // GridFS returns a GridFS value representing collections in db that
  470. // follow the standard GridFS specification.
  471. // The provided prefix (sometimes known as root) will determine which
  472. // collections to use, and is usually set to "fs" when there is a
  473. // single GridFS in the database.
  474. //
  475. // See the GridFS Create, Open, and OpenId methods for more details.
  476. //
  477. // Relevant documentation:
  478. //
  479. // http://www.mongodb.org/display/DOCS/GridFS
  480. // http://www.mongodb.org/display/DOCS/GridFS+Tools
  481. // http://www.mongodb.org/display/DOCS/GridFS+Specification
  482. //
  483. func (db *Database) GridFS(prefix string) *GridFS {
  484. return newGridFS(db, prefix)
  485. }
  486. // Run issues the provided command on the db database and unmarshals
  487. // its result in the respective argument. The cmd argument may be either
  488. // a string with the command name itself, in which case an empty document of
  489. // the form bson.M{cmd: 1} will be used, or it may be a full command document.
  490. //
  491. // Note that MongoDB considers the first marshalled key as the command
  492. // name, so when providing a command with options, it's important to
  493. // use an ordering-preserving document, such as a struct value or an
  494. // instance of bson.D. For instance:
  495. //
  496. // db.Run(bson.D{{"create", "mycollection"}, {"size", 1024}})
  497. //
  498. // For privilleged commands typically run on the "admin" database, see
  499. // the Run method in the Session type.
  500. //
  501. // Relevant documentation:
  502. //
  503. // http://www.mongodb.org/display/DOCS/Commands
  504. // http://www.mongodb.org/display/DOCS/List+of+Database+CommandSkips
  505. //
  506. func (db *Database) Run(cmd interface{}, result interface{}) error {
  507. if name, ok := cmd.(string); ok {
  508. cmd = bson.D{{name, 1}}
  509. }
  510. return db.C("$cmd").Find(cmd).One(result)
  511. }
  512. // Credential holds details to authenticate with a MongoDB server.
  513. type Credential struct {
  514. // Username and Password hold the basic details for authentication.
  515. // Password is optional with some authentication mechanisms.
  516. Username string
  517. Password string
  518. // Source is the database used to establish credentials and privileges
  519. // with a MongoDB server. Defaults to the default database provided
  520. // during dial, or "admin" if that was unset.
  521. Source string
  522. // Service defines the service name to use when authenticating with the GSSAPI
  523. // mechanism. Defaults to "mongodb".
  524. Service string
  525. // Mechanism defines the protocol for credential negotiation.
  526. // Defaults to "MONGODB-CR".
  527. Mechanism string
  528. }
  529. // Login authenticates with MongoDB using the provided credential. The
  530. // authentication is valid for the whole session and will stay valid until
  531. // Logout is explicitly called for the same database, or the session is
  532. // closed.
  533. func (db *Database) Login(user, pass string) error {
  534. return db.Session.Login(&Credential{Username: user, Password: pass, Source: db.Name})
  535. }
  536. // Login authenticates with MongoDB using the provided credential. The
  537. // authentication is valid for the whole session and will stay valid until
  538. // Logout is explicitly called for the same database, or the session is
  539. // closed.
  540. func (s *Session) Login(cred *Credential) error {
  541. socket, err := s.acquireSocket(true)
  542. if err != nil {
  543. return err
  544. }
  545. defer socket.Release()
  546. credCopy := *cred
  547. if cred.Source == "" {
  548. if cred.Mechanism == "GSSAPI" {
  549. credCopy.Source = "$external"
  550. } else {
  551. credCopy.Source = s.sourcedb
  552. }
  553. }
  554. err = socket.Login(credCopy)
  555. if err != nil {
  556. return err
  557. }
  558. s.m.Lock()
  559. s.creds = append(s.creds, credCopy)
  560. s.m.Unlock()
  561. return nil
  562. }
  563. func (s *Session) socketLogin(socket *mongoSocket) error {
  564. for _, cred := range s.creds {
  565. if err := socket.Login(cred); err != nil {
  566. return err
  567. }
  568. }
  569. return nil
  570. }
  571. // Logout removes any established authentication credentials for the database.
  572. func (db *Database) Logout() {
  573. session := db.Session
  574. dbname := db.Name
  575. session.m.Lock()
  576. found := false
  577. for i, cred := range session.creds {
  578. if cred.Source == dbname {
  579. copy(session.creds[i:], session.creds[i+1:])
  580. session.creds = session.creds[:len(session.creds)-1]
  581. found = true
  582. break
  583. }
  584. }
  585. if found {
  586. if session.masterSocket != nil {
  587. session.masterSocket.Logout(dbname)
  588. }
  589. if session.slaveSocket != nil {
  590. session.slaveSocket.Logout(dbname)
  591. }
  592. }
  593. session.m.Unlock()
  594. }
  595. // LogoutAll removes all established authentication credentials for the session.
  596. func (s *Session) LogoutAll() {
  597. s.m.Lock()
  598. for _, cred := range s.creds {
  599. if s.masterSocket != nil {
  600. s.masterSocket.Logout(cred.Source)
  601. }
  602. if s.slaveSocket != nil {
  603. s.slaveSocket.Logout(cred.Source)
  604. }
  605. }
  606. s.creds = s.creds[0:0]
  607. s.m.Unlock()
  608. }
  609. // User represents a MongoDB user.
  610. //
  611. // Relevant documentation:
  612. //
  613. // http://docs.mongodb.org/manual/reference/privilege-documents/
  614. // http://docs.mongodb.org/manual/reference/user-privileges/
  615. //
  616. type User struct {
  617. // Username is how the user identifies itself to the system.
  618. Username string `bson:"user"`
  619. // Password is the plaintext password for the user. If set,
  620. // the UpsertUser method will hash it into PasswordHash and
  621. // unset it before the user is added to the database.
  622. Password string `bson:",omitempty"`
  623. // PasswordHash is the MD5 hash of Username+":mongo:"+Password.
  624. PasswordHash string `bson:"pwd,omitempty"`
  625. // CustomData holds arbitrary data admins decide to associate
  626. // with this user, such as the full name or employee id.
  627. CustomData interface{} `bson:"customData,omitempty"`
  628. // Roles indicates the set of roles the user will be provided.
  629. // See the Role constants.
  630. Roles []Role `bson:"roles"`
  631. // OtherDBRoles allows assigning roles in other databases from
  632. // user documents inserted in the admin database. This field
  633. // only works in the admin database.
  634. OtherDBRoles map[string][]Role `bson:"otherDBRoles,omitempty"`
  635. // UserSource indicates where to look for this user's credentials.
  636. // It may be set to a database name, or to "$external" for
  637. // consulting an external resource such as Kerberos. UserSource
  638. // must not be set if Password or PasswordHash are present.
  639. //
  640. // WARNING: This setting was only ever supported in MongoDB 2.4,
  641. // and is now obsolete.
  642. UserSource string `bson:"userSource,omitempty"`
  643. }
  644. type Role string
  645. const (
  646. // Relevant documentation:
  647. //
  648. // http://docs.mongodb.org/manual/reference/user-privileges/
  649. //
  650. RoleRoot Role = "root"
  651. RoleRead Role = "read"
  652. RoleReadAny Role = "readAnyDatabase"
  653. RoleReadWrite Role = "readWrite"
  654. RoleReadWriteAny Role = "readWriteAnyDatabase"
  655. RoleDBAdmin Role = "dbAdmin"
  656. RoleDBAdminAny Role = "dbAdminAnyDatabase"
  657. RoleUserAdmin Role = "userAdmin"
  658. RoleUserAdminAny Role = "userAdminAnyDatabase"
  659. RoleClusterAdmin Role = "clusterAdmin"
  660. )
  661. // UpsertUser updates the authentication credentials and the roles for
  662. // a MongoDB user within the db database. If the named user doesn't exist
  663. // it will be created.
  664. //
  665. // This method should only be used from MongoDB 2.4 and on. For older
  666. // MongoDB releases, use the obsolete AddUser method instead.
  667. //
  668. // Relevant documentation:
  669. //
  670. // http://docs.mongodb.org/manual/reference/user-privileges/
  671. // http://docs.mongodb.org/manual/reference/privilege-documents/
  672. //
  673. func (db *Database) UpsertUser(user *User) error {
  674. if user.Username == "" {
  675. return fmt.Errorf("user has no Username")
  676. }
  677. if (user.Password != "" || user.PasswordHash != "") && user.UserSource != "" {
  678. return fmt.Errorf("user has both Password/PasswordHash and UserSource set")
  679. }
  680. if len(user.OtherDBRoles) > 0 && db.Name != "admin" {
  681. return fmt.Errorf("user with OtherDBRoles is only supported in admin database")
  682. }
  683. // Attempt to run this using 2.6+ commands.
  684. rundb := db
  685. if user.UserSource != "" {
  686. // Compatibility logic for the userSource field of MongoDB <= 2.4.X
  687. rundb = db.Session.DB(user.UserSource)
  688. }
  689. err := rundb.runUserCmd("updateUser", user)
  690. if e, ok := err.(*QueryError); ok && e.Code == 11 {
  691. return rundb.runUserCmd("createUser", user)
  692. }
  693. if !isNoCmd(err) {
  694. return err
  695. }
  696. // Command does not exist. Fallback to pre-2.6 behavior.
  697. var set, unset bson.D
  698. if user.Password != "" {
  699. psum := md5.New()
  700. psum.Write([]byte(user.Username + ":mongo:" + user.Password))
  701. set = append(set, bson.DocElem{"pwd", hex.EncodeToString(psum.Sum(nil))})
  702. unset = append(unset, bson.DocElem{"userSource", 1})
  703. } else if user.PasswordHash != "" {
  704. set = append(set, bson.DocElem{"pwd", user.PasswordHash})
  705. unset = append(unset, bson.DocElem{"userSource", 1})
  706. }
  707. if user.UserSource != "" {
  708. set = append(set, bson.DocElem{"userSource", user.UserSource})
  709. unset = append(unset, bson.DocElem{"pwd", 1})
  710. }
  711. if user.Roles != nil || user.OtherDBRoles != nil {
  712. set = append(set, bson.DocElem{"roles", user.Roles})
  713. if len(user.OtherDBRoles) > 0 {
  714. set = append(set, bson.DocElem{"otherDBRoles", user.OtherDBRoles})
  715. } else {
  716. unset = append(unset, bson.DocElem{"otherDBRoles", 1})
  717. }
  718. }
  719. users := db.C("system.users")
  720. err = users.Update(bson.D{{"user", user.Username}}, bson.D{{"$unset", unset}, {"$set", set}})
  721. if err == ErrNotFound {
  722. set = append(set, bson.DocElem{"user", user.Username})
  723. if user.Roles == nil && user.OtherDBRoles == nil {
  724. // Roles must be sent, as it's the way MongoDB distinguishes
  725. // old-style documents from new-style documents in pre-2.6.
  726. set = append(set, bson.DocElem{"roles", user.Roles})
  727. }
  728. err = users.Insert(set)
  729. }
  730. return err
  731. }
  732. func isNoCmd(err error) bool {
  733. e, ok := err.(*QueryError)
  734. return ok && strings.HasPrefix(e.Message, "no such cmd:")
  735. }
  736. func (db *Database) runUserCmd(cmdName string, user *User) error {
  737. //if user.UserSource != "" && (user.UserSource != "$external" || db.Name != "$external") {
  738. // return fmt.Errorf("MongoDB 2.6+ does not support the UserSource setting")
  739. //}
  740. cmd := make(bson.D, 0, 16)
  741. cmd = append(cmd, bson.DocElem{cmdName, user.Username})
  742. if user.Password != "" {
  743. cmd = append(cmd, bson.DocElem{"pwd", user.Password})
  744. }
  745. var roles []interface{}
  746. for _, role := range user.Roles {
  747. roles = append(roles, role)
  748. }
  749. for db, dbroles := range user.OtherDBRoles {
  750. for _, role := range dbroles {
  751. roles = append(roles, bson.D{{"role", role}, {"db", db}})
  752. }
  753. }
  754. if roles != nil || user.Roles != nil || cmdName == "createUser" {
  755. cmd = append(cmd, bson.DocElem{"roles", roles})
  756. }
  757. return db.Run(cmd, nil)
  758. }
  759. // AddUser creates or updates the authentication credentials of user within
  760. // the db database.
  761. //
  762. // WARNING: This method is obsolete and should only be used with MongoDB 2.2
  763. // or earlier. For MongoDB 2.4 and on, use UpsertUser instead.
  764. func (db *Database) AddUser(username, password string, readOnly bool) error {
  765. // Try to emulate the old behavior on 2.6+
  766. user := &User{Username: username, Password: password}
  767. if db.Name == "admin" {
  768. if readOnly {
  769. user.Roles = []Role{RoleReadAny}
  770. } else {
  771. user.Roles = []Role{RoleReadWriteAny}
  772. }
  773. } else {
  774. if readOnly {
  775. user.Roles = []Role{RoleRead}
  776. } else {
  777. user.Roles = []Role{RoleReadWrite}
  778. }
  779. }
  780. err := db.runUserCmd("updateUser", user)
  781. if e, ok := err.(*QueryError); ok && e.Code == 11 {
  782. return db.runUserCmd("createUser", user)
  783. }
  784. if !isNoCmd(err) {
  785. return err
  786. }
  787. // Command doesn't exist. Fallback to pre-2.6 behavior.
  788. psum := md5.New()
  789. psum.Write([]byte(username + ":mongo:" + password))
  790. digest := hex.EncodeToString(psum.Sum(nil))
  791. c := db.C("system.users")
  792. _, err = c.Upsert(bson.M{"user": username}, bson.M{"$set": bson.M{"user": username, "pwd": digest, "readOnly": readOnly}})
  793. return err
  794. }
  795. // RemoveUser removes the authentication credentials of user from the database.
  796. func (db *Database) RemoveUser(user string) error {
  797. err := db.Run(bson.D{{"dropUser", user}}, nil)
  798. if isNoCmd(err) {
  799. users := db.C("system.users")
  800. return users.Remove(bson.M{"user": user})
  801. }
  802. return err
  803. }
  804. type indexSpec struct {
  805. Name, NS string
  806. Key bson.D
  807. Unique bool ",omitempty"
  808. DropDups bool "dropDups,omitempty"
  809. Background bool ",omitempty"
  810. Sparse bool ",omitempty"
  811. Bits, Min, Max int ",omitempty"
  812. ExpireAfter int "expireAfterSeconds,omitempty"
  813. }
  814. type Index struct {
  815. Key []string // Index key fields; prefix name with dash (-) for descending order
  816. Unique bool // Prevent two documents from having the same index key
  817. DropDups bool // Drop documents with the same index key as a previously indexed one
  818. Background bool // Build index in background and return immediately
  819. Sparse bool // Only index documents containing the Key fields
  820. ExpireAfter time.Duration // Periodically delete docs with indexed time.Time older than that.
  821. Name string // Index name, computed by EnsureIndex
  822. Bits, Min, Max int // Properties for spatial indexes
  823. }
  824. func parseIndexKey(key []string) (name string, realKey bson.D, err error) {
  825. var order interface{}
  826. for _, field := range key {
  827. raw := field
  828. if name != "" {
  829. name += "_"
  830. }
  831. var kind string
  832. if field != "" {
  833. if field[0] == '$' {
  834. if c := strings.Index(field, ":"); c > 1 && c < len(field)-1 {
  835. kind = field[1:c]
  836. field = field[c+1:]
  837. name += field + "_" + kind
  838. }
  839. }
  840. switch field[0] {
  841. case '$':
  842. // Logic above failed. Reset and error.
  843. field = ""
  844. case '@':
  845. order = "2d"
  846. field = field[1:]
  847. // The shell used to render this field as key_ instead of key_2d,
  848. // and mgo followed suit. This has been fixed in recent server
  849. // releases, and mgo followed as well.
  850. name += field + "_2d"
  851. case '-':
  852. order = -1
  853. field = field[1:]
  854. name += field + "_-1"
  855. case '+':
  856. field = field[1:]
  857. fallthrough
  858. default:
  859. if kind == "" {
  860. order = 1
  861. name += field + "_1"
  862. } else {
  863. order = kind
  864. }
  865. }
  866. }
  867. if field == "" || kind != "" && order != kind {
  868. return "", nil, fmt.Errorf(`invalid index key: want "[$<kind>:][-]<field name>", got %q`, raw)
  869. }
  870. realKey = append(realKey, bson.DocElem{field, order})
  871. }
  872. if name == "" {
  873. return "", nil, errors.New("invalid index key: no fields provided")
  874. }
  875. return
  876. }
  877. // EnsureIndexKey ensures an index with the given key exists, creating it
  878. // if necessary.
  879. //
  880. // This example:
  881. //
  882. // err := collection.EnsureIndexKey("a", "b")
  883. //
  884. // Is equivalent to:
  885. //
  886. // err := collection.EnsureIndex(mgo.Index{Key: []string{"a", "b"}})
  887. //
  888. // See the EnsureIndex method for more details.
  889. func (c *Collection) EnsureIndexKey(key ...string) error {
  890. return c.EnsureIndex(Index{Key: key})
  891. }
  892. // EnsureIndex ensures an index with the given key exists, creating it with
  893. // the provided parameters if necessary.
  894. //
  895. // Once EnsureIndex returns successfully, following requests for the same index
  896. // will not contact the server unless Collection.DropIndex is used to drop the
  897. // same index, or Session.ResetIndexCache is called.
  898. //
  899. // For example:
  900. //
  901. // index := Index{
  902. // Key: []string{"lastname", "firstname"},
  903. // Unique: true,
  904. // DropDups: true,
  905. // Background: true, // See notes.
  906. // Sparse: true,
  907. // }
  908. // err := collection.EnsureIndex(index)
  909. //
  910. // The Key value determines which fields compose the index. The index ordering
  911. // will be ascending by default. To obtain an index with a descending order,
  912. // the field name should be prefixed by a dash (e.g. []string{"-time"}).
  913. //
  914. // If Unique is true, the index must necessarily contain only a single
  915. // document per Key. With DropDups set to true, documents with the same key
  916. // as a previously indexed one will be dropped rather than an error returned.
  917. //
  918. // If Background is true, other connections will be allowed to proceed using
  919. // the collection without the index while it's being built. Note that the
  920. // session executing EnsureIndex will be blocked for as long as it takes for
  921. // the index to be built.
  922. //
  923. // If Sparse is true, only documents containing the provided Key fields will be
  924. // included in the index. When using a sparse index for sorting, only indexed
  925. // documents will be returned.
  926. //
  927. // If ExpireAfter is non-zero, the server will periodically scan the collection
  928. // and remove documents containing an indexed time.Time field with a value
  929. // older than ExpireAfter. See the documentation for details:
  930. //
  931. // http://docs.mongodb.org/manual/tutorial/expire-data
  932. //
  933. // Other kinds of indexes are also supported through that API. Here is an example:
  934. //
  935. // index := Index{
  936. // Key: []string{"$2d:loc"},
  937. // Bits: 26,
  938. // }
  939. // err := collection.EnsureIndex(index)
  940. //
  941. // The example above requests the creation of a "2d" index for the "loc" field.
  942. //
  943. // The 2D index bounds may be changed using the Min and Max attributes of the
  944. // Index value. The default bound setting of (-180, 180) is suitable for
  945. // latitude/longitude pairs.
  946. //
  947. // The Bits parameter sets the precision of the 2D geohash values. If not
  948. // provided, 26 bits are used, which is roughly equivalent to 1 foot of
  949. // precision for the default (-180, 180) index bounds.
  950. //
  951. // Relevant documentation:
  952. //
  953. // http://www.mongodb.org/display/DOCS/Indexes
  954. // http://www.mongodb.org/display/DOCS/Indexing+Advice+and+FAQ
  955. // http://www.mongodb.org/display/DOCS/Indexing+as+a+Background+Operation
  956. // http://www.mongodb.org/display/DOCS/Geospatial+Indexing
  957. // http://www.mongodb.org/display/DOCS/Multikeys
  958. //
  959. func (c *Collection) EnsureIndex(index Index) error {
  960. name, realKey, err := parseIndexKey(index.Key)
  961. if err != nil {
  962. return err
  963. }
  964. session := c.Database.Session
  965. cacheKey := c.FullName + "\x00" + name
  966. if session.cluster().HasCachedIndex(cacheKey) {
  967. return nil
  968. }
  969. spec := indexSpec{
  970. Name: name,
  971. NS: c.FullName,
  972. Key: realKey,
  973. Unique: index.Unique,
  974. DropDups: index.DropDups,
  975. Background: index.Background,
  976. Sparse: index.Sparse,
  977. Bits: index.Bits,
  978. Min: index.Min,
  979. Max: index.Max,
  980. ExpireAfter: int(index.ExpireAfter / time.Second),
  981. }
  982. session = session.Clone()
  983. defer session.Close()
  984. session.SetMode(Strong, false)
  985. session.EnsureSafe(&Safe{})
  986. db := c.Database.With(session)
  987. err = db.C("system.indexes").Insert(&spec)
  988. if err == nil {
  989. session.cluster().CacheIndex(cacheKey, true)
  990. }
  991. session.Close()
  992. return err
  993. }
  994. // DropIndex removes the index with key from the collection.
  995. //
  996. // The key value determines which fields compose the index. The index ordering
  997. // will be ascending by default. To obtain an index with a descending order,
  998. // the field name should be prefixed by a dash (e.g. []string{"-time"}).
  999. //
  1000. // For example:
  1001. //
  1002. // err := collection.DropIndex("lastname", "firstname")
  1003. //
  1004. // See the EnsureIndex method for more details on indexes.
  1005. func (c *Collection) DropIndex(key ...string) error {
  1006. name, _, err := parseIndexKey(key)
  1007. if err != nil {
  1008. return err
  1009. }
  1010. session := c.Database.Session
  1011. cacheKey := c.FullName + "\x00" + name
  1012. session.cluster().CacheIndex(cacheKey, false)
  1013. session = session.Clone()
  1014. defer session.Close()
  1015. session.SetMode(Strong, false)
  1016. db := c.Database.With(session)
  1017. result := struct {
  1018. ErrMsg string
  1019. Ok bool
  1020. }{}
  1021. err = db.Run(bson.D{{"dropIndexes", c.Name}, {"index", name}}, &result)
  1022. if err != nil {
  1023. return err
  1024. }
  1025. if !result.Ok {
  1026. return errors.New(result.ErrMsg)
  1027. }
  1028. return nil
  1029. }
  1030. // Indexes returns a list of all indexes for the collection.
  1031. //
  1032. // For example, this snippet would drop all available indexes:
  1033. //
  1034. // indexes, err := collection.Indexes()
  1035. // if err != nil {
  1036. // return err
  1037. // }
  1038. // for _, index := range indexes {
  1039. // err = collection.DropIndex(index.Key...)
  1040. // if err != nil {
  1041. // return err
  1042. // }
  1043. // }
  1044. //
  1045. // See the EnsureIndex method for more details on indexes.
  1046. func (c *Collection) Indexes() (indexes []Index, err error) {
  1047. query := c.Database.C("system.indexes").Find(bson.M{"ns": c.FullName})
  1048. iter := query.Sort("name").Iter()
  1049. for {
  1050. var spec indexSpec
  1051. if !iter.Next(&spec) {
  1052. break
  1053. }
  1054. index := Index{
  1055. Name: spec.Name,
  1056. Key: simpleIndexKey(spec.Key),
  1057. Unique: spec.Unique,
  1058. DropDups: spec.DropDups,
  1059. Background: spec.Background,
  1060. Sparse: spec.Sparse,
  1061. ExpireAfter: time.Duration(spec.ExpireAfter) * time.Second,
  1062. }
  1063. indexes = append(indexes, index)
  1064. }
  1065. err = iter.Close()
  1066. return
  1067. }
  1068. func simpleIndexKey(realKey bson.D) (key []string) {
  1069. for i := range realKey {
  1070. field := realKey[i].Name
  1071. vi, ok := realKey[i].Value.(int)
  1072. if !ok {
  1073. vf, _ := realKey[i].Value.(float64)
  1074. vi = int(vf)
  1075. }
  1076. if vi == 1 {
  1077. key = append(key, field)
  1078. continue
  1079. }
  1080. if vi == -1 {
  1081. key = append(key, "-"+field)
  1082. continue
  1083. }
  1084. if vs, ok := realKey[i].Value.(string); ok {
  1085. key = append(key, "$"+vs+":"+field)
  1086. continue
  1087. }
  1088. panic("Got unknown index key type for field " + field)
  1089. }
  1090. return
  1091. }
  1092. // ResetIndexCache() clears the cache of previously ensured indexes.
  1093. // Following requests to EnsureIndex will contact the server.
  1094. func (s *Session) ResetIndexCache() {
  1095. s.cluster().ResetIndexCache()
  1096. }
  1097. // New creates a new session with the same parameters as the original
  1098. // session, including consistency, batch size, prefetching, safety mode,
  1099. // etc. The returned session will use sockets from the pool, so there's
  1100. // a chance that writes just performed in another session may not yet
  1101. // be visible.
  1102. //
  1103. // Login information from the original session will not be copied over
  1104. // into the new session unless it was provided through the initial URL
  1105. // for the Dial function.
  1106. //
  1107. // See the Copy and Clone methods.
  1108. //
  1109. func (s *Session) New() *Session {
  1110. s.m.Lock()
  1111. scopy := copySession(s, false)
  1112. s.m.Unlock()
  1113. scopy.Refresh()
  1114. return scopy
  1115. }
  1116. // Copy works just like New, but preserves the exact authentication
  1117. // information from the original session.
  1118. func (s *Session) Copy() *Session {
  1119. s.m.Lock()
  1120. scopy := copySession(s, true)
  1121. s.m.Unlock()
  1122. scopy.Refresh()
  1123. return scopy
  1124. }
  1125. // Clone works just like Copy, but also reuses the same socket as the original
  1126. // session, in case it had already reserved one due to its consistency
  1127. // guarantees. This behavior ensures that writes performed in the old session
  1128. // are necessarily observed when using the new session, as long as it was a
  1129. // strong or monotonic session. That said, it also means that long operations
  1130. // may cause other goroutines using the original session to wait.
  1131. func (s *Session) Clone() *Session {
  1132. s.m.Lock()
  1133. scopy := copySession(s, true)
  1134. s.m.Unlock()
  1135. return scopy
  1136. }
  1137. // Close terminates the session. It's a runtime error to use a session
  1138. // after it has been closed.
  1139. func (s *Session) Close() {
  1140. s.m.Lock()
  1141. if s.cluster_ != nil {
  1142. debugf("Closing session %p", s)
  1143. s.unsetSocket()
  1144. s.cluster_.Release()
  1145. s.cluster_ = nil
  1146. }
  1147. s.m.Unlock()
  1148. }
  1149. func (s *Session) cluster() *mongoCluster {
  1150. if s.cluster_ == nil {
  1151. panic("Session already closed")
  1152. }
  1153. return s.cluster_
  1154. }
  1155. // Refresh puts back any reserved sockets in use and restarts the consistency
  1156. // guarantees according to the current consistency setting for the session.
  1157. func (s *Session) Refresh() {
  1158. s.m.Lock()
  1159. s.slaveOk = s.consistency != Strong
  1160. s.unsetSocket()
  1161. s.m.Unlock()
  1162. }
  1163. // SetMode changes the consistency mode for the session.
  1164. //
  1165. // In the Strong consistency mode reads and writes will always be made to
  1166. // the primary server using a unique connection so that reads and writes are
  1167. // fully consistent, ordered, and observing the most up-to-date data.
  1168. // This offers the least benefits in terms of distributing load, but the
  1169. // most guarantees. See also Monotonic and Eventual.
  1170. //
  1171. // In the Monotonic consistency mode reads may not be entirely up-to-date,
  1172. // but they will always see the history of changes moving forward, the data
  1173. // read will be consistent across sequential queries in the same session,
  1174. // and modifications made within the session will be observed in following
  1175. // queries (read-your-writes).
  1176. //
  1177. // In practice, the Monotonic mode is obtained by performing initial reads
  1178. // on a unique connection to an arbitrary secondary, if one is available,
  1179. // and once the first write happens, the session connection is switched over
  1180. // to the primary server. This manages to distribute some of the reading
  1181. // load with secondaries, while maintaining some useful guarantees.
  1182. //
  1183. // In the Eventual consistency mode reads will be made to any secondary in the
  1184. // cluster, if one is available, and sequential reads will not necessarily
  1185. // be made with the same connection. This means that data may be observed
  1186. // out of order. Writes will of course be issued to the primary, but
  1187. // independent writes in the same Eventual session may also be made with
  1188. // independent connections, so there are also no guarantees in terms of
  1189. // write ordering (no read-your-writes guarantees either).
  1190. //
  1191. // The Eventual mode is the fastest and most resource-friendly, but is
  1192. // also the one offering the least guarantees about ordering of the data
  1193. // read and written.
  1194. //
  1195. // If refresh is true, in addition to ensuring the session is in the given
  1196. // consistency mode, the consistency guarantees will also be reset (e.g.
  1197. // a Monotonic session will be allowed to read from secondaries again).
  1198. // This is equivalent to calling the Refresh function.
  1199. //
  1200. // Shifting between Monotonic and Strong modes will keep a previously
  1201. // reserved connection for the session unless refresh is true or the
  1202. // connection is unsuitable (to a secondary server in a Strong session).
  1203. func (s *Session) SetMode(consistency mode, refresh bool) {
  1204. s.m.Lock()
  1205. debugf("Session %p: setting mode %d with refresh=%v (master=%p, slave=%p)", s, consistency, refresh, s.masterSocket, s.slaveSocket)
  1206. s.consistency = consistency
  1207. if refresh {
  1208. s.slaveOk = s.consistency != Strong
  1209. s.unsetSocket()
  1210. } else if s.consistency == Strong {
  1211. s.slaveOk = false
  1212. } else if s.masterSocket == nil {
  1213. s.slaveOk = true
  1214. }
  1215. s.m.Unlock()
  1216. }
  1217. // Mode returns the current consistency mode for the session.
  1218. func (s *Session) Mode() mode {
  1219. s.m.RLock()
  1220. mode := s.consistency
  1221. s.m.RUnlock()
  1222. return mode
  1223. }
  1224. // SetSyncTimeout sets the amount of time an operation with this session
  1225. // will wait before returning an error in case a connection to a usable
  1226. // server can't be established. Set it to zero to wait forever. The
  1227. // default value is 7 seconds.
  1228. func (s *Session) SetSyncTimeout(d time.Duration) {
  1229. s.m.Lock()
  1230. s.syncTimeout = d
  1231. s.m.Unlock()
  1232. }
  1233. // SetSocketTimeout sets the amount of time to wait for a non-responding
  1234. // socket to the database before it is forcefully closed.
  1235. func (s *Session) SetSocketTimeout(d time.Duration) {
  1236. s.m.Lock()
  1237. s.sockTimeout = d
  1238. if s.masterSocket != nil {
  1239. s.masterSocket.SetTimeout(d)
  1240. }
  1241. if s.slaveSocket != nil {
  1242. s.slaveSocket.SetTimeout(d)
  1243. }
  1244. s.m.Unlock()
  1245. }
  1246. // SetCursorTimeout changes the standard timeout period that the server
  1247. // enforces on created cursors. The only supported value right now is
  1248. // 0, which disables the timeout. The standard server timeout is 10 minutes.
  1249. func (s *Session) SetCursorTimeout(d time.Duration) {
  1250. s.m.Lock()
  1251. if d == 0 {
  1252. s.queryConfig.op.flags |= flagNoCursorTimeout
  1253. } else {
  1254. panic("SetCursorTimeout: only 0 (disable timeout) supported for now")
  1255. }
  1256. s.m.Unlock()
  1257. }
  1258. // SetBatch sets the default batch size used when fetching documents from the
  1259. // database. It's possible to change this setting on a per-query basis as
  1260. // well, using the Query.Batch method.
  1261. //
  1262. // The default batch size is defined by the database itself. As of this
  1263. // writing, MongoDB will use an initial size of min(100 docs, 4MB) on the
  1264. // first batch, and 4MB on remaining ones.
  1265. func (s *Session) SetBatch(n int) {
  1266. if n == 1 {
  1267. // Server interprets 1 as -1 and closes the cursor (!?)
  1268. n = 2
  1269. }
  1270. s.m.Lock()
  1271. s.queryConfig.op.limit = int32(n)
  1272. s.m.Unlock()
  1273. }
  1274. // SetPrefetch sets the default point at which the next batch of results will be
  1275. // requested. When there are p*batch_size remaining documents cached in an
  1276. // Iter, the next batch will be requested in background. For instance, when
  1277. // using this:
  1278. //
  1279. // session.SetBatch(200)
  1280. // session.SetPrefetch(0.25)
  1281. //
  1282. // and there are only 50 documents cached in the Iter to be processed, the
  1283. // next batch of 200 will be requested. It's possible to change this setting on
  1284. // a per-query basis as well, using the Prefetch method of Query.
  1285. //
  1286. // The default prefetch value is 0.25.
  1287. func (s *Session) SetPrefetch(p float64) {
  1288. s.m.Lock()
  1289. s.queryConfig.prefetch = p
  1290. s.m.Unlock()
  1291. }
  1292. // See SetSafe for details on the Safe type.
  1293. type Safe struct {
  1294. W int // Min # of servers to ack before success
  1295. WMode string // Write mode for MongoDB 2.0+ (e.g. "majority")
  1296. WTimeout int // Milliseconds to wait for W before timing out
  1297. FSync bool // Should servers sync to disk before returning success
  1298. J bool // Wait for next group commit if journaling; no effect otherwise
  1299. }
  1300. // Safe returns the current safety mode for the session.
  1301. func (s *Session) Safe() (safe *Safe) {
  1302. s.m.Lock()
  1303. defer s.m.Unlock()
  1304. if s.safeOp != nil {
  1305. cmd := s.safeOp.query.(*getLastError)
  1306. safe = &Safe{WTimeout: cmd.WTimeout, FSync: cmd.FSync, J: cmd.J}
  1307. switch w := cmd.W.(type) {
  1308. case string:
  1309. safe.WMode = w
  1310. case int:
  1311. safe.W = w
  1312. }
  1313. }
  1314. return
  1315. }
  1316. // SetSafe changes the session safety mode.
  1317. //
  1318. // If the safe parameter is nil, the session is put in unsafe mode, and writes
  1319. // become fire-and-forget, without error checking. The unsafe mode is faster
  1320. // since operations won't hold on waiting for a confirmation.
  1321. //
  1322. // If the safe parameter is not nil, any changing query (insert, update, ...)
  1323. // will be followed by a getLastError command with the specified parameters,
  1324. // to ensure the request was correctly processed.
  1325. //
  1326. // The safe.W parameter determines how many servers should confirm a write
  1327. // before the operation is considered successful. If set to 0 or 1, the
  1328. // command will return as soon as the primary is done with the request.
  1329. // If safe.WTimeout is greater than zero, it determines how many milliseconds
  1330. // to wait for the safe.W servers to respond before returning an error.
  1331. //
  1332. // Starting with MongoDB 2.0.0 the safe.WMode parameter can be used instead
  1333. // of W to request for richer semantics. If set to "majority" the server will
  1334. // wait for a majority of members from the replica set to respond before
  1335. // returning. Custom modes may also be defined within the server to create
  1336. // very detailed placement schemas. See the data awareness documentation in
  1337. // the links below for more details (note that MongoDB internally reuses the
  1338. // "w" field name for WMode).
  1339. //
  1340. // If safe.FSync is true and journaling is disabled, the servers will be
  1341. // forced to sync all files to disk immediately before returning. If the
  1342. // same option is true but journaling is enabled, the server will instead
  1343. // await for the next group commit before returning.
  1344. //
  1345. // Since MongoDB 2.0.0, the safe.J option can also be used instead of FSync
  1346. // to force the server to wait for a group commit in case journaling is
  1347. // enabled. The option has no effect if the server has journaling disabled.
  1348. //
  1349. // For example, the following statement will make the session check for
  1350. // errors, without imposing further constraints:
  1351. //
  1352. // session.SetSafe(&mgo.Safe{})
  1353. //
  1354. // The following statement will force the server to wait for a majority of
  1355. // members of a replica set to return (MongoDB 2.0+ only):
  1356. //
  1357. // session.SetSafe(&mgo.Safe{WMode: "majority"})
  1358. //
  1359. // The following statement, on the other hand, ensures that at least two
  1360. // servers have flushed the change to disk before confirming the success
  1361. // of operations:
  1362. //
  1363. // session.EnsureSafe(&mgo.Safe{W: 2, FSync: true})
  1364. //
  1365. // The following statement, on the other hand, disables the verification
  1366. // of errors entirely:
  1367. //
  1368. // session.SetSafe(nil)
  1369. //
  1370. // See also the EnsureSafe method.
  1371. //
  1372. // Relevant documentation:
  1373. //
  1374. // http://www.mongodb.org/display/DOCS/getLastError+Command
  1375. // http://www.mongodb.org/display/DOCS/Verifying+Propagation+of+Writes+with+getLastError
  1376. // http://www.mongodb.org/display/DOCS/Data+Center+Awareness
  1377. //
  1378. func (s *Session) SetSafe(safe *Safe) {
  1379. s.m.Lock()
  1380. s.safeOp = nil
  1381. s.ensureSafe(safe)
  1382. s.m.Unlock()
  1383. }
  1384. // EnsureSafe compares the provided safety parameters with the ones
  1385. // currently in use by the session and picks the most conservative
  1386. // choice for each setting.
  1387. //
  1388. // That is:
  1389. //
  1390. // - safe.WMode is always used if set.
  1391. // - safe.W is used if larger than the current W and WMode is empty.
  1392. // - safe.FSync is always used if true.
  1393. // - safe.J is used if FSync is false.
  1394. // - safe.WTimeout is used if set and smaller than the current WTimeout.
  1395. //
  1396. // For example, the following statement will ensure the session is
  1397. // at least checking for errors, without enforcing further constraints.
  1398. // If a more conservative SetSafe or EnsureSafe call was previously done,
  1399. // the following call will be ignored.
  1400. //
  1401. // session.EnsureSafe(&mgo.Safe{})
  1402. //
  1403. // See also the SetSafe method for details on what each option means.
  1404. //
  1405. // Relevant documentation:
  1406. //
  1407. // http://www.mongodb.org/display/DOCS/getLastError+Command
  1408. // http://www.mongodb.org/display/DOCS/Verifying+Propagation+of+Writes+with+getLastError
  1409. // http://www.mongodb.org/display/DOCS/Data+Center+Awareness
  1410. //
  1411. func (s *Session) EnsureSafe(safe *Safe) {
  1412. s.m.Lock()
  1413. s.ensureSafe(safe)
  1414. s.m.Unlock()
  1415. }
  1416. func (s *Session) ensureSafe(safe *Safe) {
  1417. if safe == nil {
  1418. return
  1419. }
  1420. var w interface{}
  1421. if safe.WMode != "" {
  1422. w = safe.WMode
  1423. } else if safe.W > 0 {
  1424. w = safe.W
  1425. }
  1426. var cmd getLastError
  1427. if s.safeOp == nil {
  1428. cmd = getLastError{1, w, safe.WTimeout, safe.FSync, safe.J}
  1429. } else {
  1430. // Copy. We don't want to mutate the existing query.
  1431. cmd = *(s.safeOp.query.(*getLastError))
  1432. if cmd.W == nil {
  1433. cmd.W = w
  1434. } else if safe.WMode != "" {
  1435. cmd.W = safe.WMode
  1436. } else if i, ok := cmd.W.(int); ok && safe.W > i {
  1437. cmd.W = safe.W
  1438. }
  1439. if safe.WTimeout > 0 && safe.WTimeout < cmd.WTimeout {
  1440. cmd.WTimeout = safe.WTimeout
  1441. }
  1442. if safe.FSync {
  1443. cmd.FSync = true
  1444. cmd.J = false
  1445. } else if safe.J && !cmd.FSync {
  1446. cmd.J = true
  1447. }
  1448. }
  1449. s.safeOp = &queryOp{
  1450. query: &cmd,
  1451. collection: "admin.$cmd",
  1452. limit: -1,
  1453. }
  1454. }
  1455. // Run issues the provided command on the "admin" database and
  1456. // and unmarshals its result in the respective argument. The cmd
  1457. // argument may be either a string with the command name itself, in
  1458. // which case an empty document of the form bson.M{cmd: 1} will be used,
  1459. // or it may be a full command document.
  1460. //
  1461. // Note that MongoDB considers the first marshalled key as the command
  1462. // name, so when providing a command with options, it's important to
  1463. // use an ordering-preserving document, such as a struct value or an
  1464. // instance of bson.D. For instance:
  1465. //
  1466. // db.Run(bson.D{{"create", "mycollection"}, {"size", 1024}})
  1467. //
  1468. // For commands on arbitrary databases, see the Run method in
  1469. // the Database type.
  1470. //
  1471. // Relevant documentation:
  1472. //
  1473. // http://www.mongodb.org/display/DOCS/Commands
  1474. // http://www.mongodb.org/display/DOCS/List+of+Database+CommandSkips
  1475. //
  1476. func (s *Session) Run(cmd interface{}, result interface{}) error {
  1477. return s.DB("admin").Run(cmd, result)
  1478. }
  1479. // SelectServers restricts communication to servers configured with the
  1480. // given tags. For example, the following statement restricts servers
  1481. // used for reading operations to those with both tag "disk" set to
  1482. // "ssd" and tag "rack" set to 1:
  1483. //
  1484. // session.SelectSlaves(bson.D{{"disk", "ssd"}, {"rack", 1}})
  1485. //
  1486. // Multiple sets of tags may be provided, in which case the used server
  1487. // must match all tags within any one set.
  1488. //
  1489. // If a connection was previously assigned to the session due to the
  1490. // current session mode (see Session.SetMode), the tag selection will
  1491. // only be enforced after the session is refreshed.
  1492. //
  1493. // Relevant documentation:
  1494. //
  1495. // http://docs.mongodb.org/manual/tutorial/configure-replica-set-tag-sets
  1496. //
  1497. func (s *Session) SelectServers(tags ...bson.D) {
  1498. s.m.Lock()
  1499. s.queryConfig.op.serverTags = tags
  1500. s.m.Unlock()
  1501. }
  1502. // Ping runs a trivial ping command just to get in touch with the server.
  1503. func (s *Session) Ping() error {
  1504. return s.Run("ping", nil)
  1505. }
  1506. // Fsync flushes in-memory writes to disk on the server the session
  1507. // is established with. If async is true, the call returns immediately,
  1508. // otherwise it returns after the flush has been made.
  1509. func (s *Session) Fsync(async bool) error {
  1510. return s.Run(bson.D{{"fsync", 1}, {"async", async}}, nil)
  1511. }
  1512. // FsyncLock locks all writes in the specific server the session is
  1513. // established with and returns. Any writes attempted to the server
  1514. // after it is successfully locked will block until FsyncUnlock is
  1515. // called for the same server.
  1516. //
  1517. // This method works on secondaries as well, preventing the oplog from
  1518. // being flushed while the server is locked, but since only the server
  1519. // connected to is locked, for locking specific secondaries it may be
  1520. // necessary to establish a connection directly to the secondary (see
  1521. // Dial's connect=direct option).
  1522. //
  1523. // As an important caveat, note that once a write is attempted and
  1524. // blocks, follow up reads will block as well due to the way the
  1525. // lock is internally implemented in the server. More details at:
  1526. //
  1527. // https://jira.mongodb.org/browse/SERVER-4243
  1528. //
  1529. // FsyncLock is often used for performing consistent backups of
  1530. // the database files on disk.
  1531. //
  1532. // Relevant documentation:
  1533. //
  1534. // http://www.mongodb.org/display/DOCS/fsync+Command
  1535. // http://www.mongodb.org/display/DOCS/Backups
  1536. //
  1537. func (s *Session) FsyncLock() error {
  1538. return s.Run(bson.D{{"fsync", 1}, {"lock", true}}, nil)
  1539. }
  1540. // FsyncUnlock releases the server for writes. See FsyncLock for details.
  1541. func (s *Session) FsyncUnlock() error {
  1542. return s.DB("admin").C("$cmd.sys.unlock").Find(nil).One(nil) // WTF?
  1543. }
  1544. // Find prepares a query using the provided document. The document may be a
  1545. // map or a struct value capable of being marshalled with bson. The map
  1546. // may be a generic one using interface{} for its key and/or values, such as
  1547. // bson.M, or it may be a properly typed map. Providing nil as the document
  1548. // is equivalent to providing an empty document such as bson.M{}.
  1549. //
  1550. // Further details of the query may be tweaked using the resulting Query value,
  1551. // and then executed to retrieve results using methods such as One, For,
  1552. // Iter, or Tail.
  1553. //
  1554. // In case the resulting document includes a field named $err or errmsg, which
  1555. // are standard ways for MongoDB to return query errors, the returned err will
  1556. // be set to a *QueryError value including the Err message and the Code. In
  1557. // those cases, the result argument is still unmarshalled into with the
  1558. // received document so that any other custom values may be obtained if
  1559. // desired.
  1560. //
  1561. // Relevant documentation:
  1562. //
  1563. // http://www.mongodb.org/display/DOCS/Querying
  1564. // http://www.mongodb.org/display/DOCS/Advanced+Queries
  1565. //
  1566. func (c *Collection) Find(query interface{}) *Query {
  1567. session := c.Database.Session
  1568. session.m.RLock()
  1569. q := &Query{session: session, query: session.queryConfig}
  1570. session.m.RUnlock()
  1571. q.op.query = query
  1572. q.op.collection = c.FullName
  1573. return q
  1574. }
  1575. // FindId is a convenience helper equivalent to:
  1576. //
  1577. // query := collection.Find(bson.M{"_id": id})
  1578. //
  1579. // See the Find method for more details.
  1580. func (c *Collection) FindId(id interface{}) *Query {
  1581. return c.Find(bson.D{{"_id", id}})
  1582. }
  1583. type Pipe struct {
  1584. session *Session
  1585. collection *Collection
  1586. pipeline interface{}
  1587. }
  1588. // Pipe prepares a pipeline to aggregate. The pipeline document
  1589. // must be a slice built in terms of the aggregation framework language.
  1590. //
  1591. // For example:
  1592. //
  1593. // pipe := collection.Pipe([]bson.M{{"$match": bson.M{"name": "Otavio"}}})
  1594. // iter := pipe.Iter()
  1595. //
  1596. // Relevant documentation:
  1597. //
  1598. // http://docs.mongodb.org/manual/reference/aggregation
  1599. // http://docs.mongodb.org/manual/applications/aggregation
  1600. // http://docs.mongodb.org/manual/tutorial/aggregation-examples
  1601. //
  1602. func (c *Collection) Pipe(pipeline interface{}) *Pipe {
  1603. session := c.Database.Session
  1604. return &Pipe{
  1605. session: session,
  1606. collection: c,
  1607. pipeline: pipeline,
  1608. }
  1609. }
  1610. // Iter executes the pipeline and returns an iterator capable of going
  1611. // over all the generated results.
  1612. func (p *Pipe) Iter() *Iter {
  1613. iter := &Iter{
  1614. session: p.session,
  1615. timeout: -1,
  1616. }
  1617. iter.gotReply.L = &iter.m
  1618. var result struct{ Result []bson.Raw }
  1619. c := p.collection
  1620. iter.err = c.Database.Run(bson.D{{"aggregate", c.Name}, {"pipeline", p.pipeline}}, &result)
  1621. if iter.err != nil {
  1622. return iter
  1623. }
  1624. for i := range result.Result {
  1625. iter.docData.Push(result.Result[i].Data)
  1626. }
  1627. return iter
  1628. }
  1629. // All works like Iter.All.
  1630. func (p *Pipe) All(result interface{}) error {
  1631. return p.Iter().All(result)
  1632. }
  1633. // One executes the pipeline and unmarshals the first item from the
  1634. // result set into the result parameter.
  1635. // It returns ErrNotFound if no items are generated by the pipeline.
  1636. func (p *Pipe) One(result interface{}) error {
  1637. iter := p.Iter()
  1638. if iter.Next(result) {
  1639. return nil
  1640. }
  1641. if err := iter.Err(); err != nil {
  1642. return err
  1643. }
  1644. return ErrNotFound
  1645. }
  1646. type LastError struct {
  1647. Err string
  1648. Code, N, Waited int
  1649. FSyncFiles int `bson:"fsyncFiles"`
  1650. WTimeout bool
  1651. UpdatedExisting bool `bson:"updatedExisting"`
  1652. UpsertedId interface{} `bson:"upserted"`
  1653. }
  1654. func (err *LastError) Error() string {
  1655. return err.Err
  1656. }
  1657. type queryError struct {
  1658. Err string "$err"
  1659. ErrMsg string
  1660. Assertion string
  1661. Code int
  1662. AssertionCode int "assertionCode"
  1663. LastError *LastError "lastErrorObject"
  1664. }
  1665. type QueryError struct {
  1666. Code int
  1667. Message string
  1668. Assertion bool
  1669. }
  1670. func (err *QueryError) Error() string {
  1671. return err.Message
  1672. }
  1673. // IsDup returns whether err informs of a duplicate key error because
  1674. // a primary key index or a secondary unique index already has an entry
  1675. // with the given value.
  1676. func IsDup(err error) bool {
  1677. // Besides being handy, helps with MongoDB bugs SERVER-7164 and SERVER-11493.
  1678. // What follows makes me sad. Hopefully conventions will be more clear over time.
  1679. switch e := err.(type) {
  1680. case *LastError:
  1681. return e.Code == 11000 || e.Code == 11001 || e.Code == 12582 || e.Code == 16460 && strings.Contains(e.Err, " E11000 ")
  1682. case *QueryError:
  1683. return e.Code == 11000 || e.Code == 11001 || e.Code == 12582
  1684. }
  1685. return false
  1686. }
  1687. // Insert inserts one or more documents in the respective collection. In
  1688. // case the session is in safe mode (see the SetSafe method) and an error
  1689. // happens while inserting the provided documents, the returned error will
  1690. // be of type *LastError.
  1691. func (c *Collection) Insert(docs ...interface{}) error {
  1692. _, err := c.writeQuery(&insertOp{c.FullName, docs, 0})
  1693. return err
  1694. }
  1695. // Update finds a single document matching the provided selector document
  1696. // and modifies it according to the update document.
  1697. // If the session is in safe mode (see SetSafe) a ErrNotFound error is
  1698. // returned if a document isn't found, or a value of type *LastError
  1699. // when some other error is detected.
  1700. //
  1701. // Relevant documentation:
  1702. //
  1703. // http://www.mongodb.org/display/DOCS/Updating
  1704. // http://www.mongodb.org/display/DOCS/Atomic+Operations
  1705. //
  1706. func (c *Collection) Update(selector interface{}, update interface{}) error {
  1707. lerr, err := c.writeQuery(&updateOp{c.FullName, selector, update, 0})
  1708. if err == nil && lerr != nil && !lerr.UpdatedExisting {
  1709. return ErrNotFound
  1710. }
  1711. return err
  1712. }
  1713. // UpdateId is a convenience helper equivalent to:
  1714. //
  1715. // err := collection.Update(bson.M{"_id": id}, update)
  1716. //
  1717. // See the Update method for more details.
  1718. func (c *Collection) UpdateId(id interface{}, update interface{}) error {
  1719. return c.Update(bson.D{{"_id", id}}, update)
  1720. }
  1721. // ChangeInfo holds details about the outcome of an update operation.
  1722. type ChangeInfo struct {
  1723. Updated int // Number of existing documents updated
  1724. Removed int // Number of documents removed
  1725. UpsertedId interface{} // Upserted _id field, when not explicitly provided
  1726. }
  1727. // UpdateAll finds all documents matching the provided selector document
  1728. // and modifies them according to the update document.
  1729. // If the session is in safe mode (see SetSafe) details of the executed
  1730. // operation are returned in info or an error of type *LastError when
  1731. // some problem is detected. It is not an error for the update to not be
  1732. // applied on any documents because the selector doesn't match.
  1733. //
  1734. // Relevant documentation:
  1735. //
  1736. // http://www.mongodb.org/display/DOCS/Updating
  1737. // http://www.mongodb.org/display/DOCS/Atomic+Operations
  1738. //
  1739. func (c *Collection) UpdateAll(selector interface{}, update interface{}) (info *ChangeInfo, err error) {
  1740. lerr, err := c.writeQuery(&updateOp{c.FullName, selector, update, 2})
  1741. if err == nil && lerr != nil {
  1742. info = &ChangeInfo{Updated: lerr.N}
  1743. }
  1744. return info, err
  1745. }
  1746. // Upsert finds a single document matching the provided selector document
  1747. // and modifies it according to the update document. If no document matching
  1748. // the selector is found, the update document is applied to the selector
  1749. // document and the result is inserted in the collection.
  1750. // If the session is in safe mode (see SetSafe) details of the executed
  1751. // operation are returned in info, or an error of type *LastError when
  1752. // some problem is detected.
  1753. //
  1754. // Relevant documentation:
  1755. //
  1756. // http://www.mongodb.org/display/DOCS/Updating
  1757. // http://www.mongodb.org/display/DOCS/Atomic+Operations
  1758. //
  1759. func (c *Collection) Upsert(selector interface{}, update interface{}) (info *ChangeInfo, err error) {
  1760. lerr, err := c.writeQuery(&updateOp{c.FullName, selector, update, 1})
  1761. if err == nil && lerr != nil {
  1762. info = &ChangeInfo{}
  1763. if lerr.UpdatedExisting {
  1764. info.Updated = lerr.N
  1765. } else {
  1766. info.UpsertedId = lerr.UpsertedId
  1767. }
  1768. }
  1769. return info, err
  1770. }
  1771. // UpsertId is a convenience helper equivalent to:
  1772. //
  1773. // info, err := collection.Upsert(bson.M{"_id": id}, update)
  1774. //
  1775. // See the Upsert method for more details.
  1776. func (c *Collection) UpsertId(id interface{}, update interface{}) (info *ChangeInfo, err error) {
  1777. return c.Upsert(bson.D{{"_id", id}}, update)
  1778. }
  1779. // Remove finds a single document matching the provided selector document
  1780. // and removes it from the database.
  1781. // If the session is in safe mode (see SetSafe) a ErrNotFound error is
  1782. // returned if a document isn't found, or a value of type *LastError
  1783. // when some other error is detected.
  1784. //
  1785. // Relevant documentation:
  1786. //
  1787. // http://www.mongodb.org/display/DOCS/Removing
  1788. //
  1789. func (c *Collection) Remove(selector interface{}) error {
  1790. lerr, err := c.writeQuery(&deleteOp{c.FullName, selector, 1})
  1791. if err == nil && lerr != nil && lerr.N == 0 {
  1792. return ErrNotFound
  1793. }
  1794. return err
  1795. }
  1796. // RemoveId is a convenience helper equivalent to:
  1797. //
  1798. // err := collection.Remove(bson.M{"_id": id})
  1799. //
  1800. // See the Remove method for more details.
  1801. func (c *Collection) RemoveId(id interface{}) error {
  1802. return c.Remove(bson.D{{"_id", id}})
  1803. }
  1804. // RemoveAll finds all documents matching the provided selector document
  1805. // and removes them from the database. In case the session is in safe mode
  1806. // (see the SetSafe method) and an error happens when attempting the change,
  1807. // the returned error will be of type *LastError.
  1808. //
  1809. // Relevant documentation:
  1810. //
  1811. // http://www.mongodb.org/display/DOCS/Removing
  1812. //
  1813. func (c *Collection) RemoveAll(selector interface{}) (info *ChangeInfo, err error) {
  1814. lerr, err := c.writeQuery(&deleteOp{c.FullName, selector, 0})
  1815. if err == nil && lerr != nil {
  1816. info = &ChangeInfo{Removed: lerr.N}
  1817. }
  1818. return info, err
  1819. }
  1820. // DropDatabase removes the entire database including all of its collections.
  1821. func (db *Database) DropDatabase() error {
  1822. return db.Run(bson.D{{"dropDatabase", 1}}, nil)
  1823. }
  1824. // DropCollection removes the entire collection including all of its documents.
  1825. func (c *Collection) DropCollection() error {
  1826. return c.Database.Run(bson.D{{"drop", c.Name}}, nil)
  1827. }
  1828. // The CollectionInfo type holds metadata about a collection.
  1829. //
  1830. // Relevant documentation:
  1831. //
  1832. // http://www.mongodb.org/display/DOCS/createCollection+Command
  1833. // http://www.mongodb.org/display/DOCS/Capped+Collections
  1834. //
  1835. type CollectionInfo struct {
  1836. // DisableIdIndex prevents the automatic creation of the index
  1837. // on the _id field for the collection.
  1838. DisableIdIndex bool
  1839. // ForceIdIndex enforces the automatic creation of the index
  1840. // on the _id field for the collection. Capped collections,
  1841. // for example, do not have such an index by default.
  1842. ForceIdIndex bool
  1843. // If Capped is true new documents will replace old ones when
  1844. // the collection is full. MaxBytes must necessarily be set
  1845. // to define the size when the collection wraps around.
  1846. // MaxDocs optionally defines the number of documents when it
  1847. // wraps, but MaxBytes still needs to be set.
  1848. Capped bool
  1849. MaxBytes int
  1850. MaxDocs int
  1851. }
  1852. // Create explicitly creates the c collection with details of info.
  1853. // MongoDB creates collections automatically on use, so this method
  1854. // is only necessary when creating collection with non-default
  1855. // characteristics, such as capped collections.
  1856. //
  1857. // Relevant documentation:
  1858. //
  1859. // http://www.mongodb.org/display/DOCS/createCollection+Command
  1860. // http://www.mongodb.org/display/DOCS/Capped+Collections
  1861. //
  1862. func (c *Collection) Create(info *CollectionInfo) error {
  1863. cmd := make(bson.D, 0, 4)
  1864. cmd = append(cmd, bson.DocElem{"create", c.Name})
  1865. if info.Capped {
  1866. if info.MaxBytes < 1 {
  1867. return fmt.Errorf("Collection.Create: with Capped, MaxBytes must also be set")
  1868. }
  1869. cmd = append(cmd, bson.DocElem{"capped", true})
  1870. cmd = append(cmd, bson.DocElem{"size", info.MaxBytes})
  1871. if info.MaxDocs > 0 {
  1872. cmd = append(cmd, bson.DocElem{"max", info.MaxDocs})
  1873. }
  1874. }
  1875. if info.DisableIdIndex {
  1876. cmd = append(cmd, bson.DocElem{"autoIndexId", false})
  1877. }
  1878. if info.ForceIdIndex {
  1879. cmd = append(cmd, bson.DocElem{"autoIndexId", true})
  1880. }
  1881. return c.Database.Run(cmd, nil)
  1882. }
  1883. // Batch sets the batch size used when fetching documents from the database.
  1884. // It's possible to change this setting on a per-session basis as well, using
  1885. // the Batch method of Session.
  1886. //
  1887. // The default batch size is defined by the database itself. As of this
  1888. // writing, MongoDB will use an initial size of min(100 docs, 4MB) on the
  1889. // first batch, and 4MB on remaining ones.
  1890. func (q *Query) Batch(n int) *Query {
  1891. if n == 1 {
  1892. // Server interprets 1 as -1 and closes the cursor (!?)
  1893. n = 2
  1894. }
  1895. q.m.Lock()
  1896. q.op.limit = int32(n)
  1897. q.m.Unlock()
  1898. return q
  1899. }
  1900. // Prefetch sets the point at which the next batch of results will be requested.
  1901. // When there are p*batch_size remaining documents cached in an Iter, the next
  1902. // batch will be requested in background. For instance, when using this:
  1903. //
  1904. // query.Batch(200).Prefetch(0.25)
  1905. //
  1906. // and there are only 50 documents cached in the Iter to be processed, the
  1907. // next batch of 200 will be requested. It's possible to change this setting on
  1908. // a per-session basis as well, using the SetPrefetch method of Session.
  1909. //
  1910. // The default prefetch value is 0.25.
  1911. func (q *Query) Prefetch(p float64) *Query {
  1912. q.m.Lock()
  1913. q.prefetch = p
  1914. q.m.Unlock()
  1915. return q
  1916. }
  1917. // Skip skips over the n initial documents from the query results. Note that
  1918. // this only makes sense with capped collections where documents are naturally
  1919. // ordered by insertion time, or with sorted results.
  1920. func (q *Query) Skip(n int) *Query {
  1921. q.m.Lock()
  1922. q.op.skip = int32(n)
  1923. q.m.Unlock()
  1924. return q
  1925. }
  1926. // Limit restricts the maximum number of documents retrieved to n, and also
  1927. // changes the batch size to the same value. Once n documents have been
  1928. // returned by Next, the following call will return ErrNotFound.
  1929. func (q *Query) Limit(n int) *Query {
  1930. q.m.Lock()
  1931. switch {
  1932. case n == 1:
  1933. q.limit = 1
  1934. q.op.limit = -1
  1935. case n == math.MinInt32: // -MinInt32 == -MinInt32
  1936. q.limit = math.MaxInt32
  1937. q.op.limit = math.MinInt32 + 1
  1938. case n < 0:
  1939. q.limit = int32(-n)
  1940. q.op.limit = int32(n)
  1941. default:
  1942. q.limit = int32(n)
  1943. q.op.limit = int32(n)
  1944. }
  1945. q.m.Unlock()
  1946. return q
  1947. }
  1948. // Select enables selecting which fields should be retrieved for the results
  1949. // found. For example, the following query would only retrieve the name field:
  1950. //
  1951. // err := collection.Find(nil).Select(bson.M{"name": 1}).One(&result)
  1952. //
  1953. // Relevant documentation:
  1954. //
  1955. // http://www.mongodb.org/display/DOCS/Retrieving+a+Subset+of+Fields
  1956. //
  1957. func (q *Query) Select(selector interface{}) *Query {
  1958. q.m.Lock()
  1959. q.op.selector = selector
  1960. q.m.Unlock()
  1961. return q
  1962. }
  1963. // Sort asks the database to order returned documents according to the
  1964. // provided field names. A field name may be prefixed by - (minus) for
  1965. // it to be sorted in reverse order.
  1966. //
  1967. // For example:
  1968. //
  1969. // query1 := collection.Find(nil).Sort("firstname", "lastname")
  1970. // query2 := collection.Find(nil).Sort("-age")
  1971. // query3 := collection.Find(nil).Sort("$natural")
  1972. //
  1973. // Relevant documentation:
  1974. //
  1975. // http://www.mongodb.org/display/DOCS/Sorting+and+Natural+Order
  1976. //
  1977. func (q *Query) Sort(fields ...string) *Query {
  1978. // TODO // query4 := collection.Find(nil).Sort("score:{$meta:textScore}")
  1979. q.m.Lock()
  1980. var order bson.D
  1981. for _, field := range fields {
  1982. n := 1
  1983. if field != "" {
  1984. switch field[0] {
  1985. case '+':
  1986. field = field[1:]
  1987. case '-':
  1988. n = -1
  1989. field = field[1:]
  1990. }
  1991. }
  1992. if field == "" {
  1993. panic("Sort: empty field name")
  1994. }
  1995. order = append(order, bson.DocElem{field, n})
  1996. }
  1997. q.op.options.OrderBy = order
  1998. q.op.hasOptions = true
  1999. q.m.Unlock()
  2000. return q
  2001. }
  2002. // Explain returns a number of details about how the MongoDB server would
  2003. // execute the requested query, such as the number of objects examined,
  2004. // the number of time the read lock was yielded to allow writes to go in,
  2005. // and so on.
  2006. //
  2007. // For example:
  2008. //
  2009. // m := bson.M{}
  2010. // err := collection.Find(bson.M{"filename": name}).Explain(m)
  2011. // if err == nil {
  2012. // fmt.Printf("Explain: %#v\n", m)
  2013. // }
  2014. //
  2015. // Relevant documentation:
  2016. //
  2017. // http://www.mongodb.org/display/DOCS/Optimization
  2018. // http://www.mongodb.org/display/DOCS/Query+Optimizer
  2019. //
  2020. func (q *Query) Explain(result interface{}) error {
  2021. q.m.Lock()
  2022. clone := &Query{session: q.session, query: q.query}
  2023. q.m.Unlock()
  2024. clone.op.options.Explain = true
  2025. clone.op.hasOptions = true
  2026. if clone.op.limit > 0 {
  2027. clone.op.limit = -q.op.limit
  2028. }
  2029. iter := clone.Iter()
  2030. if iter.Next(result) {
  2031. return nil
  2032. }
  2033. return iter.Close()
  2034. }
  2035. // Hint will include an explicit "hint" in the query to force the server
  2036. // to use a specified index, potentially improving performance in some
  2037. // situations. The provided parameters are the fields that compose the
  2038. // key of the index to be used. For details on how the indexKey may be
  2039. // built, see the EnsureIndex method.
  2040. //
  2041. // For example:
  2042. //
  2043. // query := collection.Find(bson.M{"firstname": "Joe", "lastname": "Winter"})
  2044. // query.Hint("lastname", "firstname")
  2045. //
  2046. // Relevant documentation:
  2047. //
  2048. // http://www.mongodb.org/display/DOCS/Optimization
  2049. // http://www.mongodb.org/display/DOCS/Query+Optimizer
  2050. //
  2051. func (q *Query) Hint(indexKey ...string) *Query {
  2052. q.m.Lock()
  2053. _, realKey, err := parseIndexKey(indexKey)
  2054. q.op.options.Hint = realKey
  2055. q.op.hasOptions = true
  2056. q.m.Unlock()
  2057. if err != nil {
  2058. panic(err)
  2059. }
  2060. return q
  2061. }
  2062. // Snapshot will force the performed query to make use of an available
  2063. // index on the _id field to prevent the same document from being returned
  2064. // more than once in a single iteration. This might happen without this
  2065. // setting in situations when the document changes in size and thus has to
  2066. // be moved while the iteration is running.
  2067. //
  2068. // Because snapshot mode traverses the _id index, it may not be used with
  2069. // sorting or explicit hints. It also cannot use any other index for the
  2070. // query.
  2071. //
  2072. // Even with snapshot mode, items inserted or deleted during the query may
  2073. // or may not be returned; that is, this mode is not a true point-in-time
  2074. // snapshot.
  2075. //
  2076. // The same effect of Snapshot may be obtained by using any unique index on
  2077. // field(s) that will not be modified (best to use Hint explicitly too).
  2078. // A non-unique index (such as creation time) may be made unique by
  2079. // appending _id to the index when creating it.
  2080. //
  2081. // Relevant documentation:
  2082. //
  2083. // http://www.mongodb.org/display/DOCS/How+to+do+Snapshotted+Queries+in+the+Mongo+Database
  2084. //
  2085. func (q *Query) Snapshot() *Query {
  2086. q.m.Lock()
  2087. q.op.options.Snapshot = true
  2088. q.op.hasOptions = true
  2089. q.m.Unlock()
  2090. return q
  2091. }
  2092. // LogReplay enables an option that optimizes queries that are typically
  2093. // made on the MongoDB oplog for replaying it. This is an internal
  2094. // implementation aspect and most likely uninteresting for other uses.
  2095. // It has seen at least one use case, though, so it's exposed via the API.
  2096. func (q *Query) LogReplay() *Query {
  2097. q.m.Lock()
  2098. q.op.flags |= flagLogReplay
  2099. q.m.Unlock()
  2100. return q
  2101. }
  2102. func checkQueryError(fullname string, d []byte) error {
  2103. l := len(d)
  2104. if l < 16 {
  2105. return nil
  2106. }
  2107. if d[5] == '$' && d[6] == 'e' && d[7] == 'r' && d[8] == 'r' && d[9] == '\x00' && d[4] == '\x02' {
  2108. goto Error
  2109. }
  2110. if len(fullname) < 5 || fullname[len(fullname)-5:] != ".$cmd" {
  2111. return nil
  2112. }
  2113. for i := 0; i+8 < l; i++ {
  2114. if d[i] == '\x02' && d[i+1] == 'e' && d[i+2] == 'r' && d[i+3] == 'r' && d[i+4] == 'm' && d[i+5] == 's' && d[i+6] == 'g' && d[i+7] == '\x00' {
  2115. goto Error
  2116. }
  2117. }
  2118. return nil
  2119. Error:
  2120. result := &queryError{}
  2121. bson.Unmarshal(d, result)
  2122. logf("queryError: %#v\n", result)
  2123. if result.LastError != nil {
  2124. return result.LastError
  2125. }
  2126. if result.Err == "" && result.ErrMsg == "" {
  2127. return nil
  2128. }
  2129. if result.AssertionCode != 0 && result.Assertion != "" {
  2130. return &QueryError{Code: result.AssertionCode, Message: result.Assertion, Assertion: true}
  2131. }
  2132. if result.Err != "" {
  2133. return &QueryError{Code: result.Code, Message: result.Err}
  2134. }
  2135. return &QueryError{Code: result.Code, Message: result.ErrMsg}
  2136. }
  2137. // One executes the query and unmarshals the first obtained document into the
  2138. // result argument. The result must be a struct or map value capable of being
  2139. // unmarshalled into by gobson. This function blocks until either a result
  2140. // is available or an error happens. For example:
  2141. //
  2142. // err := collection.Find(bson.M{"a", 1}).One(&result)
  2143. //
  2144. // In case the resulting document includes a field named $err or errmsg, which
  2145. // are standard ways for MongoDB to return query errors, the returned err will
  2146. // be set to a *QueryError value including the Err message and the Code. In
  2147. // those cases, the result argument is still unmarshalled into with the
  2148. // received document so that any other custom values may be obtained if
  2149. // desired.
  2150. //
  2151. func (q *Query) One(result interface{}) (err error) {
  2152. q.m.Lock()
  2153. session := q.session
  2154. op := q.op // Copy.
  2155. q.m.Unlock()
  2156. socket, err := session.acquireSocket(true)
  2157. if err != nil {
  2158. return err
  2159. }
  2160. defer socket.Release()
  2161. op.flags |= session.slaveOkFlag()
  2162. op.limit = -1
  2163. data, err := socket.SimpleQuery(&op)
  2164. if err != nil {
  2165. return err
  2166. }
  2167. if data == nil {
  2168. return ErrNotFound
  2169. }
  2170. if result != nil {
  2171. err = bson.Unmarshal(data, result)
  2172. if err == nil {
  2173. debugf("Query %p document unmarshaled: %#v", q, result)
  2174. } else {
  2175. debugf("Query %p document unmarshaling failed: %#v", q, err)
  2176. return err
  2177. }
  2178. }
  2179. return checkQueryError(op.collection, data)
  2180. }
  2181. // The DBRef type implements support for the database reference MongoDB
  2182. // convention as supported by multiple drivers. This convention enables
  2183. // cross-referencing documents between collections and databases using
  2184. // a structure which includes a collection name, a document id, and
  2185. // optionally a database name.
  2186. //
  2187. // See the FindRef methods on Session and on Database.
  2188. //
  2189. // Relevant documentation:
  2190. //
  2191. // http://www.mongodb.org/display/DOCS/Database+References
  2192. //
  2193. type DBRef struct {
  2194. Collection string `bson:"$ref"`
  2195. Id interface{} `bson:"$id"`
  2196. Database string `bson:"$db,omitempty"`
  2197. }
  2198. // NOTE: Order of fields for DBRef above does matter, per documentation.
  2199. // FindRef returns a query that looks for the document in the provided
  2200. // reference. If the reference includes the DB field, the document will
  2201. // be retrieved from the respective database.
  2202. //
  2203. // See also the DBRef type and the FindRef method on Session.
  2204. //
  2205. // Relevant documentation:
  2206. //
  2207. // http://www.mongodb.org/display/DOCS/Database+References
  2208. //
  2209. func (db *Database) FindRef(ref *DBRef) *Query {
  2210. var c *Collection
  2211. if ref.Database == "" {
  2212. c = db.C(ref.Collection)
  2213. } else {
  2214. c = db.Session.DB(ref.Database).C(ref.Collection)
  2215. }
  2216. return c.FindId(ref.Id)
  2217. }
  2218. // FindRef returns a query that looks for the document in the provided
  2219. // reference. For a DBRef to be resolved correctly at the session level
  2220. // it must necessarily have the optional DB field defined.
  2221. //
  2222. // See also the DBRef type and the FindRef method on Database.
  2223. //
  2224. // Relevant documentation:
  2225. //
  2226. // http://www.mongodb.org/display/DOCS/Database+References
  2227. //
  2228. func (s *Session) FindRef(ref *DBRef) *Query {
  2229. if ref.Database == "" {
  2230. panic(errors.New(fmt.Sprintf("Can't resolve database for %#v", ref)))
  2231. }
  2232. c := s.DB(ref.Database).C(ref.Collection)
  2233. return c.FindId(ref.Id)
  2234. }
  2235. // CollectionNames returns the collection names present in database.
  2236. func (db *Database) CollectionNames() (names []string, err error) {
  2237. c := len(db.Name) + 1
  2238. iter := db.C("system.namespaces").Find(nil).Iter()
  2239. var result *struct{ Name string }
  2240. for iter.Next(&result) {
  2241. if strings.Index(result.Name, "$") < 0 || strings.Index(result.Name, ".oplog.$") >= 0 {
  2242. names = append(names, result.Name[c:])
  2243. }
  2244. }
  2245. if err := iter.Close(); err != nil {
  2246. return nil, err
  2247. }
  2248. sort.Strings(names)
  2249. return names, nil
  2250. }
  2251. type dbNames struct {
  2252. Databases []struct {
  2253. Name string
  2254. Empty bool
  2255. }
  2256. }
  2257. // DatabaseNames returns the names of non-empty databases present in the cluster.
  2258. func (s *Session) DatabaseNames() (names []string, err error) {
  2259. var result dbNames
  2260. err = s.Run("listDatabases", &result)
  2261. if err != nil {
  2262. return nil, err
  2263. }
  2264. for _, db := range result.Databases {
  2265. if !db.Empty {
  2266. names = append(names, db.Name)
  2267. }
  2268. }
  2269. sort.Strings(names)
  2270. return names, nil
  2271. }
  2272. // Iter executes the query and returns an iterator capable of going over all
  2273. // the results. Results will be returned in batches of configurable
  2274. // size (see the Batch method) and more documents will be requested when a
  2275. // configurable number of documents is iterated over (see the Prefetch method).
  2276. func (q *Query) Iter() *Iter {
  2277. q.m.Lock()
  2278. session := q.session
  2279. op := q.op
  2280. prefetch := q.prefetch
  2281. limit := q.limit
  2282. q.m.Unlock()
  2283. iter := &Iter{
  2284. session: session,
  2285. prefetch: prefetch,
  2286. limit: limit,
  2287. timeout: -1,
  2288. }
  2289. iter.gotReply.L = &iter.m
  2290. iter.op.collection = op.collection
  2291. iter.op.limit = op.limit
  2292. iter.op.replyFunc = iter.replyFunc()
  2293. iter.docsToReceive++
  2294. op.replyFunc = iter.op.replyFunc
  2295. op.flags |= session.slaveOkFlag()
  2296. socket, err := session.acquireSocket(true)
  2297. if err != nil {
  2298. iter.err = err
  2299. } else {
  2300. iter.server = socket.Server()
  2301. err = socket.Query(&op)
  2302. if err != nil {
  2303. // Must lock as the query above may call replyFunc.
  2304. iter.m.Lock()
  2305. iter.err = err
  2306. iter.m.Unlock()
  2307. }
  2308. socket.Release()
  2309. }
  2310. return iter
  2311. }
  2312. // Tail returns a tailable iterator. Unlike a normal iterator, a
  2313. // tailable iterator may wait for new values to be inserted in the
  2314. // collection once the end of the current result set is reached,
  2315. // A tailable iterator may only be used with capped collections.
  2316. //
  2317. // The timeout parameter indicates how long Next will block waiting
  2318. // for a result before timing out. If set to -1, Next will not
  2319. // timeout, and will continue waiting for a result for as long as
  2320. // the cursor is valid and the session is not closed. If set to 0,
  2321. // Next times out as soon as it reaches the end of the result set.
  2322. // Otherwise, Next will wait for at least the given number of
  2323. // seconds for a new document to be available before timing out.
  2324. //
  2325. // On timeouts, Next will unblock and return false, and the Timeout
  2326. // method will return true if called. In these cases, Next may still
  2327. // be called again on the same iterator to check if a new value is
  2328. // available at the current cursor position, and again it will block
  2329. // according to the specified timeoutSecs. If the cursor becomes
  2330. // invalid, though, both Next and Timeout will return false and
  2331. // the query must be restarted.
  2332. //
  2333. // The following example demonstrates timeout handling and query
  2334. // restarting:
  2335. //
  2336. // iter := collection.Find(nil).Sort("$natural").Tail(5 * time.Second)
  2337. // for {
  2338. // for iter.Next(&result) {
  2339. // fmt.Println(result.Id)
  2340. // lastId = result.Id
  2341. // }
  2342. // if iter.Err() != nil {
  2343. // return iter.Close()
  2344. // }
  2345. // if iter.Timeout() {
  2346. // continue
  2347. // }
  2348. // query := collection.Find(bson.M{"_id": bson.M{"$gt": lastId}})
  2349. // iter = query.Sort("$natural").Tail(5 * time.Second)
  2350. // }
  2351. // iter.Close()
  2352. //
  2353. // Relevant documentation:
  2354. //
  2355. // http://www.mongodb.org/display/DOCS/Tailable+Cursors
  2356. // http://www.mongodb.org/display/DOCS/Capped+Collections
  2357. // http://www.mongodb.org/display/DOCS/Sorting+and+Natural+Order
  2358. //
  2359. func (q *Query) Tail(timeout time.Duration) *Iter {
  2360. q.m.Lock()
  2361. session := q.session
  2362. op := q.op
  2363. prefetch := q.prefetch
  2364. q.m.Unlock()
  2365. iter := &Iter{session: session, prefetch: prefetch}
  2366. iter.gotReply.L = &iter.m
  2367. iter.timeout = timeout
  2368. iter.op.collection = op.collection
  2369. iter.op.limit = op.limit
  2370. iter.op.replyFunc = iter.replyFunc()
  2371. iter.docsToReceive++
  2372. op.replyFunc = iter.op.replyFunc
  2373. op.flags |= flagTailable | flagAwaitData | session.slaveOkFlag()
  2374. socket, err := session.acquireSocket(true)
  2375. if err != nil {
  2376. iter.err = err
  2377. } else {
  2378. iter.server = socket.Server()
  2379. err = socket.Query(&op)
  2380. if err != nil {
  2381. // Must lock as the query above may call replyFunc.
  2382. iter.m.Lock()
  2383. iter.err = err
  2384. iter.m.Unlock()
  2385. }
  2386. socket.Release()
  2387. }
  2388. return iter
  2389. }
  2390. func (s *Session) slaveOkFlag() (flag queryOpFlags) {
  2391. s.m.RLock()
  2392. if s.slaveOk {
  2393. flag = flagSlaveOk
  2394. }
  2395. s.m.RUnlock()
  2396. return
  2397. }
  2398. // Err returns nil if no errors happened during iteration, or the actual
  2399. // error otherwise.
  2400. //
  2401. // In case a resulting document included a field named $err or errmsg, which are
  2402. // standard ways for MongoDB to report an improper query, the returned value has
  2403. // a *QueryError type, and includes the Err message and the Code.
  2404. func (iter *Iter) Err() error {
  2405. iter.m.Lock()
  2406. err := iter.err
  2407. iter.m.Unlock()
  2408. if err == ErrNotFound {
  2409. return nil
  2410. }
  2411. return err
  2412. }
  2413. // Close kills the server cursor used by the iterator, if any, and returns
  2414. // nil if no errors happened during iteration, or the actual error otherwise.
  2415. //
  2416. // Server cursors are automatically closed at the end of an iteration, which
  2417. // means close will do nothing unless the iteration was interrupted before
  2418. // the server finished sending results to the driver. If Close is not called
  2419. // in such a situation, the cursor will remain available at the server until
  2420. // the default cursor timeout period is reached. No further problems arise.
  2421. //
  2422. // Close is idempotent. That means it can be called repeatedly and will
  2423. // return the same result every time.
  2424. //
  2425. // In case a resulting document included a field named $err or errmsg, which are
  2426. // standard ways for MongoDB to report an improper query, the returned value has
  2427. // a *QueryError type.
  2428. func (iter *Iter) Close() error {
  2429. iter.m.Lock()
  2430. iter.killCursor()
  2431. err := iter.err
  2432. iter.m.Unlock()
  2433. if err == ErrNotFound {
  2434. return nil
  2435. }
  2436. return err
  2437. }
  2438. func (iter *Iter) killCursor() error {
  2439. if iter.op.cursorId != 0 {
  2440. socket, err := iter.acquireSocket()
  2441. if err == nil {
  2442. // TODO Batch kills.
  2443. err = socket.Query(&killCursorsOp{[]int64{iter.op.cursorId}})
  2444. socket.Release()
  2445. }
  2446. if err != nil && (iter.err == nil || iter.err == ErrNotFound) {
  2447. iter.err = err
  2448. }
  2449. iter.op.cursorId = 0
  2450. return err
  2451. }
  2452. return nil
  2453. }
  2454. // Timeout returns true if Next returned false due to a timeout of
  2455. // a tailable cursor. In those cases, Next may be called again to continue
  2456. // the iteration at the previous cursor position.
  2457. func (iter *Iter) Timeout() bool {
  2458. iter.m.Lock()
  2459. result := iter.timedout
  2460. iter.m.Unlock()
  2461. return result
  2462. }
  2463. // Next retrieves the next document from the result set, blocking if necessary.
  2464. // This method will also automatically retrieve another batch of documents from
  2465. // the server when the current one is exhausted, or before that in background
  2466. // if pre-fetching is enabled (see the Query.Prefetch and Session.SetPrefetch
  2467. // methods).
  2468. //
  2469. // Next returns true if a document was successfully unmarshalled onto result,
  2470. // and false at the end of the result set or if an error happened.
  2471. // When Next returns false, the Err method should be called to verify if
  2472. // there was an error during iteration.
  2473. //
  2474. // For example:
  2475. //
  2476. // iter := collection.Find(nil).Iter()
  2477. // for iter.Next(&result) {
  2478. // fmt.Printf("Result: %v\n", result.Id)
  2479. // }
  2480. // if err := iter.Close(); err != nil {
  2481. // return err
  2482. // }
  2483. //
  2484. func (iter *Iter) Next(result interface{}) bool {
  2485. iter.m.Lock()
  2486. iter.timedout = false
  2487. timeout := time.Time{}
  2488. for iter.err == nil && iter.docData.Len() == 0 && (iter.docsToReceive > 0 || iter.op.cursorId != 0) {
  2489. if iter.docsToReceive == 0 {
  2490. if iter.timeout >= 0 {
  2491. if timeout.IsZero() {
  2492. timeout = time.Now().Add(iter.timeout)
  2493. }
  2494. if time.Now().After(timeout) {
  2495. iter.timedout = true
  2496. iter.m.Unlock()
  2497. return false
  2498. }
  2499. }
  2500. iter.getMore()
  2501. if iter.err != nil {
  2502. break
  2503. }
  2504. }
  2505. iter.gotReply.Wait()
  2506. }
  2507. // Exhaust available data before reporting any errors.
  2508. if docData, ok := iter.docData.Pop().([]byte); ok {
  2509. if iter.limit > 0 {
  2510. iter.limit--
  2511. if iter.limit == 0 {
  2512. if iter.docData.Len() > 0 {
  2513. iter.m.Unlock()
  2514. panic(fmt.Errorf("data remains after limit exhausted: %d", iter.docData.Len()))
  2515. }
  2516. iter.err = ErrNotFound
  2517. if iter.killCursor() != nil {
  2518. iter.m.Unlock()
  2519. return false
  2520. }
  2521. }
  2522. }
  2523. if iter.op.cursorId != 0 && iter.err == nil {
  2524. if iter.docsBeforeMore == 0 {
  2525. iter.getMore()
  2526. }
  2527. iter.docsBeforeMore-- // Goes negative.
  2528. }
  2529. iter.m.Unlock()
  2530. err := bson.Unmarshal(docData, result)
  2531. if err != nil {
  2532. debugf("Iter %p document unmarshaling failed: %#v", iter, err)
  2533. iter.m.Lock()
  2534. if iter.err == nil {
  2535. iter.err = err
  2536. }
  2537. iter.m.Unlock()
  2538. return false
  2539. }
  2540. debugf("Iter %p document unmarshaled: %#v", iter, result)
  2541. // XXX Only have to check first document for a query error?
  2542. err = checkQueryError(iter.op.collection, docData)
  2543. if err != nil {
  2544. iter.m.Lock()
  2545. if iter.err == nil {
  2546. iter.err = err
  2547. }
  2548. iter.m.Unlock()
  2549. return false
  2550. }
  2551. return true
  2552. } else if iter.err != nil {
  2553. debugf("Iter %p returning false: %s", iter, iter.err)
  2554. iter.m.Unlock()
  2555. return false
  2556. } else if iter.op.cursorId == 0 {
  2557. iter.err = ErrNotFound
  2558. debugf("Iter %p exhausted with cursor=0", iter)
  2559. iter.m.Unlock()
  2560. return false
  2561. }
  2562. panic("unreachable")
  2563. }
  2564. // All retrieves all documents from the result set into the provided slice
  2565. // and closes the iterator.
  2566. //
  2567. // The result argument must necessarily be the address for a slice. The slice
  2568. // may be nil or previously allocated.
  2569. //
  2570. // WARNING: Obviously, All must not be used with result sets that may be
  2571. // potentially large, since it may consume all memory until the system
  2572. // crashes. Consider building the query with a Limit clause to ensure the
  2573. // result size is bounded.
  2574. //
  2575. // For instance:
  2576. //
  2577. // var result []struct{ Value int }
  2578. // iter := collection.Find(nil).Limit(100).Iter()
  2579. // err := iter.All(&result)
  2580. // if err != nil {
  2581. // return err
  2582. // }
  2583. //
  2584. func (iter *Iter) All(result interface{}) error {
  2585. resultv := reflect.ValueOf(result)
  2586. if resultv.Kind() != reflect.Ptr || resultv.Elem().Kind() != reflect.Slice {
  2587. panic("result argument must be a slice address")
  2588. }
  2589. slicev := resultv.Elem()
  2590. slicev = slicev.Slice(0, slicev.Cap())
  2591. elemt := slicev.Type().Elem()
  2592. i := 0
  2593. for {
  2594. if slicev.Len() == i {
  2595. elemp := reflect.New(elemt)
  2596. if !iter.Next(elemp.Interface()) {
  2597. break
  2598. }
  2599. slicev = reflect.Append(slicev, elemp.Elem())
  2600. slicev = slicev.Slice(0, slicev.Cap())
  2601. } else {
  2602. if !iter.Next(slicev.Index(i).Addr().Interface()) {
  2603. break
  2604. }
  2605. }
  2606. i++
  2607. }
  2608. resultv.Elem().Set(slicev.Slice(0, i))
  2609. return iter.Close()
  2610. }
  2611. // All works like Iter.All.
  2612. func (q *Query) All(result interface{}) error {
  2613. return q.Iter().All(result)
  2614. }
  2615. // The For method is obsolete and will be removed in a future release.
  2616. // See Iter as an elegant replacement.
  2617. func (q *Query) For(result interface{}, f func() error) error {
  2618. return q.Iter().For(result, f)
  2619. }
  2620. // The For method is obsolete and will be removed in a future release.
  2621. // See Iter as an elegant replacement.
  2622. func (iter *Iter) For(result interface{}, f func() error) (err error) {
  2623. valid := false
  2624. v := reflect.ValueOf(result)
  2625. if v.Kind() == reflect.Ptr {
  2626. v = v.Elem()
  2627. switch v.Kind() {
  2628. case reflect.Map, reflect.Ptr, reflect.Interface, reflect.Slice:
  2629. valid = v.IsNil()
  2630. }
  2631. }
  2632. if !valid {
  2633. panic("For needs a pointer to nil reference value. See the documentation.")
  2634. }
  2635. zero := reflect.Zero(v.Type())
  2636. for {
  2637. v.Set(zero)
  2638. if !iter.Next(result) {
  2639. break
  2640. }
  2641. err = f()
  2642. if err != nil {
  2643. return err
  2644. }
  2645. }
  2646. return iter.Err()
  2647. }
  2648. func (iter *Iter) acquireSocket() (*mongoSocket, error) {
  2649. socket, err := iter.session.acquireSocket(true)
  2650. if err != nil {
  2651. return nil, err
  2652. }
  2653. if socket.Server() != iter.server {
  2654. // Socket server changed during iteration. This may happen
  2655. // with Eventual sessions, if a Refresh is done, or if a
  2656. // monotonic session gets a write and shifts from secondary
  2657. // to primary. Our cursor is in a specific server, though.
  2658. iter.session.m.Lock()
  2659. sockTimeout := iter.session.sockTimeout
  2660. iter.session.m.Unlock()
  2661. socket.Release()
  2662. socket, _, err = iter.server.AcquireSocket(0, sockTimeout)
  2663. if err != nil {
  2664. return nil, err
  2665. }
  2666. err := iter.session.socketLogin(socket)
  2667. if err != nil {
  2668. socket.Release()
  2669. return nil, err
  2670. }
  2671. }
  2672. return socket, nil
  2673. }
  2674. func (iter *Iter) getMore() {
  2675. socket, err := iter.acquireSocket()
  2676. if err != nil {
  2677. iter.err = err
  2678. return
  2679. }
  2680. defer socket.Release()
  2681. debugf("Iter %p requesting more documents", iter)
  2682. if iter.limit > 0 {
  2683. limit := iter.limit - int32(iter.docsToReceive) - int32(iter.docData.Len())
  2684. if limit < iter.op.limit {
  2685. iter.op.limit = limit
  2686. }
  2687. }
  2688. if err := socket.Query(&iter.op); err != nil {
  2689. iter.err = err
  2690. }
  2691. iter.docsToReceive++
  2692. }
  2693. type countCmd struct {
  2694. Count string
  2695. Query interface{}
  2696. Limit int32 ",omitempty"
  2697. Skip int32 ",omitempty"
  2698. }
  2699. // Count returns the total number of documents in the result set.
  2700. func (q *Query) Count() (n int, err error) {
  2701. q.m.Lock()
  2702. session := q.session
  2703. op := q.op
  2704. limit := q.limit
  2705. q.m.Unlock()
  2706. c := strings.Index(op.collection, ".")
  2707. if c < 0 {
  2708. return 0, errors.New("Bad collection name: " + op.collection)
  2709. }
  2710. dbname := op.collection[:c]
  2711. cname := op.collection[c+1:]
  2712. result := struct{ N int }{}
  2713. err = session.DB(dbname).Run(countCmd{cname, op.query, limit, op.skip}, &result)
  2714. return result.N, err
  2715. }
  2716. // Count returns the total number of documents in the collection.
  2717. func (c *Collection) Count() (n int, err error) {
  2718. return c.Find(nil).Count()
  2719. }
  2720. type distinctCmd struct {
  2721. Collection string "distinct"
  2722. Key string
  2723. Query interface{} ",omitempty"
  2724. }
  2725. // Distinct returns a list of distinct values for the given key within
  2726. // the result set. The list of distinct values will be unmarshalled
  2727. // in the "values" key of the provided result parameter.
  2728. //
  2729. // For example:
  2730. //
  2731. // var result []int
  2732. // err := collection.Find(bson.M{"gender": "F"}).Distinct("age", &result)
  2733. //
  2734. // Relevant documentation:
  2735. //
  2736. // http://www.mongodb.org/display/DOCS/Aggregation
  2737. //
  2738. func (q *Query) Distinct(key string, result interface{}) error {
  2739. q.m.Lock()
  2740. session := q.session
  2741. op := q.op // Copy.
  2742. q.m.Unlock()
  2743. c := strings.Index(op.collection, ".")
  2744. if c < 0 {
  2745. return errors.New("Bad collection name: " + op.collection)
  2746. }
  2747. dbname := op.collection[:c]
  2748. cname := op.collection[c+1:]
  2749. var doc struct{ Values bson.Raw }
  2750. err := session.DB(dbname).Run(distinctCmd{cname, key, op.query}, &doc)
  2751. if err != nil {
  2752. return err
  2753. }
  2754. return doc.Values.Unmarshal(result)
  2755. }
  2756. type mapReduceCmd struct {
  2757. Collection string "mapreduce"
  2758. Map string ",omitempty"
  2759. Reduce string ",omitempty"
  2760. Finalize string ",omitempty"
  2761. Limit int32 ",omitempty"
  2762. Out interface{}
  2763. Query interface{} ",omitempty"
  2764. Sort interface{} ",omitempty"
  2765. Scope interface{} ",omitempty"
  2766. Verbose bool ",omitempty"
  2767. }
  2768. type mapReduceResult struct {
  2769. Results bson.Raw
  2770. Result bson.Raw
  2771. TimeMillis int64 "timeMillis"
  2772. Counts struct{ Input, Emit, Output int }
  2773. Ok bool
  2774. Err string
  2775. Timing *MapReduceTime
  2776. }
  2777. type MapReduce struct {
  2778. Map string // Map Javascript function code (required)
  2779. Reduce string // Reduce Javascript function code (required)
  2780. Finalize string // Finalize Javascript function code (optional)
  2781. Out interface{} // Output collection name or document. If nil, results are inlined into the result parameter.
  2782. Scope interface{} // Optional global scope for Javascript functions
  2783. Verbose bool
  2784. }
  2785. type MapReduceInfo struct {
  2786. InputCount int // Number of documents mapped
  2787. EmitCount int // Number of times reduce called emit
  2788. OutputCount int // Number of documents in resulting collection
  2789. Database string // Output database, if results are not inlined
  2790. Collection string // Output collection, if results are not inlined
  2791. Time int64 // Time to run the job, in nanoseconds
  2792. VerboseTime *MapReduceTime // Only defined if Verbose was true
  2793. }
  2794. type MapReduceTime struct {
  2795. Total int64 // Total time, in nanoseconds
  2796. Map int64 "mapTime" // Time within map function, in nanoseconds
  2797. EmitLoop int64 "emitLoop" // Time within the emit/map loop, in nanoseconds
  2798. }
  2799. // MapReduce executes a map/reduce job for documents covered by the query.
  2800. // That kind of job is suitable for very flexible bulk aggregation of data
  2801. // performed at the server side via Javascript functions.
  2802. //
  2803. // Results from the job may be returned as a result of the query itself
  2804. // through the result parameter in case they'll certainly fit in memory
  2805. // and in a single document. If there's the possibility that the amount
  2806. // of data might be too large, results must be stored back in an alternative
  2807. // collection or even a separate database, by setting the Out field of the
  2808. // provided MapReduce job. In that case, provide nil as the result parameter.
  2809. //
  2810. // These are some of the ways to set Out:
  2811. //
  2812. // nil
  2813. // Inline results into the result parameter.
  2814. //
  2815. // bson.M{"replace": "mycollection"}
  2816. // The output will be inserted into a collection which replaces any
  2817. // existing collection with the same name.
  2818. //
  2819. // bson.M{"merge": "mycollection"}
  2820. // This option will merge new data into the old output collection. In
  2821. // other words, if the same key exists in both the result set and the
  2822. // old collection, the new key will overwrite the old one.
  2823. //
  2824. // bson.M{"reduce": "mycollection"}
  2825. // If documents exist for a given key in the result set and in the old
  2826. // collection, then a reduce operation (using the specified reduce
  2827. // function) will be performed on the two values and the result will be
  2828. // written to the output collection. If a finalize function was
  2829. // provided, this will be run after the reduce as well.
  2830. //
  2831. // bson.M{...., "db": "mydb"}
  2832. // Any of the above options can have the "db" key included for doing
  2833. // the respective action in a separate database.
  2834. //
  2835. // The following is a trivial example which will count the number of
  2836. // occurrences of a field named n on each document in a collection, and
  2837. // will return results inline:
  2838. //
  2839. // job := &mgo.MapReduce{
  2840. // Map: "function() { emit(this.n, 1) }",
  2841. // Reduce: "function(key, values) { return Array.sum(values) }",
  2842. // }
  2843. // var result []struct { Id int "_id"; Value int }
  2844. // _, err := collection.Find(nil).MapReduce(job, &result)
  2845. // if err != nil {
  2846. // return err
  2847. // }
  2848. // for _, item := range result {
  2849. // fmt.Println(item.Value)
  2850. // }
  2851. //
  2852. // This function is compatible with MongoDB 1.7.4+.
  2853. //
  2854. // Relevant documentation:
  2855. //
  2856. // http://www.mongodb.org/display/DOCS/MapReduce
  2857. //
  2858. func (q *Query) MapReduce(job *MapReduce, result interface{}) (info *MapReduceInfo, err error) {
  2859. q.m.Lock()
  2860. session := q.session
  2861. op := q.op // Copy.
  2862. limit := q.limit
  2863. q.m.Unlock()
  2864. c := strings.Index(op.collection, ".")
  2865. if c < 0 {
  2866. return nil, errors.New("Bad collection name: " + op.collection)
  2867. }
  2868. dbname := op.collection[:c]
  2869. cname := op.collection[c+1:]
  2870. cmd := mapReduceCmd{
  2871. Collection: cname,
  2872. Map: job.Map,
  2873. Reduce: job.Reduce,
  2874. Finalize: job.Finalize,
  2875. Out: fixMROut(job.Out),
  2876. Scope: job.Scope,
  2877. Verbose: job.Verbose,
  2878. Query: op.query,
  2879. Sort: op.options.OrderBy,
  2880. Limit: limit,
  2881. }
  2882. if cmd.Out == nil {
  2883. cmd.Out = bson.M{"inline": 1}
  2884. }
  2885. var doc mapReduceResult
  2886. err = session.DB(dbname).Run(&cmd, &doc)
  2887. if err != nil {
  2888. return nil, err
  2889. }
  2890. if doc.Err != "" {
  2891. return nil, errors.New(doc.Err)
  2892. }
  2893. info = &MapReduceInfo{
  2894. InputCount: doc.Counts.Input,
  2895. EmitCount: doc.Counts.Emit,
  2896. OutputCount: doc.Counts.Output,
  2897. Time: doc.TimeMillis * 1e6,
  2898. }
  2899. if doc.Result.Kind == 0x02 {
  2900. err = doc.Result.Unmarshal(&info.Collection)
  2901. info.Database = dbname
  2902. } else if doc.Result.Kind == 0x03 {
  2903. var v struct{ Collection, Db string }
  2904. err = doc.Result.Unmarshal(&v)
  2905. info.Collection = v.Collection
  2906. info.Database = v.Db
  2907. }
  2908. if doc.Timing != nil {
  2909. info.VerboseTime = doc.Timing
  2910. info.VerboseTime.Total *= 1e6
  2911. info.VerboseTime.Map *= 1e6
  2912. info.VerboseTime.EmitLoop *= 1e6
  2913. }
  2914. if err != nil {
  2915. return nil, err
  2916. }
  2917. if result != nil {
  2918. return info, doc.Results.Unmarshal(result)
  2919. }
  2920. return info, nil
  2921. }
  2922. // The "out" option in the MapReduce command must be ordered. This was
  2923. // found after the implementation was accepting maps for a long time,
  2924. // so rather than breaking the API, we'll fix the order if necessary.
  2925. // Details about the order requirement may be seen in MongoDB's code:
  2926. //
  2927. // http://goo.gl/L8jwJX
  2928. //
  2929. func fixMROut(out interface{}) interface{} {
  2930. outv := reflect.ValueOf(out)
  2931. if outv.Kind() != reflect.Map || outv.Type().Key() != reflect.TypeOf("") {
  2932. return out
  2933. }
  2934. outs := make(bson.D, outv.Len())
  2935. outTypeIndex := -1
  2936. for i, k := range outv.MapKeys() {
  2937. ks := k.String()
  2938. outs[i].Name = ks
  2939. outs[i].Value = outv.MapIndex(k).Interface()
  2940. switch ks {
  2941. case "normal", "replace", "merge", "reduce", "inline":
  2942. outTypeIndex = i
  2943. }
  2944. }
  2945. if outTypeIndex > 0 {
  2946. outs[0], outs[outTypeIndex] = outs[outTypeIndex], outs[0]
  2947. }
  2948. return outs
  2949. }
  2950. // Change holds fields for running a findAndModify MongoDB command via
  2951. // the Query.Apply method.
  2952. type Change struct {
  2953. Update interface{} // The update document
  2954. Upsert bool // Whether to insert in case the document isn't found
  2955. Remove bool // Whether to remove the document found rather than updating
  2956. ReturnNew bool // Should the modified document be returned rather than the old one
  2957. }
  2958. type findModifyCmd struct {
  2959. Collection string "findAndModify"
  2960. Query, Update, Sort, Fields interface{} ",omitempty"
  2961. Upsert, Remove, New bool ",omitempty"
  2962. }
  2963. type valueResult struct {
  2964. Value bson.Raw
  2965. LastError LastError "lastErrorObject"
  2966. }
  2967. // Apply runs the findAndModify MongoDB command, which allows updating, upserting
  2968. // or removing a document matching a query and atomically returning either the old
  2969. // version (the default) or the new version of the document (when ReturnNew is true).
  2970. // If no objects are found Apply returns ErrNotFound.
  2971. //
  2972. // The Sort and Select query methods affect the result of Apply. In case
  2973. // multiple documents match the query, Sort enables selecting which document to
  2974. // act upon by ordering it first. Select enables retrieving only a selection
  2975. // of fields of the new or old document.
  2976. //
  2977. // This simple example increments a counter and prints its new value:
  2978. //
  2979. // change := mgo.Change{
  2980. // Update: bson.M{"$inc": bson.M{"n": 1}},
  2981. // ReturnNew: true,
  2982. // }
  2983. // info, err = col.Find(M{"_id": id}).Apply(change, &doc)
  2984. // fmt.Println(doc.N)
  2985. //
  2986. // This method depends on MongoDB >= 2.0 to work properly.
  2987. //
  2988. // Relevant documentation:
  2989. //
  2990. // http://www.mongodb.org/display/DOCS/findAndModify+Command
  2991. // http://www.mongodb.org/display/DOCS/Updating
  2992. // http://www.mongodb.org/display/DOCS/Atomic+Operations
  2993. //
  2994. func (q *Query) Apply(change Change, result interface{}) (info *ChangeInfo, err error) {
  2995. q.m.Lock()
  2996. session := q.session
  2997. op := q.op // Copy.
  2998. q.m.Unlock()
  2999. c := strings.Index(op.collection, ".")
  3000. if c < 0 {
  3001. return nil, errors.New("bad collection name: " + op.collection)
  3002. }
  3003. dbname := op.collection[:c]
  3004. cname := op.collection[c+1:]
  3005. cmd := findModifyCmd{
  3006. Collection: cname,
  3007. Update: change.Update,
  3008. Upsert: change.Upsert,
  3009. Remove: change.Remove,
  3010. New: change.ReturnNew,
  3011. Query: op.query,
  3012. Sort: op.options.OrderBy,
  3013. Fields: op.selector,
  3014. }
  3015. session = session.Clone()
  3016. defer session.Close()
  3017. session.SetMode(Strong, false)
  3018. var doc valueResult
  3019. err = session.DB(dbname).Run(&cmd, &doc)
  3020. if err != nil {
  3021. if qerr, ok := err.(*QueryError); ok && qerr.Message == "No matching object found" {
  3022. return nil, ErrNotFound
  3023. }
  3024. return nil, err
  3025. }
  3026. if doc.LastError.N == 0 {
  3027. return nil, ErrNotFound
  3028. }
  3029. if doc.Value.Kind != 0x0A {
  3030. err = doc.Value.Unmarshal(result)
  3031. if err != nil {
  3032. return nil, err
  3033. }
  3034. }
  3035. info = &ChangeInfo{}
  3036. lerr := &doc.LastError
  3037. if lerr.UpdatedExisting {
  3038. info.Updated = lerr.N
  3039. } else if change.Remove {
  3040. info.Removed = lerr.N
  3041. } else if change.Upsert {
  3042. info.UpsertedId = lerr.UpsertedId
  3043. }
  3044. return info, nil
  3045. }
  3046. // The BuildInfo type encapsulates details about the running MongoDB server.
  3047. //
  3048. // Note that the VersionArray field was introduced in MongoDB 2.0+, but it is
  3049. // internally assembled from the Version information for previous versions.
  3050. // In both cases, VersionArray is guaranteed to have at least 4 entries.
  3051. type BuildInfo struct {
  3052. Version string
  3053. VersionArray []int `bson:"versionArray"` // On MongoDB 2.0+; assembled from Version otherwise
  3054. GitVersion string `bson:"gitVersion"`
  3055. SysInfo string `bson:"sysInfo"`
  3056. Bits int
  3057. Debug bool
  3058. MaxObjectSize int `bson:"maxBsonObjectSize"`
  3059. }
  3060. // VersionAtLeast returns whether the BuildInfo version is greater than or
  3061. // equal to the provided version number. If more than one number is
  3062. // provided, numbers will be considered as major, minor, and so on.
  3063. func (bi *BuildInfo) VersionAtLeast(version ...int) bool {
  3064. for i := range version {
  3065. if i == len(bi.VersionArray) {
  3066. return false
  3067. }
  3068. if bi.VersionArray[i] < version[i] {
  3069. return false
  3070. }
  3071. }
  3072. return true
  3073. }
  3074. // BuildInfo retrieves the version and other details about the
  3075. // running MongoDB server.
  3076. func (s *Session) BuildInfo() (info BuildInfo, err error) {
  3077. err = s.Run(bson.D{{"buildInfo", "1"}}, &info)
  3078. if len(info.VersionArray) == 0 {
  3079. for _, a := range strings.Split(info.Version, ".") {
  3080. i, err := strconv.Atoi(a)
  3081. if err != nil {
  3082. break
  3083. }
  3084. info.VersionArray = append(info.VersionArray, i)
  3085. }
  3086. }
  3087. for len(info.VersionArray) < 4 {
  3088. info.VersionArray = append(info.VersionArray, 0)
  3089. }
  3090. return
  3091. }
  3092. // ---------------------------------------------------------------------------
  3093. // Internal session handling helpers.
  3094. func (s *Session) acquireSocket(slaveOk bool) (*mongoSocket, error) {
  3095. // Read-only lock to check for previously reserved socket.
  3096. s.m.RLock()
  3097. if s.masterSocket != nil {
  3098. socket := s.masterSocket
  3099. socket.Acquire()
  3100. s.m.RUnlock()
  3101. return socket, nil
  3102. }
  3103. if s.slaveSocket != nil && s.slaveOk && slaveOk {
  3104. socket := s.slaveSocket
  3105. socket.Acquire()
  3106. s.m.RUnlock()
  3107. return socket, nil
  3108. }
  3109. s.m.RUnlock()
  3110. // No go. We may have to request a new socket and change the session,
  3111. // so try again but with an exclusive lock now.
  3112. s.m.Lock()
  3113. defer s.m.Unlock()
  3114. if s.masterSocket != nil {
  3115. s.masterSocket.Acquire()
  3116. return s.masterSocket, nil
  3117. }
  3118. if s.slaveSocket != nil && s.slaveOk && slaveOk {
  3119. s.slaveSocket.Acquire()
  3120. return s.slaveSocket, nil
  3121. }
  3122. // Still not good. We need a new socket.
  3123. sock, err := s.cluster().AcquireSocket(slaveOk && s.slaveOk, s.syncTimeout, s.sockTimeout, s.queryConfig.op.serverTags)
  3124. if err != nil {
  3125. return nil, err
  3126. }
  3127. // Authenticate the new socket.
  3128. if err = s.socketLogin(sock); err != nil {
  3129. sock.Release()
  3130. return nil, err
  3131. }
  3132. // Keep track of the new socket, if necessary.
  3133. // Note that, as a special case, if the Eventual session was
  3134. // not refreshed (s.slaveSocket != nil), it means the developer
  3135. // asked to preserve an existing reserved socket, so we'll
  3136. // keep a master one around too before a Refresh happens.
  3137. if s.consistency != Eventual || s.slaveSocket != nil {
  3138. s.setSocket(sock)
  3139. }
  3140. // Switch over a Monotonic session to the master.
  3141. if !slaveOk && s.consistency == Monotonic {
  3142. s.slaveOk = false
  3143. }
  3144. return sock, nil
  3145. }
  3146. // setSocket binds socket to this section.
  3147. func (s *Session) setSocket(socket *mongoSocket) {
  3148. info := socket.Acquire()
  3149. if info.Master {
  3150. if s.masterSocket != nil {
  3151. panic("setSocket(master) with existing master socket reserved")
  3152. }
  3153. s.masterSocket = socket
  3154. } else {
  3155. if s.slaveSocket != nil {
  3156. panic("setSocket(slave) with existing slave socket reserved")
  3157. }
  3158. s.slaveSocket = socket
  3159. }
  3160. }
  3161. // unsetSocket releases any slave and/or master sockets reserved.
  3162. func (s *Session) unsetSocket() {
  3163. if s.masterSocket != nil {
  3164. s.masterSocket.Release()
  3165. }
  3166. if s.slaveSocket != nil {
  3167. s.slaveSocket.Release()
  3168. }
  3169. s.masterSocket = nil
  3170. s.slaveSocket = nil
  3171. }
  3172. func (iter *Iter) replyFunc() replyFunc {
  3173. return func(err error, op *replyOp, docNum int, docData []byte) {
  3174. iter.m.Lock()
  3175. iter.docsToReceive--
  3176. if err != nil {
  3177. iter.err = err
  3178. debugf("Iter %p received an error: %s", iter, err.Error())
  3179. } else if docNum == -1 {
  3180. debugf("Iter %p received no documents (cursor=%d).", iter, op.cursorId)
  3181. if op != nil && op.cursorId != 0 {
  3182. // It's a tailable cursor.
  3183. iter.op.cursorId = op.cursorId
  3184. } else {
  3185. iter.err = ErrNotFound
  3186. }
  3187. } else {
  3188. rdocs := int(op.replyDocs)
  3189. if docNum == 0 {
  3190. iter.docsToReceive += rdocs - 1
  3191. docsToProcess := iter.docData.Len() + rdocs
  3192. if iter.limit == 0 || int32(docsToProcess) < iter.limit {
  3193. iter.docsBeforeMore = docsToProcess - int(iter.prefetch*float64(rdocs))
  3194. } else {
  3195. iter.docsBeforeMore = -1
  3196. }
  3197. iter.op.cursorId = op.cursorId
  3198. }
  3199. // XXX Handle errors and flags.
  3200. debugf("Iter %p received reply document %d/%d (cursor=%d)", iter, docNum+1, rdocs, op.cursorId)
  3201. iter.docData.Push(docData)
  3202. }
  3203. iter.gotReply.Broadcast()
  3204. iter.m.Unlock()
  3205. }
  3206. }
  3207. // writeQuery runs the given modifying operation, potentially followed up
  3208. // by a getLastError command in case the session is in safe mode. The
  3209. // LastError result is made available in lerr, and if lerr.Err is set it
  3210. // will also be returned as err.
  3211. func (c *Collection) writeQuery(op interface{}) (lerr *LastError, err error) {
  3212. s := c.Database.Session
  3213. dbname := c.Database.Name
  3214. socket, err := s.acquireSocket(dbname == "local")
  3215. if err != nil {
  3216. return nil, err
  3217. }
  3218. defer socket.Release()
  3219. s.m.RLock()
  3220. safeOp := s.safeOp
  3221. s.m.RUnlock()
  3222. if safeOp == nil {
  3223. return nil, socket.Query(op)
  3224. } else {
  3225. var mutex sync.Mutex
  3226. var replyData []byte
  3227. var replyErr error
  3228. mutex.Lock()
  3229. query := *safeOp // Copy the data.
  3230. query.collection = dbname + ".$cmd"
  3231. query.replyFunc = func(err error, reply *replyOp, docNum int, docData []byte) {
  3232. replyData = docData
  3233. replyErr = err
  3234. mutex.Unlock()
  3235. }
  3236. err = socket.Query(op, &query)
  3237. if err != nil {
  3238. return nil, err
  3239. }
  3240. mutex.Lock() // Wait.
  3241. if replyErr != nil {
  3242. return nil, replyErr // XXX TESTME
  3243. }
  3244. if hasErrMsg(replyData) {
  3245. // Looks like getLastError itself failed.
  3246. err = checkQueryError(query.collection, replyData)
  3247. if err != nil {
  3248. return nil, err
  3249. }
  3250. }
  3251. result := &LastError{}
  3252. bson.Unmarshal(replyData, &result)
  3253. debugf("Result from writing query: %#v", result)
  3254. if result.Err != "" {
  3255. return result, result
  3256. }
  3257. return result, nil
  3258. }
  3259. panic("unreachable")
  3260. }
  3261. func hasErrMsg(d []byte) bool {
  3262. l := len(d)
  3263. for i := 0; i+8 < l; i++ {
  3264. if d[i] == '\x02' && d[i+1] == 'e' && d[i+2] == 'r' && d[i+3] == 'r' && d[i+4] == 'm' && d[i+5] == 's' && d[i+6] == 'g' && d[i+7] == '\x00' {
  3265. return true
  3266. }
  3267. }
  3268. return false
  3269. }