extra/mongodb/connection/connection.factor

   1 USING: accessors assocs fry io.encodings.binary io.sockets kernel math
   2 math.parser mongodb.msg mongodb.operations namespaces destructors
   3 constructors sequences splitting checksums checksums.md5 formatting
   4 io.streams.duplex io.encodings.utf8 io.encodings.string combinators.smart
   5 arrays hashtables sequences.deep vectors locals ;
   6
   7 IN: mongodb.connection
   8
   9 : md5-checksum ( string -- digest )
  10     utf8 encode md5 checksum-bytes hex-string ; inline
  11
  12 TUPLE: mdb-db name username pwd-digest nodes collections ;
  13
  14 TUPLE: mdb-node master? { address inet } remote ;
  15
  16 CONSTRUCTOR: mdb-node ( address master? -- mdb-node ) ;
  17
  18 TUPLE: mdb-connection instance node handle remote local ;
  19
  20 CONSTRUCTOR: mdb-connection ( instance -- mdb-connection ) ;
  21
  22 : check-ok ( result -- errmsg ? )
  23     [ [ "errmsg" ] dip at ]
  24     [ [ "ok" ] dip at >integer 1 = ] bi ; inline
  25
  26 : <mdb-db> ( name nodes -- mdb-db )
  27     mdb-db new swap >>nodes swap >>name H{ } clone >>collections ;
  28
  29 : master-node ( mdb -- node )
  30     nodes>> t swap at ;
  31
  32 : slave-node ( mdb -- node )
  33     nodes>> f swap at ;
  34
  35 : with-connection ( connection quot -- * )
  36     [ mdb-connection set ] prepose with-scope ; inline
  37
  38 : mdb-instance ( -- mdb )
  39     mdb-connection get instance>> ; inline
  40
  41 : index-collection ( -- ns )
  42     mdb-instance name>> "%s.system.indexes" sprintf ; inline
  43
  44 : namespaces-collection ( -- ns )
  45     mdb-instance name>> "%s.system.namespaces" sprintf ; inline
  46
  47 : cmd-collection ( -- ns )
  48     mdb-instance name>> "%s.$cmd" sprintf ; inline
  49
  50 : index-ns ( colname -- index-ns )
  51     [ mdb-instance name>> ] dip "%s.%s" sprintf ; inline
  52
  53 : send-message ( message -- )
  54     [ mdb-connection get handle>> ] dip '[ _ write-message ] with-stream* ;
  55
  56 : send-query-plain ( query-message -- result )
  57     [ mdb-connection get handle>> ] dip
  58     '[ _ write-message read-message ] with-stream* ;
  59
  60 : send-query-1result ( collection assoc -- result )
  61     <mdb-query-msg>
  62         1 >>return#
  63     send-query-plain objects>>
  64     [ f ] [ first ] if-empty ;
  65
  66 <PRIVATE
  67
  68 : get-nonce ( -- nonce )
  69     cmd-collection H{ { "getnonce" 1 } } send-query-1result
  70     [ "nonce" swap at ] [ f ] if* ;
  71
  72 : auth? ( mdb -- ? )
  73     [ username>> ] [ pwd-digest>> ] bi and ;
  74
  75 : calculate-key-digest ( nonce -- digest )
  76     mdb-instance
  77     [ username>> ]
  78     [ pwd-digest>> ] bi
  79     3array concat md5-checksum ; inline
  80
  81 : build-auth-query ( -- query-assoc )
  82     { "authenticate" 1 }
  83     "user"  mdb-instance username>> 2array
  84     "nonce" get-nonce 2array
  85     3array >hashtable
  86     [ [ "nonce" ] dip at calculate-key-digest "key" ] keep
  87     [ set-at ] keep ;
  88
  89 : perform-authentication ( --  )
  90     cmd-collection build-auth-query send-query-1result
  91     check-ok [ drop ] [ throw ] if ; inline
  92
  93 : authenticate-connection ( mdb-connection -- )
  94    [ mdb-connection get instance>> auth?
  95      [ perform-authentication ] when
  96    ] with-connection ; inline
  97
  98 : open-connection ( mdb-connection node -- mdb-connection )
  99    [ >>node ] [ address>> ] bi
 100    [ >>remote ] keep binary <client>
 101    [ >>handle ] dip >>local ;
 102
 103 : get-ismaster ( -- result )
 104     "admin.$cmd" H{ { "ismaster" 1 } } send-query-1result ;
 105
 106 : split-host-str ( hoststr -- host port )
 107    ":" split [ first ] [ second string>number ] bi ; inline
 108
 109 : eval-ismaster-result ( node result -- )
 110    [ [ "ismaster" ] dip at >integer 1 = >>master? drop ]
 111    [ [ "remote" ] dip at
 112      [ split-host-str <inet> f <mdb-node> >>remote ] when*
 113      drop ] 2bi ;
 114
 115 : check-node ( mdb node --  )
 116    [ <mdb-connection> &dispose ] dip
 117    [ open-connection ] keep swap
 118    [ get-ismaster eval-ismaster-result ] with-connection ;
 119
 120 : nodelist>table ( seq -- assoc )
 121    [ [ master?>> ] keep 2array ] map >hashtable ;
 122
 123 PRIVATE>
 124
 125 :: verify-nodes ( mdb -- )
 126     [ [let* | acc [ V{ } clone ]
 127               node1 [ mdb dup master-node [ check-node ] keep ]
 128               node2 [ mdb node1 remote>>
 129                       [ [ check-node ] keep ]
 130                       [ drop f ] if*  ]
 131               | node1 [ acc push ] when*
 132                 node2 [ acc push ] when*
 133                 mdb acc nodelist>table >>nodes drop
 134               ]
 135     ] with-destructors ;
 136
 137 : mdb-open ( mdb -- mdb-connection )
 138     clone [ <mdb-connection> ] keep
 139     master-node open-connection
 140     [ authenticate-connection ] keep ;
 141
 142 : mdb-close ( mdb-connection -- )
 143      [ dispose f ] change-handle drop ;
 144
 145 M: mdb-connection dispose
 146      mdb-close ;