Merge pull request pandas-dev#263 from manahl/mongo_quota_improvement

Mongo quota improvements

Author: jamesblackburn

arctic/arctic.py

@@ -454,20 +454,25 @@ def to_gigabytes(bytes_):
         size = stats['totals']['size']
         count = stats['totals']['count']
         if size >= self.quota:
-            raise QuotaExceededException("Quota Exceeded: %.3f / %.0f GB used" %
-                                         (to_gigabytes(size),
-                                          to_gigabytes(self.quota)))
+            raise QuotaExceededException("Mongo Quota Exceeded: %s %.3f / %.0f GB used" % (
+                                         '.'.join([self.database_name, self.library]),
+                                         to_gigabytes(size),
+                                         to_gigabytes(self.quota)))
 
         # Quota not exceeded, print an informational message and return
         avg_size = size // count if count > 1 else 100 * 1024
         remaining = self.quota - size
         remaining_count = remaining / avg_size
-        if remaining_count < 100:
-            logger.warning("Mongo Quota: %.3f / %.0f GB used" % (to_gigabytes(size),
-                                                              to_gigabytes(self.quota)))
+        if remaining_count < 100 or float(remaining) / self.quota < 0.1:
+            logger.warning("Mongo Quota: %s %.3f / %.0f GB used" % (
+                            '.'.join([self.database_name, self.library]),
+                            to_gigabytes(size),
+                            to_gigabytes(self.quota)))
         else:
-            logger.info("Mongo Quota: %.3f / %.0f GB used" % (to_gigabytes(size),
-                                                              to_gigabytes(self.quota)))
+            logger.info("Mongo Quota: %s %.3f / %.0f GB used" % (
+                            '.'.join([self.database_name, self.library]),
+                            to_gigabytes(size),
+                            to_gigabytes(self.quota)))
 
         # Set-up a timer to prevent us for checking for a few writes.
         # This will check every average half-life

tests/unit/test_arctic.py

@@ -196,7 +196,8 @@ def test_check_quota_countdown():
 
 
 def test_check_quota():
-    self = create_autospec(ArcticLibraryBinding)
+    self = create_autospec(ArcticLibraryBinding, database_name='arctic_db',
+                           library='lib')
     self.arctic = create_autospec(Arctic)
     self.get_library_metadata.return_value = 1024 * 1024 * 1024
     self.quota_countdown = 0
@@ -208,12 +209,30 @@ def test_check_quota():
     with patch('arctic.arctic.logger.warning') as warn:
         ArcticLibraryBinding.check_quota(self)
     self.arctic.__getitem__.assert_called_once_with(self.get_name.return_value)
-    warn.assert_called_once_with('Mongo Quota: 0.879 / 1 GB used')
+    warn.assert_called_once_with('Mongo Quota: arctic_db.lib 0.879 / 1 GB used')
     assert self.quota_countdown == 6
 
 
+def test_check_quota_90_percent():
+    self = create_autospec(ArcticLibraryBinding, database_name='arctic_db',
+                           library='lib')
+    self.arctic = create_autospec(Arctic)
+    self.get_library_metadata.return_value = 1024 * 1024 * 1024
+    self.quota_countdown = 0
+    self.arctic.__getitem__.return_value = Mock(stats=Mock(return_value={'totals':
+                                                                             {'size': 0.91 * 1024 * 1024 * 1024,
+                                                                              'count': 1000000,
+                                                                              }
+                                                                             }))
+    with patch('arctic.arctic.logger.warning') as warn:
+        ArcticLibraryBinding.check_quota(self)
+    self.arctic.__getitem__.assert_called_once_with(self.get_name.return_value)
+    warn.assert_called_once_with('Mongo Quota: arctic_db.lib 0.910 / 1 GB used')
+
+
 def test_check_quota_info():
-    self = create_autospec(ArcticLibraryBinding)
+    self = create_autospec(ArcticLibraryBinding, database_name='arctic_db',
+                           library='lib')
     self.arctic = create_autospec(Arctic)
     self.get_library_metadata.return_value = 1024 * 1024 * 1024
     self.quota_countdown = 0
@@ -225,12 +244,13 @@ def test_check_quota_info():
     with patch('arctic.arctic.logger.info') as info:
         ArcticLibraryBinding.check_quota(self)
     self.arctic.__getitem__.assert_called_once_with(self.get_name.return_value)
-    info.assert_called_once_with('Mongo Quota: 0.001 / 1 GB used')
+    info.assert_called_once_with('Mongo Quota: arctic_db.lib 0.001 / 1 GB used')
     assert self.quota_countdown == 51153
 
 
 def test_check_quota_exceeded():
-    self = create_autospec(ArcticLibraryBinding)
+    self = create_autospec(ArcticLibraryBinding, database_name='arctic_db',
+                           library='lib')
     self.arctic = create_autospec(Arctic)
     self.get_library_metadata.return_value = 1024 * 1024 * 1024
     self.quota_countdown = 0
@@ -241,7 +261,7 @@ def test_check_quota_exceeded():
                                                                              }))
     with pytest.raises(QuotaExceededException) as e:
         ArcticLibraryBinding.check_quota(self)
-    assert "Quota Exceeded: 1.000 / 1 GB used" in str(e)
+    assert "Quota Exceeded: arctic_db.lib 1.000 / 1 GB used" in str(e)
 
 
 def test_initialize_library():