feat: add --temp-data CLI option

hathor/builder/cli_builder.py

@@ -131,23 +131,28 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
         self.event_ws_factory: Optional[EventWebsocketFactory] = None
 
         if self._args.memory_storage:
+            self.log.warn('--memory-storage is deprecated, use --temp-data instead')
             self.check_or_raise(not self._args.data, '--data should not be used with --memory-storage')
+            self.check_or_raise(not self._args.temp_data, '--temp-data should not be used with --memory-storage')
             # if using MemoryStorage, no need to have cache
             indexes = MemoryIndexesManager()
             tx_storage = TransactionMemoryStorage(indexes, settings=settings)
             event_storage = EventMemoryStorage()
             self.check_or_raise(not self._args.x_rocksdb_indexes, 'RocksDB indexes require RocksDB data')
             self.log.info('with storage', storage_class=type(tx_storage).__name__)
         else:
-            self.check_or_raise(bool(self._args.data), '--data is expected')
-            assert self._args.data is not None
+            self.check_or_raise(
+                bool(self._args.data) or self._args.temp_data,
+                'either --data or --temp-data is expected'
+            )
             if self._args.rocksdb_storage:
                 self.log.warn('--rocksdb-storage is now implied, no need to specify it')
             cache_capacity = self._args.rocksdb_cache
             self.rocksdb_storage = RocksDBStorage(path=self._args.data, cache_capacity=cache_capacity)
 
             # Initialize indexes manager.
             if self._args.memory_indexes:
+                self.log.warn('--memory-indexes is deprecated')
                 indexes = MemoryIndexesManager()
             else:
                 indexes = RocksDBIndexesManager(self.rocksdb_storage)

hathor/cli/run_node.py

@@ -95,12 +95,14 @@ def create_parser(cls) -> ArgumentParser:
         parser.add_argument('--x-status-ipv6-interface', help='IPv6 interface to bind the status server')
         parser.add_argument('--stratum', type=int, help='Port to run stratum server')
         parser.add_argument('--x-stratum-ipv6-interface', help='IPv6 interface to bind the stratum server')
-        parser.add_argument('--data', help='Data directory')
+        data_group = parser.add_mutually_exclusive_group()
+        data_group.add_argument('--data', help='Data directory')
+        data_group.add_argument('--temp-data', action='store_true',
+                                help='Automatically create storage in a temporary directory')
         storage = parser.add_mutually_exclusive_group()
-        storage.add_argument('--rocksdb-storage', action='store_true', help='Use RocksDB storage backend (default)')
-        storage.add_argument('--memory-storage', action='store_true', help='Do not use a persistent storage')
-        parser.add_argument('--memory-indexes', action='store_true',
-                            help='Use memory indexes when using RocksDB storage (startup is significantly slower)')
+        storage.add_argument('--rocksdb-storage', action='store_true', help=SUPPRESS)  # deprecated
+        storage.add_argument('--memory-storage', action='store_true', help=SUPPRESS)  # deprecated
+        parser.add_argument('--memory-indexes', action='store_true', help=SUPPRESS)  # deprecated
         parser.add_argument('--rocksdb-cache', type=int, help='RocksDB block-table cache size (bytes)', default=None)
         parser.add_argument('--wallet', help='Set wallet type. Options are hd (Hierarchical Deterministic) or keypair',
                             default=None)

hathor/cli/run_node_args.py

@@ -43,6 +43,7 @@ class RunNodeArgs(BaseModel, extra=Extra.allow):
     rocksdb_storage: bool
     memory_storage: bool
     memory_indexes: bool
+    temp_data: bool
     rocksdb_cache: Optional[int]
     wallet: Optional[str]
     wallet_enable_api: bool

hathor/storage/rocksdb_storage.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-from typing import TYPE_CHECKING, Optional
+from __future__ import annotations
 
-if TYPE_CHECKING:
-    import rocksdb
+import os
+import tempfile
 
+import rocksdb
 from structlog import get_logger
+from typing_extensions import assert_never
 
 logger = get_logger()
 _DB_NAME = 'data_v2.db'
@@ -28,12 +29,16 @@ class RocksDBStorage:
     """ Creates a RocksDB database
         Give clients the option to create column families
     """
-    def __init__(self, path: str = './', cache_capacity: Optional[int] = None):
-        import rocksdb
+    def __init__(
+        self,
+        path: str | tempfile.TemporaryDirectory | None = None,
+        cache_capacity: int | None = None,
+    ) -> None:
         self.log = logger.new()
-        self._path = path
+        # We have to keep a reference to the TemporaryDirectory because it is cleaned up when garbage collected.
+        self.path, self.temp_dir = self._get_path_and_temp_dir(path)
 
-        db_path = os.path.join(path, _DB_NAME)
+        db_path = os.path.join(self.path, _DB_NAME)
         lru_cache = cache_capacity and rocksdb.LRUCache(cache_capacity)
         table_factory = rocksdb.BlockBasedTableFactory(block_cache=lru_cache)
         options = rocksdb.Options(
@@ -58,13 +63,28 @@ def __init__(self, path: str = './', cache_capacity: Optional[int] = None):
 
         # finally, open the database
         self._db = rocksdb.DB(db_path, options, column_families=column_families)
+        self.log.info('starting rocksdb', path=self.path)
         self.log.debug('open db', cf_list=[cf.name.decode('ascii') for cf in self._db.column_families])
 
-    def get_db(self) -> 'rocksdb.DB':
+    @staticmethod
+    def _get_path_and_temp_dir(
+        path: str | tempfile.TemporaryDirectory | None,
+    ) -> tuple[str, tempfile.TemporaryDirectory | None]:
+        match path:
+            case str():
+                return path, None
+            case tempfile.TemporaryDirectory():
+                return path.name, path
+            case None:
+                temp_dir = tempfile.TemporaryDirectory()
+                return temp_dir.name, temp_dir
+            case _:
+                assert_never(path)
+
+    def get_db(self) -> rocksdb.DB:
         return self._db
 
     def get_or_create_column_family(self, cf_name: bytes) -> 'rocksdb.ColumnFamilyHandle':
-        import rocksdb
         cf = self._db.get_column_family(cf_name)
         if cf is None:
             cf = self._db.create_column_family(cf_name, rocksdb.ColumnFamilyOptions())