last changes

mathema/core/db.py

@@ -1,4 +1,6 @@
-from neo4j import AsyncGraphDatabase
+from typing import LiteralString, cast
+
+from neo4j import AsyncGraphDatabase, Query
 
 NEO4J_CONSTRAINTS = [
     "CREATE CONSTRAINT cortex_id    IF NOT EXISTS FOR (n:cortex)    REQUIRE n.id IS UNIQUE",
@@ -30,37 +32,129 @@ class Neo4jDB:
             return await s.run(cypher, **params)
     """
 
-    async def run_read(self, cypher: str, **params):
+    async def execute_write(self, work):
+        """
+
+        Method: execute_write
+
+        Description:
+        This method is used to execute a write operation on the database using the provided work.
+
+        Parameters:
+        self: The current instance of the class.
+        work: The work to be executed as a write operation on the database.
+
+        """
+        async with self._driver.session(database=self._database) as session:
+            return await session.execute_write(work)
+
+    async def run_read(self, cypher: LiteralString | Query, **params):
+        """
+
+        Method: run_read
+
+        Description:
+        This method allows running a read operation with the provided cypher query and parameters using the underlying
+        driver session. It returns the result of the read operation.
+
+        Parameters:
+        - cypher: str or Query object representing the cypher query to be executed.
+        - **params: Additional keyword arguments that represent parameters to be passed to the cypher query.
+
+        Returns:
+        Result of the read operation based on the provided cypher query and parameters.
+
+        """
         async with self._driver.session(database=self._database) as s:
             return await s.run(cypher, **params)
 
-    async def read_single(self, cypher: str, **params):
+    async def read_single(self, cypher: LiteralString | Query, **params):
+        """
+        Reads a single record from the database using the provided Cypher query and parameters.
+
+        :param cypher: The Cypher query to execute.
+        :param params: Additional parameters to be passed to the query.
+
+        :return: A single record retrieved from the database based on the given Cypher query and parameters.
+        """
         async with self._driver.session(database=self._database) as s:
             res = await s.run(cypher, **params)
             return await res.single()
 
-    async def read_all(self, cypher: str, **params):
+    async def read_all(self, cypher: LiteralString | Query, **params):
+        """
+        Reads all records from the database based on the given cypher query and parameters.
+
+        Parameters:
+            cypher (str): The Cypher query to execute.
+            **params: Additional parameters to pass to the Cypher query.
+
+        Return Type:
+            list: A list of retrieved records from the database.
+        """
         async with self._driver.session(database=self._database) as s:
             res = await s.run(cypher, **params)
             return [r async for r in res]
 
-    async def run_consume(self, cypher: str, **params):
+    async def run_consume(self, cypher: LiteralString | Query, **params):
+        """
+
+        Run a Cypher query and consume the result.
+
+        Parameters:
+            cypher : Union[LiteralString, Query]
+                The Cypher query to be executed.
+            **params : Any
+                Keyword arguments for query parameters.
+
+        Returns:
+            None
+
+        """
         async with self._driver.session(database=self._database) as s:
             res = await s.run(cypher, **params)
             return await res.consume()
 
     async def create_schema(self):
+        """
+        Creates database schema by running Neo4j constraints queries.
+
+        Parameters:
+        - self: instance of the class
+        - database: name of the database to be used for creating schema
+
+        Returns:
+        - None
+        """
         async with self._driver.session(database=self._database) as s:
             for stmt in NEO4J_CONSTRAINTS:
-                await s.run(stmt)
+                await s.run(cast(LiteralString, stmt))
 
     async def purge_all_nodes(self):
+        """
+        Purge all nodes in the database.
+
+        Method parameters:
+        - None
+
+        Returns:
+        - None
+        """
         async with self._driver.session(database=self._database) as s:
             await s.run("MATCH (n) DETACH DELETE n")
 
     async def drop_schema(self):
+        """
+        Drop the current schema by dropping all constraints in the database.
+
+        Parameters:
+        None
+
+        Return Type:
+        None
+        """
         async with self._driver.session(database=self._database) as s:
             res = await s.run("SHOW CONSTRAINTS")
             async for record in res:
                 name = record["name"]
-                await s.run(f"DROP CONSTRAINT {name} IF EXISTS")
+                await s.run(cast(LiteralString, f"DROP CONSTRAINT {name} IF EXISTS"))