extra/mongodb/connection/connection.factor

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