Add storage mode info

gqlalchemy/vendors/memgraph.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from enum import Enum
 import os
 import sqlite3
 from typing import List, Optional, Union
@@ -49,6 +50,15 @@ class MemgraphConstants:
     UNIQUE = "unique"
 
 
+class MemgraphStorageMode(Enum):
+    IN_MEMORY_TRANSACTIONAL = "IN_MEMORY_TRANSACTIONAL"
+    IN_MEMORY_ANALYTICAL = "IN_MEMORY_ANALYTICAL"
+    ON_DISK_TRANSACTIONAL = "ON_DISK_TRANSACTIONAL"
+
+    def __str__(self):
+        return self.value
+
+
 class Memgraph(DatabaseClient):
     def __init__(
         self,
@@ -432,3 +442,13 @@ def with_power_bi(self) -> "Memgraph":
         module_name = "power_bi_stream.py"
 
         return self.add_query_module(file_path=file_path, module_name=module_name)
+
+    def get_storage_mode(self) -> str:
+        """Returns the storage mode of the Memgraph instance."""
+        result = self.execute_and_fetch("SHOW STORAGE INFO;")
+        storage_mode_value = next((item["value"] for item in result if item["storage info"] == "storage_mode"), None)
+        return MemgraphStorageMode(storage_mode_value).value
+
+    def set_storage_mode(self, storage_mode: MemgraphStorageMode) -> None:
+        """Sets the storage mode of the Memgraph instance."""
+        self.execute(f"STORAGE MODE {storage_mode};")

tests/conftest.py

@@ -40,11 +40,13 @@ def memgraph() -> Memgraph:
     memgraph.ensure_indexes([])
     memgraph.ensure_constraints([])
     memgraph.drop_database()
+    memgraph.set_storage_mode("IN_MEMORY_TRANSACTIONAL")
 
     yield memgraph
 
     memgraph.ensure_indexes([])
     memgraph.ensure_constraints([])
+    memgraph.set_storage_mode("IN_MEMORY_TRANSACTIONAL")
 
 
 @pytest.fixture

tests/integration/test_storage_modes.py

@@ -0,0 +1,10 @@
+import pytest
+from gqlalchemy.exceptions import GQLAlchemyDatabaseError
+from gqlalchemy import create
+
+
+def test_switch_to_on_disk(memgraph):
+    create().node(labels="Person", name="Leslie").execute()
+
+    with pytest.raises(GQLAlchemyDatabaseError):
+        memgraph.set_storage_mode("ON_DISK_TRANSACTIONAL")

tests/ogm/test_storage_modes.py

@@ -0,0 +1,7 @@
+def test_get_storage_mode(memgraph):
+    assert memgraph.get_storage_mode() == "IN_MEMORY_TRANSACTIONAL"
+
+
+def test_set_storage_mode(memgraph):
+    memgraph.set_storage_mode("IN_MEMORY_ANALYTICAL")
+    assert memgraph.get_storage_mode() == "IN_MEMORY_ANALYTICAL"