o ;hF@sddlmZddlZddlZddlmZmZddlmZm Z ddl m Z m Z m Z mZddlmZdd lmZmZmZGd d d eeeeZed d e dddGdddeZdS)) annotationsN)ABCabstractmethod) versionaddedversionchanged) MongoClientMongoCollection MongoDatabasecast)get_dependency)MovingWindowSupportSlidingWindowCounterSupportStoragecseZdZdZdgZ    dIdJfdd ZedKddZedLddZedMddZ edMddZ e dNd d!Z edOd#d$Z dPd%d&ZdQd(d)ZdRd+d,ZdSd.d/ZdTd1d2ZdUdVd6d7ZdWd8d9ZdXdd?ZdZdAdBZ 3dUdYdCdDZd[dEdFZdPdGdHZZS)\MongoDBStorageBasezV Rate limit storage with MongoDB as backend. Depends on :pypi:`pymongo`. pymongolimitscounterswindowsFuristr database_namecounter_collection_namewindow_collection_namewrap_exceptionsbooloptionsint | str | boolreturnNonec s\tj|fd|i|||_||d|_|jdj|_td\|_}||_ ||_ d|_ dS)ai :param uri: uri of the form ``mongodb://[user:password]@host:port?...``, This uri is passed directly to :class:`~pymongo.mongo_client.MongoClient` :param database_name: The database to use for storing the rate limit collections. :param counter_collection_name: The collection name to use for individual counters used in fixed window strategies :param window_collection_name: The collection name to use for sliding & moving window storage :param wrap_exceptions: Whether to wrap storage exceptions in :exc:`limits.errors.StorageError` before raising it. :param options: all remaining keyword arguments are passed to the constructor of :class:`~pymongo.mongo_client.MongoClient` :raise ConfigurationError: when the :pypi:`pymongo` library is not available r)rrrzpymongo.errorsN) super__init___database_name_collection_mapping dependenciesmodulelibr lib_errors _storage_uri_storage_options_storage)selfrrrrrr_ __class__I/var/www/html/venv/lib/python3.10/site-packages/limits/storage/mongodb.pyr"s zMongoDBStorageBase.__init__rcCs0|jdur|j|jfi|j|_||jSN)r+_init_mongo_clientr)r*(_MongoDBStorageBase__initialize_databaser,r0r0r1storageDs zMongoDBStorageBase.storager cCs |j|jSr2)r6r#r5r0r0r1 _databaseMs zMongoDBStorageBase._databasercC|j|jdS)Nrr7r$r5r0r0r1rQzMongoDBStorageBase.counterscCr8)Nrr9r5r0r0r1rUr:zMongoDBStorageBase.windows str | NonecKstr2)NotImplementedErrorr,rrr0r0r1r3Ysz%MongoDBStorageBase._init_mongo_client-type[Exception] | tuple[type[Exception], ...]cCs|jjSr2)r(Z PyMongoErrorr5r0r0r1base_exceptions_sz"MongoDBStorageBase.base_exceptionscCs$|jjddd|jjddddS)NexpireAtr)ZexpireAfterSeconds)rZ create_indexrr5r0r0r1Z__initialize_databaseesz(MongoDBStorageBase.__initialize_database int | NonecCs4|ji|ji}|j|jt|S)z^ Delete all rate limit keys in the rate limit collections (counters, windows) )rZcount_documentsrZdropint)r,Znum_keysr0r0r1resetis  zMongoDBStorageBase.resetkeycCs$|jd|i|jd|idS)z> :param key: the key to clear rate limits for _idN)rZfind_one_and_deleter)r,rDr0r0r1clearsszMongoDBStorageBase.clearfloatcCs6|jd|i}|r|dntjjtjjdS)z; :param key: the key to get the expiry for rEr@tzinfo)rfind_onedatetimenowreplacetimezoneutc timestampr,rDZcounterr0r0r1 get_expiryzs zMongoDBStorageBase.get_expiryrBcCs8|jj|dtjtjjiddgd}|r|dpdS)zB :param key: the key to get the counter value for $gte)rEr@count) projectionr)rrJrKrLrNrOrQr0r0r1getszMongoDBStorageBase.getr expiryamountc Cstjtjjtj|d}t|jjd|idddddgi|dd |gid iddddgi|dd id igd d g|jj j dd S)z increments the counter for a given rate limit key :param key: the key to increment :param expiry: amount in seconds for the key to expire in :param amount: the number to increment by secondsrE$set$condz$lt $expireAt$$NOW$addz$countifZthenelse)rTr@TrT)upsertrUreturn_document) rKrLrNrO timedeltarBrfind_one_and_updater'ReturnDocumentAFTER)r,rDrWrXZ expirationr0r0r1incrs:   zMongoDBStorageBase.incrcCsz|jWdSYdS)zm Check if storage is healthy by calling :meth:`pymongo.mongo_client.MongoClient.server_info` TF)r6Z server_infor5r0r0r1checks  zMongoDBStorageBase.checklimittuple[float, int]c Cs|t}t|jdd|iiddddddd ||gid iiidd d id d idig}r:|dd|ddfS|dfS)z returns the starting point and the number of entries in the moving window :param key: rate limit key :param expiry: expiry of entry :return: (start of window, number of acquired entries) z$matchrEz$projectZfilteredEntriesz$filterz$entriesentryrSz$$entry)inputasZcondz$minz$filteredEntriesz$size)minrTrrprT)timelistrZ aggregate)r,rDrkrWrPresultr0r0r1get_moving_windows0   z$MongoDBStorageBase.get_moving_windowcCs||krdSt}z7d|g|d|didtjtjjtj|did}|jjd|d ||d d ||iii|d d Wd S|jj j yNYdSw)z :param key: rate limit key to acquire an entry in :param limit: amount of entries allowed :param expiry: expiry of the entry :param amount: the number of entries to acquire Fentriesr)z$eachz $positionz$slicer@rY)z$pushr[rEzentries.z$notrST)rc) rqrKrLrNrOrerZ update_oner'errorsZDuplicateKeyError)r,rDrkrWrXrPZupdatesr0r0r1 acquire_entrys8  z MongoDBStorageBase.acquire_entrytuple[int, float, int, float]c Cs|d}|jjd|iddddddgi|gid d d gid d d gid idddddgi|gid d d d gid idddddgi|gidd|gidd idig|jjjgdd}r|dri|djtjj d nt }t d |t }t d |dr||nd }|d||d|fSdS)NrEr[r\$lte $subtractr]r^$ifNull $currentCountr$previousCountr`r_) previousCount currentCountr@)rrr@)rdrUr@rHrr)rrr) rrfr'rgrhrVrMrKrNrOrPrqmax)r,rDrW expiry_msrsZ expires_atZ current_ttlZprev_ttlr0r0r1get_sliding_windowsj        , 3z%MongoDBStorageBase.get_sliding_windowcCsl|d}|jjd|iddddddd gi|gid d d gid d d gidiiidddddd gi|gid d d d gididdddd gi|gidddd gidd|gidd d|gididdididddddd ddd dddd |gigigi|gigid giiiidddddd|gi|gidd |gid diiiddddd|gi|giiiddgig|jjjdd}tt|dS)NryrEr[rr\rzr{r]r^r|r}rr~r`z$gtr_r )rr@ZcurWeightedCountz$floorz $multiplyz$dividez$maxrz$curWeightedCountZ _acquiredz$unsetT)rdrc)rrfr'rgrhr r)r,rDrkrWrXrrsr0r0r1acquire_sliding_window_entry]s         #&  qvz/MongoDBStorageBase.acquire_sliding_window_entrycCs ||Sr2)rF)r,rDrWr0r0r1clear_sliding_windows z'MongoDBStorageBase.clear_sliding_windowcCs|jr |jdSdSr2)r6closer5r0r0r1__del__szMongoDBStorageBase.__del__)rrrF)rrrrrrrrrrrrrr )rr)rr )rrrr;rrrr)rr>)rr )rrA)rDrrr )rDrrrG)rDrrrB)r )rDrrWrBrXrBrrB)rr)rDrrkrBrWrBrrl) rDrrkrBrWrBrXrBrr)rDrrWrBrrx)rDrrWrBrr )__name__ __module__ __qualname____doc__Z DEPENDENCIESr"propertyr6r7rrrr3r?r4rCrFrRrVrirjrtrwrrrr __classcell__r0r0r.r1rsD%           ) % ,F |rz2.1)versionz3.14.0zEAdded option to select custom collection names for windows & counters)rreasonc@seZdZddgZd d d Zd S) MongoDBStorageZmongodbz mongodb+srvrr;rrrrcKstt|jj|fi|Sr2)r rr'r=r0r0r1r3sz!MongoDBStorage._init_mongo_clientNr)rrrZSTORAGE_SCHEMEr3r0r0r0r1rsr) __future__rrKrqabcrrZdeprecated.sphinxrrZ limits.typingrrr r utilr baserrrrrr0r0r0r1s(   P