348 lines
		
	
	
		
			9.3 KiB
		
	
	
	
		
			Factor
		
	
	
			
		
		
	
	
			348 lines
		
	
	
		
			9.3 KiB
		
	
	
	
		
			Factor
		
	
	
| USING: accessors arrays assocs bson.constants combinators
 | |
| combinators.smart constructors destructors fry hashtables io
 | |
| io.pools io.sockets kernel linked-assocs locals math
 | |
| mongodb.cmd mongodb.connection mongodb.msg namespaces parser
 | |
| prettyprint prettyprint.custom prettyprint.sections sequences
 | |
| sets splitting strings ;
 | |
| FROM: ascii => ascii? ;
 | |
| FROM: math.bitwise => set-bit ;
 | |
| IN: mongodb.driver
 | |
| 
 | |
| TUPLE: mdb-pool < pool mdb ;
 | |
| 
 | |
| TUPLE: mdb-cursor id query ;
 | |
| 
 | |
| TUPLE: mdb-collection
 | |
| { name string }
 | |
| { capped boolean }
 | |
| { size integer }
 | |
| { max integer } ;
 | |
| 
 | |
| CONSTRUCTOR: mdb-collection ( name -- collection ) ;
 | |
| 
 | |
| TUPLE: index-spec
 | |
| { ns string } { name string } { key hashtable } { unique? boolean initial: f } ;
 | |
| 
 | |
| CONSTRUCTOR: index-spec ( ns name key -- index-spec ) ;
 | |
| 
 | |
| M: mdb-pool make-connection
 | |
|     mdb>> mdb-open ;
 | |
| 
 | |
| : <mdb-pool> ( mdb -- pool ) [ mdb-pool <pool> ] dip >>mdb ; inline
 | |
| 
 | |
| CONSTANT: PARTIAL? "partial?"
 | |
| 
 | |
| ERROR: mdb-error msg ;
 | |
| 
 | |
| M: mdb-error pprint* ( obj -- )
 | |
|     msg>> text ;
 | |
| 
 | |
| : >pwd-digest ( user password -- digest )
 | |
|     "mongo" swap 3array ":" join md5-checksum ; 
 | |
| 
 | |
| <PRIVATE
 | |
| 
 | |
| GENERIC: <mdb-cursor> ( id mdb-query-msg/mdb-getmore-msg -- mdb-cursor )
 | |
| 
 | |
| M: mdb-query-msg <mdb-cursor>
 | |
|     mdb-cursor boa ;
 | |
| 
 | |
| M: mdb-getmore-msg <mdb-cursor>
 | |
|     query>> mdb-cursor boa ;
 | |
| 
 | |
| : >mdbregexp ( value -- regexp )
 | |
|    first <mdbregexp> ; inline
 | |
| 
 | |
| GENERIC: update-query ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- )
 | |
| 
 | |
| M: mdb-query-msg update-query 
 | |
|     swap [ start#>> ] [ returned#>> ] bi + >>skip# drop ;
 | |
| 
 | |
| M: mdb-getmore-msg update-query
 | |
|     query>> update-query ; 
 | |
|       
 | |
| : make-cursor ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- mdb-cursor/f )
 | |
|     over cursor>> 0 >
 | |
|     [ [ update-query ]
 | |
|       [ [ cursor>> ] dip <mdb-cursor> ] 2bi
 | |
|     ] [ 2drop f ] if ;
 | |
| 
 | |
| DEFER: send-query
 | |
| 
 | |
| GENERIC: verify-query-result ( mdb-result-msg mdb-query-msg/mdb-getmore-msg -- mdb-result-msg mdb-query-msg/mdb-getmore-msg ) 
 | |
| 
 | |
| M: mdb-query-msg verify-query-result ;
 | |
| 
 | |
| M: mdb-getmore-msg verify-query-result
 | |
|     over flags>> ResultFlag_CursorNotFound =
 | |
|     [ nip query>> [ send-query-plain ] keep ] when ;
 | |
|     
 | |
| : send-query ( mdb-query-msg/mdb-getmore-msg -- mdb-cursor/f seq )
 | |
|     [ send-query-plain ] keep
 | |
|     verify-query-result 
 | |
|     [ collection>> >>collection drop ]
 | |
|     [ return#>> >>requested# ] 
 | |
|     [ make-cursor ] 2tri
 | |
|     swap objects>> ;
 | |
| 
 | |
| 
 | |
| PRIVATE>
 | |
| 
 | |
| SYNTAX: r/
 | |
|     \ / [ >mdbregexp ] parse-literal ; 
 | |
| 
 | |
| : with-db ( mdb quot -- )
 | |
|     '[ _ mdb-open &dispose _ with-connection ] with-destructors ; inline
 | |
| 
 | |
| : with-mdb ( mdb quot -- )
 | |
|     [ <mdb-pool> ] dip
 | |
|     [ mdb-pool swap with-variable ] curry with-disposal ; inline
 | |
| 
 | |
| : with-mdb-pool ( ..a mdb-pool quot -- ..b )
 | |
|     '[ _ with-connection ] with-pooled-connection ; inline
 | |
| 
 | |
| : with-mdb-connection ( quot -- )
 | |
|     [ mdb-pool get ] dip with-mdb-pool ; inline
 | |
| 
 | |
| : >id-selector ( assoc -- selector )
 | |
|     [ MDB_OID_FIELD swap at ] keep
 | |
|     H{ } clone [ set-at ] keep ;
 | |
| 
 | |
| : <mdb> ( db host port -- mdb )
 | |
|    <inet> t [ <mdb-node> ] keep
 | |
|    H{ } clone [ set-at ] keep <mdb-db>
 | |
|    [ verify-nodes ] keep ;
 | |
| 
 | |
| GENERIC: create-collection ( name/collection -- )
 | |
| 
 | |
| M: string create-collection
 | |
|     <mdb-collection> create-collection ;
 | |
| 
 | |
| M: mdb-collection create-collection ( collection -- )
 | |
|     create-cmd make-cmd over
 | |
|     {
 | |
|         [ name>> "create" set-cmd-opt ]
 | |
|         [ capped>> [ "capped" set-cmd-opt ] when* ]
 | |
|         [ max>> [ "max" set-cmd-opt ] when* ]
 | |
|         [ size>> [ "size" set-cmd-opt ] when* ]
 | |
|     } cleave send-cmd check-ok
 | |
|     [ drop [ ] [ name>> ] bi mdb-instance collections>> set-at ]
 | |
|     [ throw ] if ;
 | |
|   
 | |
| : load-collection-list ( -- collection-list )
 | |
|     namespaces-collection
 | |
|     H{ } clone <mdb-query-msg> send-query-plain objects>> ;
 | |
| 
 | |
| <PRIVATE
 | |
| 
 | |
| : ensure-valid-collection-name ( collection -- )
 | |
|     [
 | |
|         [ ";$." intersect length 0 > ] keep
 | |
|         '[ _ "contains invalid characters ( . $ ; )" ":" glue throw ] when
 | |
|     ] [
 | |
|         [ ascii? ] all? [ "collection names must only contain ascii characters" throw ] unless
 | |
|     ] bi ; inline
 | |
| 
 | |
| : build-collection-map ( -- assoc )
 | |
|     H{ } clone load-collection-list      
 | |
|     [ [ "name" ] dip at "." split second <mdb-collection> ] map
 | |
|     over '[ [ ] [ name>> ] bi _ set-at ] each ;
 | |
| 
 | |
| : ensure-collection-map ( mdb-instance -- assoc )
 | |
|     dup collections>> dup keys length 0 = 
 | |
|     [ drop build-collection-map [ >>collections drop ] keep ]
 | |
|     [ nip ] if ; 
 | |
| 
 | |
| : (ensure-collection) ( collection mdb-instance -- collection )
 | |
|     ensure-collection-map [ dup ] dip key?
 | |
|     [ ] [ [ ensure-valid-collection-name ]
 | |
|           [ create-collection ]
 | |
|           [ ] tri ] if ; 
 | |
|       
 | |
| : reserved-namespace? ( name -- ? )
 | |
|     [ "$cmd" = ] [ "system" head? ] bi or ;
 | |
| 
 | |
| : check-collection ( collection -- fq-collection )
 | |
|     [let
 | |
|         mdb-instance :> instance
 | |
|         instance name>> :> instance-name
 | |
|         dup mdb-collection? [ name>> ] when
 | |
|         "." split1 over instance-name =
 | |
|         [ nip ] [ drop ] if
 | |
|         [ ] [ reserved-namespace? ] bi
 | |
|         [ instance (ensure-collection) ] unless
 | |
|         [ instance-name ] dip "." glue
 | |
|     ] ;
 | |
| 
 | |
| : fix-query-collection ( mdb-query -- mdb-query )
 | |
|     [ check-collection ] change-collection ; inline
 | |
| 
 | |
| : get-more ( mdb-cursor -- mdb-cursor seq )
 | |
|     [ [ query>> dup [ collection>> ] [ return#>> ] bi ]
 | |
|       [ id>> ] bi <mdb-getmore-msg> swap >>query send-query ] 
 | |
|     [ f f ] if* ;
 | |
| 
 | |
| PRIVATE>
 | |
| 
 | |
| : <query> ( collection assoc -- mdb-query-msg )
 | |
|     <mdb-query-msg> ; inline
 | |
| 
 | |
| : >slave-ok ( mdb-query-msg -- mdb-query-msg )
 | |
|     [ 2 set-bit ] change-flags ;
 | |
| 
 | |
| : >await-data ( mdb-query-msg -- mdb-query-msg )
 | |
|     [ 5 set-bit ] change-flags ;
 | |
| 
 | |
| : >tailable ( mdb-query-msg -- mdb-query-msg )
 | |
|     [ 1 set-bit ] change-flags ;
 | |
| 
 | |
| : limit ( mdb-query-msg limit# -- mdb-query-msg )
 | |
|     >>return# ; inline
 | |
| 
 | |
| : skip ( mdb-query-msg skip# -- mdb-query-msg )
 | |
|     >>skip# ; inline
 | |
| 
 | |
| : asc ( key -- spec ) 1 2array ; inline
 | |
| : desc ( key -- spec ) -1 2array ; inline
 | |
| 
 | |
| : sort ( mdb-query-msg sort-quot -- mdb-query-msg )
 | |
|     output>array >hashtable >>orderby ; inline
 | |
| 
 | |
| : filter-fields ( mdb-query-msg filterseq -- mdb-query-msg )
 | |
|     [ asc ] map >hashtable >>returnfields ; inline
 | |
| 
 | |
| : key-spec ( spec-quot -- spec-assoc )
 | |
|     output>array >hashtable ; inline
 | |
| 
 | |
| GENERIC# hint 1 ( mdb-query-msg index-hint -- mdb-query-msg )
 | |
| 
 | |
| M: mdb-query-msg hint 
 | |
|     >>hint ;
 | |
| 
 | |
| GENERIC: find ( selector -- mdb-cursor/f seq )
 | |
| 
 | |
| M: mdb-query-msg find
 | |
|     fix-query-collection send-query ;
 | |
| 
 | |
| M: mdb-cursor find
 | |
|     get-more ;
 | |
| 
 | |
| : each-chunk ( selector quot: ( seq -- ) -- )
 | |
|     swap find
 | |
|     [ pick call( seq -- ) ] when*
 | |
|     [ swap each-chunk ] [ drop ] if* ;
 | |
| 
 | |
| : find-all ( selector -- seq )
 | |
|     [ V{ } clone ] dip
 | |
|     over '[ _ push-all ] each-chunk >array ;
 | |
| 
 | |
| : explain. ( mdb-query-msg -- )
 | |
|     t >>explain find nip . ;
 | |
| 
 | |
| : find-one ( mdb-query-msg -- result/f )
 | |
|     fix-query-collection 1 >>return#
 | |
|     send-query-plain objects>> ?first ;
 | |
| 
 | |
| : count ( mdb-query-msg -- result )
 | |
|     [ count-cmd make-cmd ] dip
 | |
|     [ collection>> "count" set-cmd-opt ]
 | |
|     [ query>> "query" set-cmd-opt ] bi send-cmd 
 | |
|     [ check-ok nip ] keep '[ "n" _ at >fixnum ] [ f ] if ;
 | |
| 
 | |
| : lasterror ( -- error )
 | |
|     getlasterror-cmd make-cmd send-cmd
 | |
|     [ "err" ] dip at ;
 | |
| 
 | |
| GENERIC: validate. ( collection -- )
 | |
| 
 | |
| M: string validate.
 | |
|     [ validate-cmd make-cmd ] dip
 | |
|     "validate" set-cmd-opt send-cmd
 | |
|     [ check-ok nip ] keep
 | |
|     '[ "result" _ at print ] [  ] if ;
 | |
| 
 | |
| M: mdb-collection validate.
 | |
|     name>> validate. ;
 | |
| 
 | |
| <PRIVATE
 | |
| 
 | |
| : send-message-check-error ( message -- )
 | |
|     send-message lasterror [ mdb-error ] when* ;
 | |
| 
 | |
| PRIVATE>
 | |
| 
 | |
| : save ( collection assoc -- )
 | |
|     [ check-collection ] dip
 | |
|     <mdb-insert-msg> send-message-check-error ;
 | |
| 
 | |
| : save-unsafe ( collection assoc -- )
 | |
|     [ check-collection ] dip
 | |
|     <mdb-insert-msg> send-message ;
 | |
| 
 | |
| : ensure-index ( index-spec -- )
 | |
|     <linked-hash> [ [ <oid> "_id" ] dip set-at ] keep
 | |
|     [ { [ [ name>> "name" ] dip set-at ]
 | |
|         [ [ ns>> index-ns "ns" ] dip set-at ]
 | |
|         [ [ key>> "key" ] dip set-at ]
 | |
|         [ swap unique?>>
 | |
|           [ swap [ "unique" ] dip set-at ] [ drop ] if* ] } 2cleave
 | |
|     ] keep
 | |
|     [ index-collection ] dip save ;
 | |
| 
 | |
| : drop-index ( collection name -- )
 | |
|     [ delete-index-cmd make-cmd ] 2dip
 | |
|     [ "deleteIndexes" set-cmd-opt ]
 | |
|     [ "index" set-cmd-opt ] bi* send-cmd drop ;
 | |
| 
 | |
| : <update> ( collection selector object -- mdb-update-msg )
 | |
|     [ check-collection ] 2dip <mdb-update-msg> ;
 | |
| 
 | |
| : >upsert ( mdb-update-msg -- mdb-update-msg )
 | |
|     [ 0 set-bit ] change-update-flags ;
 | |
| 
 | |
| : >multi ( mdb-update-msg -- mdb-update-msg )
 | |
|     [ 1 set-bit ] change-update-flags ;
 | |
| 
 | |
| : update ( mdb-update-msg -- )
 | |
|     send-message-check-error ;
 | |
| 
 | |
| : update-unsafe ( mdb-update-msg -- )
 | |
|     send-message ;
 | |
| 
 | |
| : find-and-modify ( collection selector modifier -- mongodb-cmd )
 | |
|     [ findandmodify-cmd make-cmd ] 3dip
 | |
|     [ "findandmodify" set-cmd-opt ]
 | |
|     [ "query" set-cmd-opt ]
 | |
|     [ "update" set-cmd-opt ] tri* ; inline
 | |
| 
 | |
| : run-cmd ( cmd -- result )
 | |
|     send-cmd ; inline
 | |
| 
 | |
| : <delete> ( collection selector -- mdb-delete-msg )
 | |
|     [ check-collection ] dip <mdb-delete-msg> ;
 | |
| 
 | |
| : >single-remove ( mdb-delete-msg -- mdb-delete-msg )
 | |
|     [ 0 set-bit ] change-delete-flags ;
 | |
| 
 | |
| : delete ( mdb-delete-msg -- )
 | |
|     send-message-check-error ;
 | |
| 
 | |
| : delete-unsafe ( mdb-delete-msg -- )
 | |
|     send-message ;
 | |
| 
 | |
| : kill-cursor ( mdb-cursor -- )
 | |
|     id>> <mdb-killcursors-msg> send-message ;
 | |
| 
 | |
| : load-index-list ( -- index-list )
 | |
|     index-collection
 | |
|     H{ } clone <mdb-query-msg> find nip ;
 | |
| 
 | |
| : ensure-collection ( name -- )
 | |
|     check-collection drop ;
 | |
| 
 | |
| : drop-collection ( name -- )
 | |
|     [ drop-cmd make-cmd ] dip
 | |
|     "drop" set-cmd-opt send-cmd drop ;
 | |
| 
 | |
| 
 |