diff --git a/gqlalchemy/vendors/memgraph.py b/gqlalchemy/vendors/memgraph.py index 6e8dc1ad..fb5fb7a9 100644 --- a/gqlalchemy/vendors/memgraph.py +++ b/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 @@ -69,6 +70,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 + + def create_transaction(transaction_data) -> MemgraphTransaction: """Create a MemgraphTransaction object from transaction data. Args: @@ -481,6 +491,16 @@ def with_power_bi(self) -> "Memgraph": 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};") + def get_transactions(self) -> List[MemgraphTransaction]: """Get all transactions in the database. Returns: diff --git a/tests/conftest.py b/tests/conftest.py index 0eaa0465..d70915d1 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/integration/test_storage_modes.py b/tests/integration/test_storage_modes.py new file mode 100644 index 00000000..3276c6b8 --- /dev/null +++ b/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") diff --git a/tests/ogm/test_storage_modes.py b/tests/ogm/test_storage_modes.py new file mode 100644 index 00000000..618bdc4c --- /dev/null +++ b/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"