extra/mongodb/driver/driver.factor

   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 ;
   9 IN: mongodb.driver
  10
  11 TUPLE: mdb-pool < pool mdb ;
  12
  13 TUPLE: mdb-cursor id query ;
  14
  15 TUPLE: mdb-collection
  16 { name string }
  17 { capped boolean }
  18 { size integer }
  19 { max integer } ;
  20
  21 CONSTRUCTOR: <mdb-collection> mdb-collection ( name -- collection ) ;
  22
  23 TUPLE: index-spec
  24 { ns string } { name string } { key hashtable } { unique? boolean initial: f } ;
  25
  26 CONSTRUCTOR: <index-spec> index-spec ( ns name key -- index-spec ) ;
  27
  28 M: mdb-pool make-connection
  29     mdb>> mdb-open ;
  30
  31 : <mdb-pool> ( mdb -- pool ) [ mdb-pool <pool> ] dip >>mdb ; inline
  32
  33 CONSTANT: PARTIAL? "partial?"
  34
  35 ERROR: mdb-error msg ;
  36
  37 M: mdb-error pprint* ( obj -- )
  38     msg>> text ;
  39
  40 : >pwd-digest ( user password -- digest )
  41     "mongo" swap 3array ":" join md5-checksum ;
  42
  43 <PRIVATE
  44
  45 GENERIC: <mdb-cursor> ( id mdb-query-msg/mdb-getmore-msg -- mdb-cursor )
  46
  47 M: mdb-query-msg <mdb-cursor>
  48     mdb-cursor boa ;
  49
  50 M: mdb-getmore-msg <mdb-cursor>
  51     query>> mdb-cursor boa ;
  52
  53 : >mdbregexp ( value -- regexp )
  54    first <mdbregexp> ; inline
  55
  56 GENERIC: update-query ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- )
  57
  58 M: mdb-query-msg update-query
  59     swap [ start#>> ] [ returned#>> ] bi + >>skip# drop ;
  60
  61 M: mdb-getmore-msg update-query
  62     query>> update-query ;
  63
  64 : make-cursor ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- mdb-cursor/f )
  65     over cursor>> 0 >
  66     [ [ update-query ]
  67       [ [ cursor>> ] dip <mdb-cursor> ] 2bi
  68     ] [ 2drop f ] if ;
  69
  70 DEFER: send-query
  71
  72 GENERIC: verify-query-result ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- mdb-result-msg mdb-query-msg/mdb-getmore-msg )
  73
  74 M: mdb-query-msg verify-query-result ;
  75
  76 M: mdb-getmore-msg verify-query-result
  77     over flags>> ResultFlag_CursorNotFound =
  78     [ nip query>> [ send-query-plain ] keep ] when ;
  79
  80 : send-query ( mdb-query-msg/mdb-getmore-msg -- mdb-cursor/f seq )
  81     [ send-query-plain ] keep
  82     verify-query-result
  83     [ collection>> >>collection drop ]
  84     [ return#>> >>requested# ]
  85     [ make-cursor ] 2tri
  86     swap objects>> ;
  87
  88
  89 PRIVATE>
  90
  91 SYNTAX: r/
  92     \ / [ >mdbregexp ] parse-literal ;
  93
  94 : with-db ( mdb quot -- )
  95     '[ _ mdb-open &dispose _ with-connection ] with-destructors ; inline
  96
  97 : with-mdb ( mdb quot -- )
  98     [ <mdb-pool> ] dip
  99     [ mdb-pool swap with-variable ] curry with-disposal ; inline
 100
 101 : with-mdb-pool ( ..a mdb-pool quot -- ..b )
 102     '[ _ with-connection ] with-pooled-connection ; inline
 103
 104 : with-mdb-connection ( quot -- )
 105     [ mdb-pool get ] dip with-mdb-pool ; inline
 106
 107 : >id-selector ( assoc -- selector )
 108     [ MDB_OID_FIELD of ] keep
 109     H{ } clone [ set-at ] keep ;
 110
 111 : <mdb> ( db host port -- mdb )
 112     <inet> t [ <mdb-node> ] keep
 113     H{ } clone [ set-at ] keep <mdb-db>
 114     [ verify-nodes ] keep ;
 115
 116 GENERIC: create-collection ( name/collection -- )
 117
 118 M: string create-collection
 119     <mdb-collection> create-collection ;
 120
 121 M: mdb-collection create-collection ( collection -- )
 122     create-cmd make-cmd over
 123     {
 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 ]
 130     [ throw ] if ;
 131
 132 : load-collection-list ( -- collection-list )
 133     namespaces-collection
 134     H{ } clone <mdb-query-msg> send-query-plain objects>> ;
 135
 136 <PRIVATE
 137
 138 : ensure-valid-collection-name ( collection -- )
 139     [
 140         [ ";$." intersect length 0 > ] keep
 141         '[ _ "contains invalid characters ( . $ ; )" ":" glue throw ] when
 142     ] [
 143         [ ascii? ] all? [ "collection names must only contain ascii characters" throw ] unless
 144     ] bi ; inline
 145
 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 ;
 150
 151 : ensure-collection-map ( mdb-instance -- assoc )
 152     dup collections>> dup assoc-empty?
 153     [ drop build-collection-map [ >>collections drop ] keep ]
 154     [ nip ] if ;
 155
 156 : (ensure-collection) ( collection mdb-instance -- collection )
 157     ensure-collection-map [ dup ] dip key?
 158     [ ] [
 159         [ ensure-valid-collection-name ]
 160         [ create-collection ]
 161         [ ] tri
 162     ] if ;
 163
 164 : reserved-namespace? ( name -- ? )
 165     [ "$cmd" = ] [ "system" head? ] bi or ;
 166
 167 : check-collection ( collection -- fq-collection )
 168     [let
 169         mdb-instance :> instance
 170         instance name>> :> instance-name
 171         dup mdb-collection? [ name>> ] when
 172         "." split1 over instance-name =
 173         [ nip ] [ drop ] if
 174         [ ] [ reserved-namespace? ] bi
 175         [ instance (ensure-collection) ] unless
 176         [ instance-name ] dip "." glue
 177     ] ;
 178
 179 : fix-query-collection ( mdb-query -- mdb-query )
 180     [ check-collection ] change-collection ; inline
 181
 182 : get-more ( mdb-cursor -- mdb-cursor seq )
 183     [ [ query>> dup [ collection>> ] [ return#>> ] bi ]
 184       [ id>> ] bi <mdb-getmore-msg> swap >>query send-query ]
 185     [ f f ] if* ;
 186
 187 PRIVATE>
 188
 189 : <query> ( collection assoc -- mdb-query-msg )
 190     <mdb-query-msg> ; inline
 191
 192 : >slave-ok ( mdb-query-msg -- mdb-query-msg )
 193     [ 2 set-bit ] change-flags ;
 194
 195 : >await-data ( mdb-query-msg -- mdb-query-msg )
 196     [ 5 set-bit ] change-flags ;
 197
 198 : >tailable ( mdb-query-msg -- mdb-query-msg )
 199     [ 1 set-bit ] change-flags ;
 200
 201 : limit ( mdb-query-msg limit# -- mdb-query-msg )
 202     >>return# ; inline
 203
 204 : skip ( mdb-query-msg skip# -- mdb-query-msg )
 205     >>skip# ; inline
 206
 207 : asc ( key -- spec ) 1 2array ; inline
 208 : desc ( key -- spec ) -1 2array ; inline
 209
 210 : sort ( mdb-query-msg sort-quot -- mdb-query-msg )
 211     output>array >hashtable >>orderby ; inline
 212
 213 : filter-fields ( mdb-query-msg filterseq -- mdb-query-msg )
 214     [ asc ] map >hashtable >>returnfields ; inline
 215
 216 : key-spec ( spec-quot -- spec-assoc )
 217     output>array >hashtable ; inline
 218
 219 GENERIC#: hint 1 ( mdb-query-msg index-hint -- mdb-query-msg )
 220
 221 M: mdb-query-msg hint
 222     >>hint ;
 223
 224 GENERIC: find ( selector -- mdb-cursor/f seq )
 225
 226 M: mdb-query-msg find
 227     fix-query-collection send-query ;
 228
 229 M: mdb-cursor find
 230     get-more ;
 231
 232 : each-chunk ( selector quot: ( seq -- ) -- )
 233     swap find
 234     [ pick call( seq -- ) ] when*
 235     [ swap each-chunk ] [ drop ] if* ;
 236
 237 : find-all ( selector -- seq )
 238     [ V{ } clone ] dip
 239     over '[ _ push-all ] each-chunk >array ;
 240
 241 : explain. ( mdb-query-msg -- )
 242     t >>explain find nip . ;
 243
 244 : find-one ( mdb-query-msg -- result/f )
 245     fix-query-collection 1 >>return#
 246     send-query-plain objects>> ?first ;
 247
 248 : count ( mdb-query-msg -- result )
 249     [ count-cmd make-cmd ] dip
 250     [ collection>> "count" set-cmd-opt ]
 251     [ query>> "query" set-cmd-opt ] bi send-cmd
 252     [ check-ok nip ] keep '[ "n" _ at >fixnum ] [ f ] if ;
 253
 254 : lasterror ( -- error )
 255     getlasterror-cmd make-cmd send-cmd
 256     [ "err" ] dip at ;
 257
 258 GENERIC: validate. ( collection -- )
 259
 260 M: string validate.
 261     [ validate-cmd make-cmd ] dip
 262     "validate" set-cmd-opt send-cmd
 263     [ check-ok nip ] keep
 264     '[ "result" _ at print ] [  ] if ;
 265
 266 M: mdb-collection validate.
 267     name>> validate. ;
 268
 269 <PRIVATE
 270
 271 : send-message-check-error ( message -- )
 272     send-message lasterror [ mdb-error ] when* ;
 273
 274 PRIVATE>
 275
 276 : save ( collection assoc -- )
 277     [ check-collection ] dip
 278     <mdb-insert-msg> send-message-check-error ;
 279
 280 : save-unsafe ( collection assoc -- )
 281     [ check-collection ] dip
 282     <mdb-insert-msg> send-message ;
 283
 284 : ensure-index ( index-spec -- )
 285     <linked-hash> [ [ <oid> "_id" ] dip set-at ] keep
 286     [ { [ [ name>> "name" ] dip set-at ]
 287         [ [ ns>> index-ns "ns" ] dip set-at ]
 288         [ [ key>> "key" ] dip set-at ]
 289         [ swap unique?>>
 290           [ swap [ "unique" ] dip set-at ] [ drop ] if* ] } 2cleave
 291     ] keep
 292     [ index-collection ] dip save ;
 293
 294 : drop-index ( collection name -- )
 295     [ delete-index-cmd make-cmd ] 2dip
 296     [ "deleteIndexes" set-cmd-opt ]
 297     [ "index" set-cmd-opt ] bi* send-cmd drop ;
 298
 299 : <update> ( collection selector object -- mdb-update-msg )
 300     [ check-collection ] 2dip <mdb-update-msg> ;
 301
 302 : >upsert ( mdb-update-msg -- mdb-update-msg )
 303     [ 0 set-bit ] change-update-flags ;
 304
 305 : >multi ( mdb-update-msg -- mdb-update-msg )
 306     [ 1 set-bit ] change-update-flags ;
 307
 308 : update ( mdb-update-msg -- )
 309     send-message-check-error ;
 310
 311 : update-unsafe ( mdb-update-msg -- )
 312     send-message ;
 313
 314 : find-and-modify ( collection selector modifier -- mongodb-cmd )
 315     [ findandmodify-cmd make-cmd ] 3dip
 316     [ "findandmodify" set-cmd-opt ]
 317     [ "query" set-cmd-opt ]
 318     [ "update" set-cmd-opt ] tri* ; inline
 319
 320 : run-cmd ( cmd -- result )
 321     send-cmd ; inline
 322
 323 : <delete> ( collection selector -- mdb-delete-msg )
 324     [ check-collection ] dip <mdb-delete-msg> ;
 325
 326 : >single-remove ( mdb-delete-msg -- mdb-delete-msg )
 327     [ 0 set-bit ] change-delete-flags ;
 328
 329 : delete ( mdb-delete-msg -- )
 330     send-message-check-error ;
 331
 332 : delete-unsafe ( mdb-delete-msg -- )
 333     send-message ;
 334
 335 : kill-cursor ( mdb-cursor -- )
 336     id>> <mdb-killcursors-msg> send-message ;
 337
 338 : load-index-list ( -- index-list )
 339     index-collection
 340     H{ } clone <mdb-query-msg> find nip ;
 341
 342 : ensure-collection ( name -- )
 343     check-collection drop ;
 344
 345 : drop-collection ( name -- )
 346     [ drop-cmd make-cmd ] dip
 347     "drop" set-cmd-opt send-cmd drop ;