Repository: topoteretes/cognee
File: cognee/infrastructure/databases/graph/neo4j_driver/adapter.py (Python)

Code before:
        --------

            A tuple of nodes and edges data.
        """
        query_nodes = "MATCH (n) RETURN collect(n) AS nodes"
        nodes = await self.query(query_nodes)

        query_edges = "MATCH (n)-[r]->(m) RETURN collect([n, r, m]) AS elements"
        edges = await self.query(query_edges)

        return (nodes, edges)

    async def get_graph_data(self):
        """
        Retrieve comprehensive data about nodes and relationships within the graph.

        Returns:
        --------

            A tuple containing two lists: nodes and edges with their properties.
        """
        logger.info("Starting full graph data retrieval from Neo4j")
        import time

        start_time = time.time()

        try:
            # Retrieve nodes
            logger.debug("Executing nodes query")
            query = "MATCH (n) RETURN ID(n) AS id, labels(n) AS labels, properties(n) AS properties"
            result = await self.query(query)

            logger.debug(f"Raw nodes query returned {len(result)} records")

            nodes = []
            processed_nodes = 0
            failed_nodes = 0

            for record in result:
                try:
                    if (
                        "properties" in record
                        and record["properties"]
                        and "id" in record["properties"]
                    ):
                        nodes.append(
                            (
                                record["properties"]["id"],
                                record["properties"],
                            )
                        )

Review comment (suggestion):
I dont wouldn't do the logging of this. This as I understand is doing some kind of a structural check on the retrieved nodes and edges and doesn' log anything. We cannot have missing properties in the graph, or at least it shouldn't be tested here at every retrieval. In this part (get_graph_data) I would measure the time it takes to retreive these, that part is useful.

Code after:
        --------

            A tuple of nodes and edges data.
        """
        query_nodes = "MATCH (n) RETURN collect(n) AS nodes"
        nodes = await self.query(query_nodes)

        query_edges = "MATCH (n)-[r]->(m) RETURN collect([n, r, m]) AS elements"
        edges = await self.query(query_edges)

        return (nodes, edges)

    async def get_graph_data(self):
        """
        Retrieve comprehensive data about nodes and relationships within the graph.

        Returns:
        --------

            A tuple containing two lists: nodes and edges with their properties.
        """
        import time

        start_time = time.time()

        try:
            # Retrieve nodes
            query = "MATCH (n) RETURN ID(n) AS id, labels(n) AS labels, properties(n) AS properties"
            result = await self.query(query)

            nodes = []
            for record in result:
                nodes.append(
                    (
                        record["properties"]["id"],
                        record["properties"],
                    )
                )

            # Retrieve edges
            query = """
            MATCH (n)-[r]->(m)
            RETURN ID(n) AS source, ID(m) AS target, TYPE(r) AS type, properties(r) AS properties
            """
            result = await self.query(query)

            edges = []
            for record in result:
                edges.append(
                    (
                        record["properties"]["source_node_id"],