From 0043bd5de3df6bf092d65ef1890d4eb86c36e5e4 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Tue, 18 Sep 2018 03:40:41 +0300
Subject: [PATCH] state: limit number of rows

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/config.py       | 17 ++++++++++++++++
 dvc/project.py      |  2 +-
 dvc/state.py        | 49 ++++++++++++++++++++++++++++++++++++++++++---
 tests/test_state.py |  2 +-
 4 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/dvc/config.py b/dvc/config.py
index 15091ebc9..d8abad665 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -45,6 +45,14 @@ def to_bool(val):
     return val.lower() == 'true'
 
 
+def is_whole(val):
+    return int(val) >= 0
+
+
+def is_percent(val):
+    return int(val) >= 0 and int(val) <= 100
+
+
 class Config(object):
     CONFIG = 'config'
     CONFIG_LOCAL = 'config.local'
@@ -149,10 +157,19 @@ class Config(object):
         Optional(SECTION_REMOTE_TIMEOUT): Use(int),
     }
 
+    SECTION_STATE = 'state'
+    SECTION_STATE_ROW_LIMIT = 'row_limit'
+    SECTION_STATE_ROW_CLEANUP_QUOTA = 'row_cleanup_quota'
+    SECTION_STATE_SCHEMA = {
+        Optional(SECTION_STATE_ROW_LIMIT): And(Use(int), is_whole),
+        Optional(SECTION_STATE_ROW_CLEANUP_QUOTA): And(Use(int), is_percent),
+    }
+
     SCHEMA = {
         Optional(SECTION_CORE, default={}): SECTION_CORE_SCHEMA,
         Optional(Regex(SECTION_REMOTE_REGEX)): SECTION_REMOTE_SCHEMA,
         Optional(SECTION_CACHE, default={}): SECTION_CACHE_SCHEMA,
+        Optional(SECTION_STATE, default={}): SECTION_STATE_SCHEMA,
 
         # backward compatibility
         Optional(SECTION_AWS, default={}): SECTION_AWS_SCHEMA,
diff --git a/dvc/project.py b/dvc/project.py
index 141031993..9c60b4a4e 100644
--- a/dvc/project.py
+++ b/dvc/project.py
@@ -38,7 +38,7 @@ class Project(object):
         self.lock = Lock(self.dvc_dir)
         # NOTE: storing state and link_state in the repository itself to avoid
         # any possible state corruption in 'shared cache dir' scenario.
-        self.state = State(self)
+        self.state = State(self, self.config._config)
 
         core = self.config._config[Config.SECTION_CORE]
         self.logger = Logger(core.get(Config.SECTION_CORE_LOGLEVEL, None))
diff --git a/dvc/state.py b/dvc/state.py
index 77c47a8df..337645c3b 100644
--- a/dvc/state.py
+++ b/dvc/state.py
@@ -3,6 +3,7 @@ import time
 import sqlite3
 import nanotime
 
+from dvc.config import Config
 from dvc.system import System
 from dvc.utils import file_md5, remove
 from dvc.exceptions import DvcException
@@ -30,11 +31,23 @@ class State(object):
                               "inode INTEGER NOT NULL, " \
                               "mtime TEXT NOT NULL"
 
-    def __init__(self, project):
+    STATE_ROW_LIMIT = 10000000
+    STATE_ROW_CLEANUP_QUOTA = 50
+
+    def __init__(self, project, config):
         self.project = project
         self.dvc_dir = project.dvc_dir
         self.root_dir = project.root_dir
 
+        self.row_limit = 100
+        self.row_cleanup_quota = 50
+
+        c = config.get(Config.SECTION_STATE, {})
+        self.row_limit = c.get(Config.SECTION_STATE_ROW_LIMIT,
+                               self.STATE_ROW_LIMIT)
+        self.row_cleanup_quota = c.get(Config.SECTION_STATE_ROW_CLEANUP_QUOTA,
+                                       self.STATE_ROW_CLEANUP_QUOTA)
+
         if not self.dvc_dir:
             self.state_file = None
             return
@@ -110,9 +123,39 @@ class State(object):
     def dump(self):
         assert self.db is not None
 
-        cmd = "UPDATE {} SET count = count + {} WHERE rowid = {}"
+        cmd = "SELECT count from {} WHERE rowid={}"
         self.c.execute(cmd.format(self.STATE_INFO_TABLE,
-                                  self.inserts,
+                                  self.STATE_INFO_ROW))
+        ret = self.c.fetchall()
+        assert len(ret) == 1
+        assert len(ret[0]) == 1
+        count = ret[0][0] + self.inserts
+
+        if count > self.row_limit:
+            msg = "Cleaning up state. This might take a while."
+            self.project.logger.warn(msg)
+
+            delete = (count - self.row_limit)
+            delete += int(self.row_limit * (self.row_cleanup_quota/100.))
+            cmd = "DELETE FROM {} WHERE timestamp IN (" \
+                  "SELECT timestamp FROM {} ORDER BY timestamp ASC LIMIT {});"
+            self.c.execute(cmd.format(self.STATE_TABLE,
+                                      self.STATE_TABLE,
+                                      delete))
+
+            self.c.execute("VACUUM")
+
+            cmd = "SELECT COUNT(*) FROM {}"
+
+            self.c.execute(cmd.format(self.STATE_TABLE))
+            ret = self.c.fetchall()
+            assert len(ret) == 1
+            assert len(ret[0]) == 1
+            count = ret[0][0]
+
+        cmd = "UPDATE {} SET count = {} WHERE rowid = {}"
+        self.c.execute(cmd.format(self.STATE_INFO_TABLE,
+                                  count,
                                   self.STATE_INFO_ROW))
 
         self.db.commit()
diff --git a/tests/test_state.py b/tests/test_state.py
index ce9bcd5be..fa7894db2 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -16,7 +16,7 @@ class TestState(TestDvc):
         mtime = os.path.getmtime(path)
         inode = System.inode(path)
 
-        state = State(self.dvc)
+        state = State(self.dvc, self.dvc.config._config)
 
         with state:
             entry_md5 = state.update(path)