1 USING: accessors arrays assocs bson.constants combinators
2 combinators.smart constructors destructors fry hashtables io
3 io.pools io.sockets kernel linked-assocs locals math
4 mongodb.cmd mongodb.connection mongodb.msg namespaces parser
5 prettyprint prettyprint.custom prettyprint.sections sequences
6 sets splitting strings ;
7 FROM: ascii => ascii? ;
8 FROM: math.bitwise => set-bit ;
11 TUPLE: mdb-pool < pool mdb ;
13 TUPLE: mdb-cursor id query ;
21 CONSTRUCTOR: <mdb-collection> mdb-collection ( name -- collection ) ;
24 { ns string } { name string } { key hashtable } { unique? boolean initial: f } ;
26 CONSTRUCTOR: <index-spec> index-spec ( ns name key -- index-spec ) ;
28 M: mdb-pool make-connection
31 : <mdb-pool> ( mdb -- pool ) [ mdb-pool <pool> ] dip >>mdb ; inline
33 CONSTANT: PARTIAL? "partial?"
35 ERROR: mdb-error msg ;
37 M: mdb-error pprint* ( obj -- )
40 : >pwd-digest ( user password -- digest )
41 "mongo" swap 3array ":" join md5-checksum ;
45 GENERIC: <mdb-cursor> ( id mdb-query-msg/mdb-getmore-msg -- mdb-cursor )
47 M: mdb-query-msg <mdb-cursor>
50 M: mdb-getmore-msg <mdb-cursor>
51 query>> mdb-cursor boa ;
53 : >mdbregexp ( value -- regexp )
54 first <mdbregexp> ; inline
56 GENERIC: update-query ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- )
58 M: mdb-query-msg update-query
59 swap [ start#>> ] [ returned#>> ] bi + >>skip# drop ;
61 M: mdb-getmore-msg update-query
62 query>> update-query ;
64 : make-cursor ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- mdb-cursor/f )
67 [ [ cursor>> ] dip <mdb-cursor> ] 2bi
72 GENERIC: verify-query-result ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- mdb-result-msg mdb-query-msg/mdb-getmore-msg )
74 M: mdb-query-msg verify-query-result ;
76 M: mdb-getmore-msg verify-query-result
77 over flags>> ResultFlag_CursorNotFound =
78 [ nip query>> [ send-query-plain ] keep ] when ;
80 : send-query ( mdb-query-msg/mdb-getmore-msg -- mdb-cursor/f seq )
81 [ send-query-plain ] keep
83 [ collection>> >>collection drop ]
84 [ return#>> >>requested# ]
92 \ / [ >mdbregexp ] parse-literal ;
94 : with-db ( mdb quot -- )
95 '[ _ mdb-open &dispose _ with-connection ] with-destructors ; inline
97 : with-mdb ( mdb quot -- )
99 [ mdb-pool swap with-variable ] curry with-disposal ; inline
101 : with-mdb-pool ( ..a mdb-pool quot -- ..b )
102 '[ _ with-connection ] with-pooled-connection ; inline
104 : with-mdb-connection ( quot -- )
105 [ mdb-pool get ] dip with-mdb-pool ; inline
107 : >id-selector ( assoc -- selector )
108 [ MDB_OID_FIELD of ] keep
109 H{ } clone [ set-at ] keep ;
111 : <mdb> ( db host port -- mdb )
112 <inet> t [ <mdb-node> ] keep
113 H{ } clone [ set-at ] keep <mdb-db>
114 [ verify-nodes ] keep ;
116 GENERIC: create-collection ( name/collection -- )
118 M: string create-collection
119 <mdb-collection> create-collection ;
121 M: mdb-collection create-collection ( collection -- )
122 create-cmd make-cmd over
124 [ name>> "create" set-cmd-opt ]
125 [ capped>> [ "capped" set-cmd-opt ] when* ]
126 [ max>> [ "max" set-cmd-opt ] when* ]
127 [ size>> [ "size" set-cmd-opt ] when* ]
128 } cleave send-cmd check-ok
129 [ drop [ ] [ name>> ] bi mdb-instance collections>> set-at ]
132 : load-collection-list ( -- collection-list )
133 namespaces-collection
134 H{ } clone <mdb-query-msg> send-query-plain objects>> ;
138 : ensure-valid-collection-name ( collection -- )
140 [ ";$." intersect length 0 > ] keep
141 '[ _ "contains invalid characters ( . $ ; )" ":" glue throw ] when
143 [ ascii? ] all? [ "collection names must only contain ascii characters" throw ] unless
146 : build-collection-map ( -- assoc )
147 H{ } clone load-collection-list
148 [ [ "name" ] dip at "." split second <mdb-collection> ] map
149 over '[ [ ] [ name>> ] bi _ set-at ] each ;
151 : ensure-collection-map ( mdb-instance -- assoc )
152 dup collections>> dup assoc-empty?
153 [ drop build-collection-map [ >>collections drop ] keep ]
156 : (ensure-collection) ( collection mdb-instance -- collection )
157 ensure-collection-map [ dup ] dip key?
158 [ ] [ [ ensure-valid-collection-name ]
159 [ create-collection ]
162 : reserved-namespace? ( name -- ? )
163 [ "$cmd" = ] [ "system" head? ] bi or ;
165 : check-collection ( collection -- fq-collection )
167 mdb-instance :> instance
168 instance name>> :> instance-name
169 dup mdb-collection? [ name>> ] when
170 "." split1 over instance-name =
172 [ ] [ reserved-namespace? ] bi
173 [ instance (ensure-collection) ] unless
174 [ instance-name ] dip "." glue
177 : fix-query-collection ( mdb-query -- mdb-query )
178 [ check-collection ] change-collection ; inline
180 : get-more ( mdb-cursor -- mdb-cursor seq )
181 [ [ query>> dup [ collection>> ] [ return#>> ] bi ]
182 [ id>> ] bi <mdb-getmore-msg> swap >>query send-query ]
187 : <query> ( collection assoc -- mdb-query-msg )
188 <mdb-query-msg> ; inline
190 : >slave-ok ( mdb-query-msg -- mdb-query-msg )
191 [ 2 set-bit ] change-flags ;
193 : >await-data ( mdb-query-msg -- mdb-query-msg )
194 [ 5 set-bit ] change-flags ;
196 : >tailable ( mdb-query-msg -- mdb-query-msg )
197 [ 1 set-bit ] change-flags ;
199 : limit ( mdb-query-msg limit# -- mdb-query-msg )
202 : skip ( mdb-query-msg skip# -- mdb-query-msg )
205 : asc ( key -- spec ) 1 2array ; inline
206 : desc ( key -- spec ) -1 2array ; inline
208 : sort ( mdb-query-msg sort-quot -- mdb-query-msg )
209 output>array >hashtable >>orderby ; inline
211 : filter-fields ( mdb-query-msg filterseq -- mdb-query-msg )
212 [ asc ] map >hashtable >>returnfields ; inline
214 : key-spec ( spec-quot -- spec-assoc )
215 output>array >hashtable ; inline
217 GENERIC# hint 1 ( mdb-query-msg index-hint -- mdb-query-msg )
219 M: mdb-query-msg hint
222 GENERIC: find ( selector -- mdb-cursor/f seq )
224 M: mdb-query-msg find
225 fix-query-collection send-query ;
230 : each-chunk ( selector quot: ( seq -- ) -- )
232 [ pick call( seq -- ) ] when*
233 [ swap each-chunk ] [ drop ] if* ;
235 : find-all ( selector -- seq )
237 over '[ _ push-all ] each-chunk >array ;
239 : explain. ( mdb-query-msg -- )
240 t >>explain find nip . ;
242 : find-one ( mdb-query-msg -- result/f )
243 fix-query-collection 1 >>return#
244 send-query-plain objects>> ?first ;
246 : count ( mdb-query-msg -- result )
247 [ count-cmd make-cmd ] dip
248 [ collection>> "count" set-cmd-opt ]
249 [ query>> "query" set-cmd-opt ] bi send-cmd
250 [ check-ok nip ] keep '[ "n" _ at >fixnum ] [ f ] if ;
252 : lasterror ( -- error )
253 getlasterror-cmd make-cmd send-cmd
256 GENERIC: validate. ( collection -- )
259 [ validate-cmd make-cmd ] dip
260 "validate" set-cmd-opt send-cmd
261 [ check-ok nip ] keep
262 '[ "result" _ at print ] [ ] if ;
264 M: mdb-collection validate.
269 : send-message-check-error ( message -- )
270 send-message lasterror [ mdb-error ] when* ;
274 : save ( collection assoc -- )
275 [ check-collection ] dip
276 <mdb-insert-msg> send-message-check-error ;
278 : save-unsafe ( collection assoc -- )
279 [ check-collection ] dip
280 <mdb-insert-msg> send-message ;
282 : ensure-index ( index-spec -- )
283 <linked-hash> [ [ <oid> "_id" ] dip set-at ] keep
284 [ { [ [ name>> "name" ] dip set-at ]
285 [ [ ns>> index-ns "ns" ] dip set-at ]
286 [ [ key>> "key" ] dip set-at ]
288 [ swap [ "unique" ] dip set-at ] [ drop ] if* ] } 2cleave
290 [ index-collection ] dip save ;
292 : drop-index ( collection name -- )
293 [ delete-index-cmd make-cmd ] 2dip
294 [ "deleteIndexes" set-cmd-opt ]
295 [ "index" set-cmd-opt ] bi* send-cmd drop ;
297 : <update> ( collection selector object -- mdb-update-msg )
298 [ check-collection ] 2dip <mdb-update-msg> ;
300 : >upsert ( mdb-update-msg -- mdb-update-msg )
301 [ 0 set-bit ] change-update-flags ;
303 : >multi ( mdb-update-msg -- mdb-update-msg )
304 [ 1 set-bit ] change-update-flags ;
306 : update ( mdb-update-msg -- )
307 send-message-check-error ;
309 : update-unsafe ( mdb-update-msg -- )
312 : find-and-modify ( collection selector modifier -- mongodb-cmd )
313 [ findandmodify-cmd make-cmd ] 3dip
314 [ "findandmodify" set-cmd-opt ]
315 [ "query" set-cmd-opt ]
316 [ "update" set-cmd-opt ] tri* ; inline
318 : run-cmd ( cmd -- result )
321 : <delete> ( collection selector -- mdb-delete-msg )
322 [ check-collection ] dip <mdb-delete-msg> ;
324 : >single-remove ( mdb-delete-msg -- mdb-delete-msg )
325 [ 0 set-bit ] change-delete-flags ;
327 : delete ( mdb-delete-msg -- )
328 send-message-check-error ;
330 : delete-unsafe ( mdb-delete-msg -- )
333 : kill-cursor ( mdb-cursor -- )
334 id>> <mdb-killcursors-msg> send-message ;
336 : load-index-list ( -- index-list )
338 H{ } clone <mdb-query-msg> find nip ;
340 : ensure-collection ( name -- )
341 check-collection drop ;
343 : drop-collection ( name -- )
344 [ drop-cmd make-cmd ] dip
345 "drop" set-cmd-opt send-cmd drop ;