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     [ ] [ [ ensure-valid-collection-name ]
 159           [ create-collection ]
 160           [ ] tri ] if ;
 161
 162 : reserved-namespace? ( name -- ? )
 163     [ "$cmd" = ] [ "system" head? ] bi or ;
 164
 165 : check-collection ( collection -- fq-collection )
 166     [let
 167         mdb-instance :> instance
 168         instance name>> :> instance-name
 169         dup mdb-collection? [ name>> ] when
 170         "." split1 over instance-name =
 171         [ nip ] [ drop ] if
 172         [ ] [ reserved-namespace? ] bi
 173         [ instance (ensure-collection) ] unless
 174         [ instance-name ] dip "." glue
 175     ] ;
 176
 177 : fix-query-collection ( mdb-query -- mdb-query )
 178     [ check-collection ] change-collection ; inline
 179
 180 : get-more ( mdb-cursor -- mdb-cursor seq )
 181     [ [ query>> dup [ collection>> ] [ return#>> ] bi ]
 182       [ id>> ] bi <mdb-getmore-msg> swap >>query send-query ]
 183     [ f f ] if* ;
 184
 185 PRIVATE>
 186
 187 : <query> ( collection assoc -- mdb-query-msg )
 188     <mdb-query-msg> ; inline
 189
 190 : >slave-ok ( mdb-query-msg -- mdb-query-msg )
 191     [ 2 set-bit ] change-flags ;
 192
 193 : >await-data ( mdb-query-msg -- mdb-query-msg )
 194     [ 5 set-bit ] change-flags ;
 195
 196 : >tailable ( mdb-query-msg -- mdb-query-msg )
 197     [ 1 set-bit ] change-flags ;
 198
 199 : limit ( mdb-query-msg limit# -- mdb-query-msg )
 200     >>return# ; inline
 201
 202 : skip ( mdb-query-msg skip# -- mdb-query-msg )
 203     >>skip# ; inline
 204
 205 : asc ( key -- spec ) 1 2array ; inline
 206 : desc ( key -- spec ) -1 2array ; inline
 207
 208 : sort ( mdb-query-msg sort-quot -- mdb-query-msg )
 209     output>array >hashtable >>orderby ; inline
 210
 211 : filter-fields ( mdb-query-msg filterseq -- mdb-query-msg )
 212     [ asc ] map >hashtable >>returnfields ; inline
 213
 214 : key-spec ( spec-quot -- spec-assoc )
 215     output>array >hashtable ; inline
 216
 217 GENERIC#: hint 1 ( mdb-query-msg index-hint -- mdb-query-msg )
 218
 219 M: mdb-query-msg hint
 220     >>hint ;
 221
 222 GENERIC: find ( selector -- mdb-cursor/f seq )
 223
 224 M: mdb-query-msg find
 225     fix-query-collection send-query ;
 226
 227 M: mdb-cursor find
 228     get-more ;
 229
 230 : each-chunk ( selector quot: ( seq -- ) -- )
 231     swap find
 232     [ pick call( seq -- ) ] when*
 233     [ swap each-chunk ] [ drop ] if* ;
 234
 235 : find-all ( selector -- seq )
 236     [ V{ } clone ] dip
 237     over '[ _ push-all ] each-chunk >array ;
 238
 239 : explain. ( mdb-query-msg -- )
 240     t >>explain find nip . ;
 241
 242 : find-one ( mdb-query-msg -- result/f )
 243     fix-query-collection 1 >>return#
 244     send-query-plain objects>> ?first ;
 245
 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 ;
 251
 252 : lasterror ( -- error )
 253     getlasterror-cmd make-cmd send-cmd
 254     [ "err" ] dip at ;
 255
 256 GENERIC: validate. ( collection -- )
 257
 258 M: string validate.
 259     [ validate-cmd make-cmd ] dip
 260     "validate" set-cmd-opt send-cmd
 261     [ check-ok nip ] keep
 262     '[ "result" _ at print ] [  ] if ;
 263
 264 M: mdb-collection validate.
 265     name>> validate. ;
 266
 267 <PRIVATE
 268
 269 : send-message-check-error ( message -- )
 270     send-message lasterror [ mdb-error ] when* ;
 271
 272 PRIVATE>
 273
 274 : save ( collection assoc -- )
 275     [ check-collection ] dip
 276     <mdb-insert-msg> send-message-check-error ;
 277
 278 : save-unsafe ( collection assoc -- )
 279     [ check-collection ] dip
 280     <mdb-insert-msg> send-message ;
 281
 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 ]
 287         [ swap unique?>>
 288           [ swap [ "unique" ] dip set-at ] [ drop ] if* ] } 2cleave
 289     ] keep
 290     [ index-collection ] dip save ;
 291
 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 ;
 296
 297 : <update> ( collection selector object -- mdb-update-msg )
 298     [ check-collection ] 2dip <mdb-update-msg> ;
 299
 300 : >upsert ( mdb-update-msg -- mdb-update-msg )
 301     [ 0 set-bit ] change-update-flags ;
 302
 303 : >multi ( mdb-update-msg -- mdb-update-msg )
 304     [ 1 set-bit ] change-update-flags ;
 305
 306 : update ( mdb-update-msg -- )
 307     send-message-check-error ;
 308
 309 : update-unsafe ( mdb-update-msg -- )
 310     send-message ;
 311
 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
 317
 318 : run-cmd ( cmd -- result )
 319     send-cmd ; inline
 320
 321 : <delete> ( collection selector -- mdb-delete-msg )
 322     [ check-collection ] dip <mdb-delete-msg> ;
 323
 324 : >single-remove ( mdb-delete-msg -- mdb-delete-msg )
 325     [ 0 set-bit ] change-delete-flags ;
 326
 327 : delete ( mdb-delete-msg -- )
 328     send-message-check-error ;
 329
 330 : delete-unsafe ( mdb-delete-msg -- )
 331     send-message ;
 332
 333 : kill-cursor ( mdb-cursor -- )
 334     id>> <mdb-killcursors-msg> send-message ;
 335
 336 : load-index-list ( -- index-list )
 337     index-collection
 338     H{ } clone <mdb-query-msg> find nip ;
 339
 340 : ensure-collection ( name -- )
 341     check-collection drop ;
 342
 343 : drop-collection ( name -- )
 344     [ drop-cmd make-cmd ] dip
 345     "drop" set-cmd-opt send-cmd drop ;