extra/mongodb/driver/driver.factor

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