extra/mongodb/operations/operations.factor

   1 USING: accessors assocs bson.reader bson.writer byte-arrays
   2 byte-vectors combinators formatting fry io io.binary io.encodings.private
   3 io.encodings.binary io.encodings.string io.encodings.utf8 io.encodings.utf8.private io.files
   4 kernel locals math mongodb.msg namespaces sequences uuid bson.writer.private ;
   5
   6 IN: alien.c-types
   7
   8 M: byte-vector byte-length length ;
   9
  10 IN: mongodb.operations
  11
  12 <PRIVATE
  13
  14 PREDICATE: mdb-reply-op < integer OP_Reply = ;
  15 PREDICATE: mdb-query-op < integer OP_Query = ;
  16 PREDICATE: mdb-insert-op < integer OP_Insert = ;
  17 PREDICATE: mdb-update-op < integer OP_Update = ;
  18 PREDICATE: mdb-delete-op < integer OP_Delete = ;
  19 PREDICATE: mdb-getmore-op < integer OP_GetMore = ;
  20 PREDICATE: mdb-killcursors-op < integer OP_KillCursors = ;
  21
  22 PRIVATE>
  23
  24 GENERIC: write-message ( message -- )
  25
  26 <PRIVATE
  27
  28 CONSTANT: MSG-HEADER-SIZE 16
  29
  30 SYMBOL: msg-bytes-read
  31
  32 : bytes-read> ( -- integer )
  33     msg-bytes-read get ; inline
  34
  35 : >bytes-read ( integer -- )
  36     msg-bytes-read set ; inline
  37
  38 : change-bytes-read ( integer -- )
  39     bytes-read> [ 0 ] unless* + >bytes-read ; inline
  40
  41 : read-int32 ( -- int32 ) 4 [ read le> ] [ change-bytes-read ] bi ; inline
  42 : read-longlong ( -- longlong ) 8 [ read le> ] [ change-bytes-read ] bi ; inline
  43 : read-byte-raw ( -- byte-raw ) 1 [ read le> ] [ change-bytes-read ] bi ; inline
  44 : read-byte ( -- byte ) read-byte-raw first ; inline
  45
  46 : (read-cstring) ( acc -- )
  47     [ read-byte ] dip ! b acc
  48     2dup push             ! b acc
  49     [ 0 = ] dip      ! bool acc
  50     '[ _ (read-cstring) ] unless ; inline recursive
  51
  52 : read-cstring ( -- string )
  53     BV{ } clone
  54     [ (read-cstring) ] keep
  55     [ zero? ] trim-tail
  56     >byte-array utf8 decode ; inline
  57
  58 GENERIC: (read-message) ( message opcode -- message )
  59
  60 : copy-header ( message msg-stub -- message )
  61     [ length>> ] keep [ >>length ] dip
  62     [ req-id>> ] keep [ >>req-id ] dip
  63     [ resp-id>> ] keep [ >>resp-id ] dip
  64     [ opcode>> ] keep [ >>opcode ] dip
  65     flags>> >>flags ;
  66
  67 M: mdb-reply-op (read-message) ( msg-stub opcode -- message )
  68     drop
  69     [ <mdb-reply-msg> ] dip copy-header
  70     read-longlong >>cursor
  71     read-int32 >>start#
  72     read-int32 [ >>returned# ] keep
  73     [ H{ } stream>assoc ] accumulator [ times ] dip >>objects ;
  74
  75 : read-header ( message -- message )
  76     read-int32 >>length
  77     read-int32 >>req-id
  78     read-int32 >>resp-id
  79     read-int32 >>opcode
  80     read-int32 >>flags ; inline
  81
  82 : write-header ( message -- )
  83     [ req-id>> write-int32 ] keep
  84     [ resp-id>> write-int32 ] keep
  85     opcode>> write-int32 ; inline
  86
  87 PRIVATE>
  88
  89 : read-message ( -- message )
  90     mdb-msg new
  91     0 >bytes-read
  92     read-header
  93     [ ] [ opcode>> ] bi (read-message) ;
  94
  95 <PRIVATE
  96
  97 USE: tools.walker
  98
  99 : dump-to-file ( array -- )
 100     [ uuid1 "/tmp/mfb/%s.dump" sprintf binary ] dip
 101     '[ _ write ] with-file-writer ;
 102
 103 : (write-message) ( message quot -- )
 104     '[ [ [ _ write-header ] dip _ call ] with-length-prefix ] with-buffer
 105     ! [ dump-to-file ] keep
 106     write flush ; inline
 107
 108 :: build-query-object ( query -- selector )
 109     H{ } clone :> selector
 110     query { [ orderby>> [ "$orderby" selector set-at ] when* ]
 111       [ explain>> [ "$explain" selector set-at ] when* ]
 112       [ hint>> [ "$hint" selector set-at ] when* ]
 113       [ query>> "query" selector set-at ]
 114     } cleave
 115     selector ;
 116
 117 PRIVATE>
 118
 119 M: mdb-query-msg write-message ( message -- )
 120      dup
 121      '[ _
 122         [ flags>> write-int32 ] keep
 123         [ collection>> write-cstring ] keep
 124         [ skip#>> write-int32 ] keep
 125         [ return#>> write-int32 ] keep
 126         [ build-query-object assoc>stream ] keep
 127         returnfields>> [ assoc>stream ] when*
 128      ] (write-message) ;
 129
 130 M: mdb-insert-msg write-message ( message -- )
 131     dup
 132     '[ _
 133        [ flags>> write-int32 ] keep
 134        [ collection>> write-cstring ] keep
 135        objects>> [ assoc>stream ] each
 136     ] (write-message) ;
 137
 138 M: mdb-update-msg write-message ( message -- )
 139     dup
 140     '[ _
 141        [ flags>> write-int32 ] keep
 142        [ collection>> write-cstring ] keep
 143        [ upsert?>> write-int32 ] keep
 144        [ selector>> assoc>stream ] keep
 145        object>> assoc>stream
 146     ] (write-message) ;
 147
 148 M: mdb-delete-msg write-message ( message -- )
 149     dup
 150     '[ _
 151        [ flags>> write-int32 ] keep
 152        [ collection>> write-cstring ] keep
 153        0 write-int32
 154        selector>> assoc>stream
 155     ] (write-message) ;
 156
 157 M: mdb-getmore-msg write-message ( message -- )
 158     dup
 159     '[ _
 160        [ flags>> write-int32 ] keep
 161        [ collection>> write-cstring ] keep
 162        [ return#>> write-int32 ] keep
 163        cursor>> write-longlong
 164     ] (write-message) ;
 165
 166 M: mdb-killcursors-msg write-message ( message -- )
 167     dup
 168     '[ _
 169        [ flags>> write-int32 ] keep
 170        [ cursors#>> write-int32 ] keep
 171        cursors>> [ write-longlong ] each
 172     ] (write-message) ;
 173