Repository: topoteretes/cognee
File: cognee/tasks/ingestion/migrate_relational_database.py (Python)

Code before:

logger = logging.getLogger(__name__)


async def migrate_relational_database(
    graph_db, schema, migrate_column_data=True, schema_only=False
):
    """
    Migrates data from a relational database into a graph database.

    For each table in the schema:
      - Creates a TableType node representing the table
      - Fetches all rows and creates a TableRow node for each row
      - Links each TableRow node to its TableType node with an "is_part_of" relationship
    Then, for every foreign key defined in the schema:
      - Establishes relationships between TableRow nodes based on foreign key relationships

    Both TableType and TableRow inherit from DataPoint to maintain consistency with Cognee data model.
    """
    engine = get_migration_relational_engine()
    # Create a mapping of node_id to node objects for referencing in edge creation
    node_mapping = {}
    edge_mapping = []

    if schema_only:
        database_config = get_migration_config().to_dict()
        # Calling the ingest_database_schema function to return DataPoint subclasses
        result = await ingest_database_schema(
            database_config=database_config,
            schema_name="migrated_schema",
            max_sample_rows=5,
        )
        database_schema = result["database_schema"]
        schema_tables = result["schema_tables"]
        schema_relationships = result["relationships"]
        database_node_id = database_schema.id
        node_mapping[database_node_id] = database_schema
        for table in schema_tables:
            table_node_id = table.id
            # Add TableSchema Datapoint as a node.
            node_mapping[table_node_id] = table
            edge_mapping.append(
                (
                    table_node_id,
                    database_node_id,
                    "is_part_of",
                    dict(
                        source_node_id=table_node_id,
                        target_node_id=database_node_id,
                        relationship_name="is_part_of",
                    ),

Review comment (suggestion):
Overall a pretty nice PR already! Thank you for the contribution 

I think it would be good to create two separate functions here in this if schema_only else conditional

One should return the nodes and edges for the schema_only migration 
and the other one should return node and edge mapping for the regular migration 

It will make it more readable 

Code after:
logger = logging.getLogger(__name__)


async def migrate_relational_database(
    graph_db, schema, migrate_column_data=True, schema_only=False
):
    """
    Migrates data from a relational database into a graph database.

    For each table in the schema:
      - Creates a TableType node representing the table
      - Fetches all rows and creates a TableRow node for each row
      - Links each TableRow node to its TableType node with an "is_part_of" relationship
    Then, for every foreign key defined in the schema:
      - Establishes relationships between TableRow nodes based on foreign key relationships

    Both TableType and TableRow inherit from DataPoint to maintain consistency with Cognee data model.
    """
    # Create a mapping of node_id to node objects for referencing in edge creation
    node_mapping = {}
    edge_mapping = []

    if schema_only:
        node_mapping, edge_mapping = await schema_only_ingestion()

    else:
        node_mapping, edge_mapping = await complete_database_ingestion(schema, migrate_column_data)

    def _remove_duplicate_edges(edge_mapping):
        seen = set()
        unique_original_shape = []

        for tup in edge_mapping:
            # We go through all the tuples in the edge_mapping and we only add unique tuples to the list
            # To eliminate duplicate edges.
            source_id, target_id, rel_name, rel_dict = tup
            # We need to convert the dictionary to a frozenset to be able to compare values for it
            rel_dict_hashable = frozenset(sorted(rel_dict.items()))
            hashable_tup = (source_id, target_id, rel_name, rel_dict_hashable)

            # We use the seen set to keep track of unique edges
            if hashable_tup not in seen:
                # A list that has frozensets elements instead of dictionaries is needed to be able to compare values
                seen.add(hashable_tup)
                # append the original tuple shape (with the dictionary) if it's the first time we see it
                unique_original_shape.append(tup)

        return unique_original_shape

    # Add all nodes and edges to the graph
    # NOTE: Nodes and edges have to be added in batch for speed optimization, Especially for NetworkX.