Repository: AgnostiqHQ/covalent (language: python)

Issue:
Implement Result.persist method
This issue is dedicated to implementing the `Result.persist` method. There are several auxiliary methods that need to be implemented in order to accomplish this. Planning documentation can be found [here](https://docs.google.com/document/d/1uaqhBQ6PffsxWHeUjFx2JWd6xGG2-j4lLJbMeJmJG3A/edit#).


**Acceptance criteria**

Implement the following methods / functions:
- [ ] `Result.persist` which is analogous to `Result.save`.
    - [x] `insert_lattices_data`
    - [x] `insert_electrons_data`
    - [x] `insert_electron_dependency_data`
    - [x] `update_lattices_data`
    - [x] `update_electrons_data`
    - [ ] `are_electron_dependencies_added`
    - [x] `get_electron_type`
    - [ ] connecting lattice and sublattices via electron_id. 
- [ ] Keep design doc updated if changes are required.
- [ ] Add unit tests with coverage >= 90%.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 261d6e0b2..34b24d5a6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Added
+
+- `covalent/_results_manager/write_result_to_db.py` module and methods to insert / update data in the DB.
+- `tests/covalent_tests/results_manager_tests/write_result_to_db_test.py` containing the unit tests for corresponding functions.
+
+### Changed
+
+- Electron `type` column to a string type rather than an `ElectronType` in DB models.
+- Primary keys from `BigInteger` to `Integer` in DB models.
+
 ## [0.121.0] - 2022-07-04
 
 ### Authors
@@ -23,6 +33,7 @@ Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.gi
 ### Docs
 
 - Updated the compatibility matrix in the docs.
+
 ## [0.120.0] - 2022-07-04
 
 ### Authors
diff --git a/covalent/_data_store/models.py b/covalent/_data_store/models.py
index 22eb721a8..e9a4d279e 100644
--- a/covalent/_data_store/models.py
+++ b/covalent/_data_store/models.py
@@ -22,35 +22,15 @@
 Models for the workflows db. Based on schema v9
 """
 
-import enum
-
-from sqlalchemy import BigInteger, Column, DateTime, Enum, ForeignKey, Integer, String, Text
+from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text
 from sqlalchemy.orm import declarative_base
 
 Base = declarative_base()
 
 
-# See electron.py's import list and
-# execution.py's _run_planned_workflow
-class ElectronTypeEnum(enum.Enum):
-    electron = 1
-    parameter = 2
-    sublattice = 3
-    collection = 4
-    attribute = 5
-    generated = 6
-    subscript = 7
-
-
-class ParameterTypeEnum(enum.Enum):
-    arg = 1
-    kwarg = 2
-    NULL = 3
-
-
 class Lattice(Base):
     __tablename__ = "lattices"
-    id = Column(BigInteger, primary_key=True)
+    id = Column(Integer, primary_key=True)
     dispatch_id = Column(String(64), nullable=False)
 
     # id of node if the lattice is actually a sublattice
@@ -95,16 +75,16 @@ class Lattice(Base):
 
 class Electron(Base):
     __tablename__ = "electrons"
-    id = Column(BigInteger, primary_key=True)
+    id = Column(Integer, primary_key=True)
 
     # id of the lattice containing this electron
-    parent_lattice_id = Column(BigInteger, ForeignKey("lattices.id"), nullable=False)
+    parent_lattice_id = Column(Integer, ForeignKey("lattices.id"), nullable=False)
 
     # id of the node in the context of a transport graph
     transport_graph_node_id = Column(Integer, nullable=False)
 
     # Node type
-    type = Column(Enum(ElectronTypeEnum), nullable=False)
+    type = Column(String(24), nullable=False)
 
     # Node name
     name = Column(Text, nullable=False)
@@ -127,9 +107,6 @@ class Electron(Base):
     # Name of the file containing the serialized executor
     executor_filename = Column(Text)
 
-    # Name of the file containing an error message for the electron
-    error_filename = Column(Text)
-
     # name of the file containing the serialized output
     results_filename = Column(Text)
 
@@ -161,7 +138,7 @@ class Electron(Base):
 
 class ElectronDependency(Base):
     __tablename__ = "electron_dependency"
-    id = Column(BigInteger, primary_key=True)
+    id = Column(Integer, primary_key=True)
 
     # Unique ID of electron
     electron_id = Column(Integer, nullable=False)
@@ -172,7 +149,7 @@ class ElectronDependency(Base):
     edge_name = Column(Text, nullable=False)
 
     # "args, kwarg, null"
-    parameter_type = Column(Enum(ParameterTypeEnum), nullable=False)
+    parameter_type = Column(String(24), nullable=False)
 
     # Argument position
     arg_index = Column(Integer, nullable=False)
diff --git a/covalent/_results_manager/write_result_to_db.py b/covalent/_results_manager/write_result_to_db.py
new file mode 100644
index 000000000..e7708dfd1
--- /dev/null
+++ b/covalent/_results_manager/write_result_to_db.py
@@ -0,0 +1,305 @@
+# Copyright 2021 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
+"""This module contains all the functions required to save the decomposed result object in the database."""
+
+from datetime import datetime as dt
+
+import networkx as nx
+from sqlalchemy import update
+from sqlalchemy.orm import Session
+
+from covalent._data_store.datastore import DataStore
+from covalent._data_store.models import Electron, ElectronDependency, Lattice
+from covalent._workflow.lattice import Lattice as Workflow_lattice
+
+from .._shared_files.defaults import (
+    arg_prefix,
+    attr_prefix,
+    electron_dict_prefix,
+    electron_list_prefix,
+    generator_prefix,
+    parameter_prefix,
+    prefix_separator,
+    sublattice_prefix,
+    subscript_prefix,
+)
+
+
+class MissingLatticeRecordError(Exception):
+    pass
+
+
+class MissingElectronRecordError(Exception):
+    pass
+
+
+def insert_lattices_data(
+    db: DataStore,
+    dispatch_id: str,
+    name: str,
+    status: str,
+    storage_type: str,
+    storage_path: str,
+    function_filename: str,
+    function_string_filename: str,
+    executor_filename: str,
+    error_filename: str,
+    inputs_filename: str,
+    results_filename: str,
+    created_at: dt,
+    updated_at: dt,
+    started_at: dt,
+    completed_at: dt,
+) -> int:
+    """This funtion writes the lattice data / metadata to the Lattices table in the DB."""
+
+    lattice_row = Lattice(
+        dispatch_id=dispatch_id,
+        name=name,
+        status=status,
+        storage_type=storage_type,
+        storage_path=storage_path,
+        function_filename=function_filename,
+        function_string_filename=function_string_filename,
+        executor_filename=executor_filename,
+        error_filename=error_filename,
+        inputs_filename=inputs_filename,
+        results_filename=results_filename,
+        created_at=created_at,
+        updated_at=updated_at,
+        started_at=started_at,
+        completed_at=completed_at,
+    )
+
+    with Session(db.engine) as session:
+        session.add(lattice_row)
+        session.commit()
+        lattice_id = lattice_row.id
+
+    return lattice_id
+
+
+def insert_electrons_data(
+    db: DataStore,
+    parent_dispatch_id: str,
+    transport_graph_node_id: int,
+    type: str,
+    name: str,
+    status: str,
+    storage_type: str,
+    storage_path: str,
+    function_filename: str,
+    function_string_filename: str,
+    executor_filename: str,
+    results_filename: str,
+    value_filename: str,
+    attribute_name: str,
+    key: str,
+    stdout_filename: str,
+    stderr_filename: str,
+    info_filename: str,
+    created_at: dt,
+    updated_at: dt,
+    started_at: dt,
+    completed_at: dt,
+) -> id:
+    """This function writes the transport graph node data to the Electrons table in the DB."""
+
+    # Check that the foreign key corresponding to this table exists
+    with Session(db.engine) as session:
+        row = session.query(Lattice).where(Lattice.dispatch_id == parent_dispatch_id).all()
+    if len(row) == 0:
+        raise MissingLatticeRecordError
+
+    parent_lattice_id = row[0].id
+
+    electron_row = Electron(
+        parent_lattice_id=parent_lattice_id,
+        transport_graph_node_id=transport_graph_node_id,
+        type=type,
+        name=name,
+        status=status,
+        storage_type=storage_type,
+        storage_path=storage_path,
+        function_filename=function_filename,
+        function_string_filename=function_string_filename,
+        executor_filename=executor_filename,
+        results_filename=results_filename,
+        value_filename=value_filename,
+        attribute_name=attribute_name,
+        key=key,
+        stdout_filename=stdout_filename,
+        stderr_filename=stderr_filename,
+        info_filename=info_filename,
+        created_at=created_at,
+        updated_at=updated_at,
+        started_at=started_at,
+        completed_at=completed_at,
+    )
+
+    with Session(db.engine) as session:
+        session.add(electron_row)
+        session.commit()
+        electron_id = electron_row.id
+
+    return electron_id
+
+
+def insert_electron_dependency_data(db: DataStore, dispatch_id: str, lattice: Workflow_lattice):
+    """Extract electron dependencies from the lattice transport graph and add them to the DB."""
+
+    node_links = nx.readwrite.node_link_data(lattice.transport_graph._graph)["links"]
+
+    print(node_links)
+
+    electron_dependency_ids = []
+    with Session(db.engine) as session:
+        for edge_data in node_links:
+            electron_id = (
+                session.query(Lattice, Electron)
+                .filter(Lattice.id == Electron.parent_lattice_id)
+                .filter(Lattice.dispatch_id == dispatch_id)
+                .filter(Electron.transport_graph_node_id == edge_data["target"])
+                .first()
+                .Electron.id
+            )
+            parent_electron_id = (
+                session.query(Lattice, Electron)
+                .filter(Lattice.id == Electron.parent_lattice_id)
+                .filter(Lattice.dispatch_id == dispatch_id)
+                .filter(Electron.transport_graph_node_id == edge_data["source"])
+                .first()
+                .Electron.id
+            )
+
+            electron_dependency_row = ElectronDependency(
+                electron_id=electron_id,
+                parent_electron_id=parent_electron_id,
+                edge_name=edge_data["edge_name"],
+                parameter_type=edge_data["param_type"],
+                arg_index=edge_data["arg_index"],
+                created_at=dt.now(),
+            )
+
+            session.add(electron_dependency_row)
+            session.commit()
+            electron_dependency_ids.append(electron_dependency_row.id)
+
+    return electron_dependency_ids
+
+
+def update_lattices_data(
+    db: DataStore,
+    dispatch_id: str,
+    status: str,
+    updated_at: dt,
+    completed_at: dt,
+) -> None:
+    """This function updates the lattices record."""
+
+    with Session(db.engine) as session:
+        valid_update = (
+            session.query(Lattice).where(Lattice.dispatch_id == dispatch_id).first() is not None
+        )
+
+        if not valid_update:
+            raise MissingLatticeRecordError
+
+        session.execute(
+            update(Lattice)
+            .where(Lattice.dispatch_id == dispatch_id)
+            .values(status=status, updated_at=updated_at, completed_at=completed_at)
+        )
+        session.commit()
+
+
+def update_electrons_data(
+    db: DataStore,
+    parent_dispatch_id: str,
+    transport_graph_node_id: int,
+    status: str,
+    started_at: dt,
+    updated_at: dt,
+    completed_at: dt,
+) -> None:
+    """This function updates the electrons record."""
+
+    with Session(db.engine) as session:
+        parent_lattice_id = (
+            session.query(Lattice).where(Lattice.dispatch_id == parent_dispatch_id).all()[0].id
+        )
+        valid_update = (
+            session.query(Electron)
+            .where(
+                Electron.parent_lattice_id == parent_lattice_id,
+                Electron.transport_graph_node_id == transport_graph_node_id,
+            )
+            .first()
+            is not None
+        )
+        if not valid_update:
+            raise MissingElectronRecordError
+
+        session.execute(
+            update(Electron)
+            .where(
+                Electron.parent_lattice_id == parent_lattice_id,
+                Electron.transport_graph_node_id == transport_graph_node_id,
+            )
+            .values(
+                status=status,
+                started_at=started_at,
+                updated_at=updated_at,
+                completed_at=completed_at,
+            )
+        )
+        session.commit()
+
+
+def get_electron_type(node: dict) -> str:
+    """Get the electron type (to be written to DB) given the electron node data."""
+
+    if node["name"].startswith(arg_prefix):
+        return arg_prefix.strip(prefix_separator)
+
+    elif node["name"].startswith(attr_prefix):
+        return attr_prefix.strip(prefix_separator)
+
+    elif node["name"].startswith(electron_dict_prefix):
+        return electron_dict_prefix.strip(prefix_separator)
+
+    elif node["name"].startswith(electron_list_prefix):
+        return electron_list_prefix.strip(prefix_separator)
+
+    elif node["name"].startswith(generator_prefix):
+        return generator_prefix.strip(prefix_separator)
+
+    elif node["name"].startswith(parameter_prefix):
+        return parameter_prefix.strip(prefix_separator)
+
+    elif node["name"].startswith(sublattice_prefix):
+        return sublattice_prefix.strip(prefix_separator)
+
+    elif node["name"].startswith(subscript_prefix):
+        return subscript_prefix.strip(prefix_separator)
+
+    else:
+        return "function"


Test patch:
diff --git a/tests/covalent_tests/results_manager_tests/__init__.py b/tests/covalent_tests/results_manager_tests/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/covalent_tests/results_manager_tests/write_result_to_db_test.py b/tests/covalent_tests/results_manager_tests/write_result_to_db_test.py
new file mode 100644
index 000000000..210df1ad5
--- /dev/null
+++ b/tests/covalent_tests/results_manager_tests/write_result_to_db_test.py
@@ -0,0 +1,434 @@
+# Copyright 2021 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
+"""Unit tests for the module used to write the decomposed result object to the database."""
+
+from datetime import datetime as dt
+
+import pytest
+from sqlalchemy.orm import Session
+
+import covalent as ct
+from covalent._data_store.datastore import DataStore
+from covalent._data_store.models import Electron, ElectronDependency, Lattice
+from covalent._results_manager.write_result_to_db import (
+    MissingElectronRecordError,
+    MissingLatticeRecordError,
+    get_electron_type,
+    insert_electron_dependency_data,
+    insert_electrons_data,
+    insert_lattices_data,
+    update_electrons_data,
+    update_lattices_data,
+)
+from covalent._shared_files.defaults import (
+    arg_prefix,
+    attr_prefix,
+    electron_dict_prefix,
+    electron_list_prefix,
+    generator_prefix,
+    parameter_prefix,
+    prefix_separator,
+    sublattice_prefix,
+    subscript_prefix,
+)
+
+STORAGE_TYPE = "local"
+FUNCTION_FILENAME = "dispatch_source.pkl"
+FUNCTION_STRING_FILENAME = "dispatch_source.py"
+EXECUTOR_FILENAME = "executor.pkl"
+ERROR_FILENAME = "error.txt"
+INPUTS_FILENAME = "inputs.pkl"
+RESULTS_FILENAME = "result.pkl"
+VALUE_FILENAME = "value.pkl"
+STDOUT_FILENAME = "stdout.log"
+STDERR_FILENAME = "stderr.log"
+INFO_FILENAME = "info.log"
+
+
+@pytest.fixture
+def db():
+    """Instantiate and return an in-memory database."""
+
+    return DataStore(
+        db_URL="sqlite+pysqlite:///:memory:",
+        initialize_db=True,
+    )
+
+
+@pytest.fixture
+def workflow_lattice():
+    @ct.electron
+    def task_1(x, y):
+        return x * y
+
+    @ct.electron
+    def sublattice_task(z):
+        return z
+
+    @ct.electron
+    @ct.lattice
+    def task_2(z):
+        return sublattice_task(z)
+
+    @ct.lattice
+    def workflow_1(a, b):
+        res_1 = task_1(a, b)
+        return task_2(res_1, b)
+
+    workflow_1.build_graph(a=1, b=2)
+
+    return workflow_1
+
+
+def get_lattice_kwargs(
+    dispatch_id="dispatch_1",
+    name="workflow_1",
+    status="RUNNING",
+    storage_type=STORAGE_TYPE,
+    storage_path="results/dispatch_1/",
+    function_filename=FUNCTION_FILENAME,
+    function_string_filename=FUNCTION_STRING_FILENAME,
+    executor_filename=EXECUTOR_FILENAME,
+    error_filename=ERROR_FILENAME,
+    inputs_filename=INPUTS_FILENAME,
+    results_filename=RESULTS_FILENAME,
+    created_at=None,
+    updated_at=None,
+    started_at=None,
+    completed_at=None,
+):
+    """Create lattice kwargs."""
+
+    return {
+        "dispatch_id": dispatch_id,
+        "name": name,
+        "status": status,
+        "storage_type": storage_type,
+        "storage_path": storage_path,
+        "function_filename": function_filename,
+        "function_string_filename": function_string_filename,
+        "executor_filename": executor_filename,
+        "error_filename": error_filename,
+        "inputs_filename": inputs_filename,
+        "results_filename": results_filename,
+        "created_at": created_at,
+        "updated_at": updated_at,
+        "started_at": started_at,
+        "completed_at": completed_at,
+    }
+
+
+def get_electron_kwargs(
+    parent_dispatch_id="dispatch_1",
+    transport_graph_node_id=0,
+    type=parameter_prefix.strip(prefix_separator),
+    name=f"{parameter_prefix}0",
+    status="NEW_OBJ",
+    storage_type=STORAGE_TYPE,
+    storage_path="results/dispatch_1/node_0/",
+    function_filename=FUNCTION_STRING_FILENAME,
+    function_string_filename=FUNCTION_STRING_FILENAME,
+    executor_filename=EXECUTOR_FILENAME,
+    results_filename=RESULTS_FILENAME,
+    value_filename=VALUE_FILENAME,
+    attribute_name=None,
+    key=None,
+    stdout_filename=STDOUT_FILENAME,
+    stderr_filename=STDERR_FILENAME,
+    info_filename=INFO_FILENAME,
+    created_at=None,
+    updated_at=None,
+    started_at=None,
+    completed_at=None,
+):
+    """Create electron kwargs."""
+
+    return {
+        "parent_dispatch_id": parent_dispatch_id,
+        "transport_graph_node_id": transport_graph_node_id,
+        "type": type,
+        "name": name,
+        "status": status,
+        "storage_type": storage_type,
+        "storage_path": storage_path,
+        "function_filename": function_filename,
+        "function_string_filename": function_string_filename,
+        "executor_filename": executor_filename,
+        "results_filename": results_filename,
+        "value_filename": value_filename,
+        "attribute_name": attribute_name,
+        "key": key,
+        "stdout_filename": stdout_filename,
+        "stderr_filename": stderr_filename,
+        "info_filename": info_filename,
+        "created_at": created_at,
+        "updated_at": updated_at,
+        "started_at": started_at,
+        "completed_at": completed_at,
+    }
+
+
+def test_insert_lattices_data(db):
+    """Test the function that inserts the lattices data in the DB."""
+
+    lattice_ids = []
+    timestamps = []
+
+    for i in range(2):
+        cur_time = dt.now()
+        timestamps.append(cur_time)
+        lattice_args = get_lattice_kwargs(
+            dispatch_id=f"dispatch_{i + 1}",
+            name=f"workflow_{i + 1}",
+            storage_path=f"results/dispatch_{i+1}/",
+            created_at=cur_time,
+            updated_at=cur_time,
+            started_at=cur_time,
+        )
+        lattice_id = insert_lattices_data(db=db, **lattice_args)
+        lattice_ids.append(lattice_id)
+
+    assert lattice_ids == [1, 2]
+
+    with Session(db.engine) as session:
+        rows = session.query(Lattice).all()
+
+    for i, lattice in enumerate(rows):
+        assert lattice.id == i + 1
+        assert lattice.dispatch_id == f"dispatch_{i + 1}"
+        assert lattice.name == f"workflow_{i + 1}"
+        assert lattice.status == "RUNNING"
+        assert lattice.storage_type == STORAGE_TYPE
+        assert lattice.storage_path == f"results/dispatch_{i+1}/"
+        assert lattice.function_filename == FUNCTION_FILENAME
+        assert lattice.function_string_filename == FUNCTION_STRING_FILENAME
+        assert lattice.executor_filename == EXECUTOR_FILENAME
+        assert lattice.error_filename == ERROR_FILENAME
+        assert lattice.inputs_filename == INPUTS_FILENAME
+        assert lattice.results_filename == RESULTS_FILENAME
+        assert lattice.created_at == lattice.updated_at == lattice.started_at == timestamps[i]
+        assert lattice.completed_at is None
+
+        with Session(db.engine) as session:
+            rows = session.query(Lattice).where(Lattice.dispatch_id == "dispatch_3").all()
+
+        assert not rows
+
+
+def test_insert_electrons_data(db):
+    """Test the function that inserts the electron data to the Electrons table."""
+
+    cur_time = dt.now()
+    insert_lattices_data(
+        db=db, **get_lattice_kwargs(created_at=cur_time, updated_at=cur_time, started_at=cur_time)
+    )
+
+    electron_kwargs = {
+        **get_electron_kwargs(
+            created_at=cur_time,
+            updated_at=cur_time,
+        )
+    }
+
+    electron_id = insert_electrons_data(db=db, **electron_kwargs)
+    assert electron_id == 1
+
+    with Session(db.engine) as session:
+        rows = session.query(Electron).all()
+
+    assert len(rows) == 1
+
+    for electron in rows:
+        for key, value in electron_kwargs.items():
+            if key == "parent_dispatch_id":
+                assert electron.parent_lattice_id == 1
+            else:
+                assert getattr(electron, key) == value
+
+    electron_id = insert_electrons_data(db=db, **electron_kwargs)
+    assert electron_id == 2
+
+
+def test_insert_electrons_data_missing_lattice_record(db):
+    """Test the function that inserts the electron data to the Electrons table."""
+
+    cur_time = dt.now()
+    electron_kwargs = {
+        **get_electron_kwargs(
+            created_at=cur_time,
+            updated_at=cur_time,
+        )
+    }
+    with pytest.raises(MissingLatticeRecordError):
+        insert_electrons_data(db=db, **electron_kwargs)
+
+
+def test_insert_electron_dependency_data(db, workflow_lattice):
+    """Test the function that adds the electron dependencies of the lattice to the DB."""
+
+    cur_time = dt.now()
+    insert_lattices_data(
+        db=db, **get_lattice_kwargs(created_at=cur_time, updated_at=cur_time, started_at=cur_time)
+    )
+
+    electron_ids = []
+    cur_time = dt.now()
+    for (name, node_id) in [
+        ("task_1", 0),
+        (":parameter:1", 1),
+        (":parameter:2", 2),
+        (":sublattice:task_2", 3),
+        (":parameter:2", 4),
+    ]:
+        electron_kwargs = get_electron_kwargs(
+            name=name,
+            transport_graph_node_id=node_id,
+            created_at=cur_time,
+            updated_at=cur_time,
+        )
+        electron_ids.append(insert_electrons_data(db=db, **electron_kwargs))
+
+    dependency_ids = insert_electron_dependency_data(
+        db=db, dispatch_id="dispatch_1", lattice=workflow_lattice
+    )
+
+    assert dependency_ids == [1, 2, 3, 4]
+
+    with Session(db.engine) as session:
+        rows = session.query(ElectronDependency).all()
+
+    for electron_dependency in rows:
+        if electron_dependency.electron_id == 4 and electron_dependency.parent_electron_id == 1:
+            assert electron_dependency.edge_name == "arg[0]"
+            assert electron_dependency.arg_index == 0
+            assert electron_dependency.parameter_type == "arg"
+
+        elif electron_dependency.electron_id == 1 and electron_dependency.parent_electron_id == 2:
+            assert electron_dependency.edge_name == "x"
+            assert electron_dependency.arg_index == 0
+            assert electron_dependency.parameter_type == "arg"
+
+        elif electron_dependency.electron_id == 1 and electron_dependency.parent_electron_id == 3:
+            assert electron_dependency.edge_name == "y"
+            assert electron_dependency.arg_index == 1
+            assert electron_dependency.parameter_type == "arg"
+
+        elif electron_dependency.electron_id == 4 and electron_dependency.parent_electron_id == 5:
+            assert electron_dependency.edge_name == "arg[1]"
+            assert electron_dependency.arg_index == 1
+            assert electron_dependency.parameter_type == "arg"
+
+
+def test_update_lattices_data(db):
+    """Test the function that updates the lattice data."""
+
+    cur_time = dt.now()
+    with pytest.raises(MissingLatticeRecordError):
+        update_lattices_data(
+            db=db,
+            dispatch_id="dispatch_1",
+            status="COMPLETED",
+            updated_at=cur_time,
+            completed_at=cur_time,
+        )
+
+    insert_lattices_data(
+        db=db, **get_lattice_kwargs(created_at=cur_time, updated_at=cur_time, started_at=cur_time)
+    )
+    update_lattices_data(
+        db=db,
+        dispatch_id="dispatch_1",
+        status="COMPLETED",
+        updated_at=cur_time,
+        completed_at=cur_time,
+    )
+
+    with Session(db.engine) as session:
+        rows = session.query(Lattice).all()
+
+    for lattice in rows:
+        assert lattice.status == "COMPLETED"
+        assert lattice.updated_at == lattice.completed_at == cur_time
+        assert lattice.id == 1
+
+
+def test_update_electrons_data(db):
+    """Test the function that updates the data in the Electrons table."""
+
+    insert_lattices_data(
+        db=db, **get_lattice_kwargs(created_at=dt.now(), updated_at=dt.now(), started_at=dt.now())
+    )
+
+    with pytest.raises(MissingElectronRecordError):
+        update_electrons_data(
+            db=db,
+            parent_dispatch_id="dispatch_1",
+            transport_graph_node_id=0,
+            status="RUNNING",
+            started_at=dt.now(),
+            updated_at=dt.now(),
+            completed_at=None,
+        )
+
+    insert_electrons_data(db=db, **get_electron_kwargs(created_at=dt.now(), updated_at=dt.now()))
+    cur_time = dt.now()
+    update_electrons_data(
+        db=db,
+        parent_dispatch_id="dispatch_1",
+        transport_graph_node_id=0,
+        status="RUNNING",
+        started_at=cur_time,
+        updated_at=cur_time,
+        completed_at=None,
+    )
+
+    with Session(db.engine) as session:
+        rows = session.query(Electron).all()
+
+    for electron in rows:
+        assert electron.status == "RUNNING"
+        assert electron.started_at == electron.updated_at == cur_time
+
+
+def test_are_electron_dependencies_added():
+    pass
+
+
+@pytest.mark.parametrize(
+    "node_name,electron_type",
+    [
+        (f"{parameter_prefix}0", parameter_prefix.strip(prefix_separator)),
+        (f"{arg_prefix}1", arg_prefix.strip(prefix_separator)),
+        (f"{attr_prefix}2", attr_prefix.strip(prefix_separator)),
+        (f"{electron_dict_prefix}3", electron_dict_prefix.strip(prefix_separator)),
+        (f"{electron_list_prefix}4", electron_list_prefix.strip(prefix_separator)),
+        (f"{generator_prefix}5", generator_prefix.strip(prefix_separator)),
+        (f"{sublattice_prefix}sometask", sublattice_prefix.strip(prefix_separator)),
+        (f"{subscript_prefix}6", subscript_prefix.strip(prefix_separator)),
+        ("regular_task", "function"),
+    ],
+)
+def test_get_electron_type(node_name, electron_type):
+    """Test that given an electron node, the correct electron type is returned."""
+
+    node = {"name": node_name}
+
+    assert get_electron_type(node) == electron_type