extra/mongodb/driver/driver.factor

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