Add graph traversal

Author: BenConstable9

deploy_ai_search/text_2_sql_schema_store.py

@@ -133,29 +133,29 @@ def get_index_fields(self) -> list[SearchableField]:
                 # This is needed to enable semantic searching against the column names as complex field types are not used.
             ),
             ComplexField(
-                name="ImmediateEntityRelationships",
+                name="EntityRelationships",
                 collection=True,
                 fields=[
                     SearchableField(
-                        name="Name",
+                        name="ForeignEntity",
                         type=SearchFieldDataType.String,
                     ),
                     ComplexField(
                         name="ForeignKeys",
                         collection=True,
                         fields=[
                             SearchableField(
-                                name="SourceColumnName", type=SearchFieldDataType.String
+                                name="Column", type=SearchFieldDataType.String
                             ),
                             SearchableField(
-                                name="TargetColumnName", type=SearchFieldDataType.String
+                                name="ForeignColumn", type=SearchFieldDataType.String
                             ),
                         ],
                     ),
                 ],
             ),
             SimpleField(
-                name="CompleteEntityRelationships",
+                name="CompleteEntityRelationshipGraph",
                 type=SearchFieldDataType.String,
                 collection=True,
             ),

text_2_sql/data_dictionary/data_dictionary_creator.py

@@ -14,10 +14,52 @@
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 import random
 import re
+import networkx as nx
 
 logging.basicConfig(level=logging.INFO)
 
 
+class ForeignKeyRelationship(BaseModel):
+    source_column: str = Field(..., alias="Column")
+    foreign_column: str = Field(..., alias="ForeignColumn")
+
+
+class EntityRelationship(BaseModel):
+    foreign_entity: str = Field(..., alias="ForeignEntity")
+    foreign_keys: list[ForeignKeyRelationship] = Field(..., alias="ForeignKeys")
+
+    def pivot(self, entity: str):
+        """A method to pivot the entity relationship."""
+        return EntityRelationship(
+            foreign_entity=entity,
+            foreign_keys=[
+                ForeignKeyRelationship(
+                    source_column=foreign_key.foreign_column,
+                    foreign_column=foreign_key.source_column,
+                )
+                for foreign_key in self.foreign_keys
+            ],
+        )
+
+    def add_foreign_key(self, foreign_key: ForeignKeyRelationship):
+        """A method to add a foreign key to the entity relationship."""
+        self.foreign_keys.append(foreign_key)
+
+    @classmethod
+    def from_sql_row(cls, row, columns):
+        """A method to create an EntityRelationship from a SQL row."""
+        result = dict(zip(columns, row))
+        return cls(
+            foreign_entity=result["ForeignEntity"],
+            foreign_keys=[
+                ForeignKeyRelationship(
+                    source_column=result["Column"],
+                    foreign_column=result["ForeignColumn"],
+                )
+            ],
+        )
+
+
 class ColumnItem(BaseModel):
     """A class to represent a column item."""
 
@@ -38,7 +80,7 @@ def from_sql_row(cls, row, columns):
         result = dict(zip(columns, row))
         return cls(
             name=result["Name"],
-            type=result["DataType"],
+            data_type=result["DataType"],
             definition=result["Definition"],
         )
 
@@ -51,6 +93,16 @@ class EntityItem(BaseModel):
     name: str = Field(..., alias="Name", exclude=True)
     entity_schema: str = Field(..., alias="Schema", exclude=True)
     entity_name: Optional[str] = Field(default=None, alias="EntityName")
+    database: Optional[str] = Field(default=None, alias="Database")
+    warehouse: Optional[str] = Field(default=None, alias="Warehouse")
+
+    entity_relationships: Optional[list[EntityRelationship]] = Field(
+        None, alias="EntityRelationships"
+    )
+
+    complete_entity_relationship_graph = Optional[str] = Field(
+        None, alias="CompleteEntityRelationshipGraph"
+    )
 
     columns: Optional[list[ColumnItem]] = Field(
         ..., alias="Columns", default_factory=list
@@ -97,6 +149,9 @@ def __init__(
         self.single_file = single_file
         self.generate_descriptions = generate_descriptions
 
+        self.entity_relationships = {}
+        self.relationship_graph = nx.DiGraph()
+
         load_dotenv(find_dotenv())
 
     @property
@@ -119,6 +174,13 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
 
         Must return 3 columns: Name, DataType, Definition."""
 
+    @property
+    @abstractmethod
+    def extract_entity_relationships_sql_query(self) -> str:
+        """An abstract method to extract entity relationships from a database.
+
+        Must return 4 columns: Entity, ForeignEntity, Column, ForeignColumn."""
+
     def extract_distinct_values_sql_query(
         self, entity: EntityItem, column: ColumnItem
     ) -> str:
@@ -165,6 +227,72 @@ async def query_entities(
 
         return results
 
+    async def extract_entity_relationships(self) -> list[EntityRelationship]:
+        """A method to extract entity relationships from a database.
+
+        Returns:
+            list[EntityRelationships]: The list of entity relationships."""
+
+        relationships = await self.query_entities(
+            self.extract_entity_relationships_sql_query, cast_to=EntityRelationship
+        )
+
+        # Duplicate relationships to create a complete graph
+
+        for relationship in relationships:
+            if relationship.entity not in self.entity_relationships:
+                self.entity_relationships[relationship.entity] = {
+                    relationship.foreign_entity: relationship
+                }
+            else:
+                self.entity_relationships[relationship.entity][
+                    relationship.foreign_entity
+                ].add_foreign_key(relationship.foreign_keys[0])
+
+            if relationship.foreign_entity not in self.entity_relationships:
+                self.entity_relationships[relationship.foreign_entity] = {
+                    relationship.entity: relationship.pivot()
+                }
+            else:
+                self.entity_relationships[relationship.foreign_entity][
+                    relationship.entity
+                ].add_foreign_key(relationship.pivot().foreign_keys[0])
+
+    async def build_entity_relationship_graph(self) -> nx.DiGraph:
+        """A method to build a complete entity relationship graph."""
+
+        for entity, foreign_entities in self.entity_relationships.items():
+            for foreign_entity, relationship in foreign_entities.items():
+                self.relationship_graph.add_edge(
+                    entity, foreign_entity, relationship=relationship
+                )
+
+    def get_entity_relationships_from_graph(
+        self, entity: str, path=None, result=None, visited=None
+    ) -> nx.DiGraph:
+        if entity not in self.relationship_graph:
+            return None
+
+        if path is None:
+            path = []
+        if result is None:
+            result = []
+        if visited is None:
+            visited = set()
+
+        # Mark the current node as visited
+        visited.add(entity)
+
+        # For each successor (neighbor in the directed path)
+        for successor in self.relationship_graph.successors(entity):
+            new_path = path + [f"{entity} -> {successor}"]
+            result.append(" -> ".join(new_path))  # Add the path as a string
+            self.get_entity_relationships_from_graph(
+                self.relationship_graph, successor, new_path, result, visited
+            )
+
+        return result
+
     async def extract_entities_with_descriptions(self) -> list[EntityItem]:
         """A method to extract entities with descriptions from a database.
 
@@ -420,12 +548,27 @@ async def build_entity_entry(self, entity: EntityItem) -> EntityItem:
         if self.generate_descriptions:
             await self.generate_entity_description(entity)
 
+        # add in relationships
+        if entity.entity in self.entity_relationships:
+            entity.entity_relationships = list(
+                self.entity_relationships[entity.entity].values()
+            )
+
+        # add in the graph traversal
+        entity.complete_entity_relationship_graph = (
+            self.get_entity_relationships_from_graph(entity.entity)
+        )
+
         return entity
 
     async def create_data_dictionary(self):
         """A method to build a data dictionary from a database. Writes to file."""
         entities = await self.extract_entities_with_descriptions()
 
+        await self.extract_entity_relationships()
+
+        await self.build_entity_relationship_graph()
+
         entity_tasks = []
         for entity in entities:
             entity_tasks.append(self.build_entity_entry(entity))

text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py

@@ -34,7 +34,7 @@ def extract_table_entities_sql_query(self) -> str:
         return """SELECT
     t.TABLE_NAME AS Entity,
     t.TABLE_SCHEMA AS EntitySchema,
-    CAST(ep.value AS NVARCHAR(500)) AS Description
+    CAST(ep.value AS NVARCHAR(500)) AS Definition
 FROM
     INFORMATION_SCHEMA.TABLES t
 LEFT JOIN
@@ -80,6 +80,30 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
     c.TABLE_SCHEMA = '{entity.entity_schema}'
     AND c.TABLE_NAME = '{entity.name}';"""
 
+    @property
+    def extract_entity_relationships_sql_query() -> str:
+        """A property to extract entity relationships from a SQL Server database."""
+        return """SELECT
+            fk_tab.name AS Entity,
+            pk_tab.name AS ForeignEntity,
+            fk_col.name AS Column,
+            pk_col.name AS ForeignColumn
+        FROM
+            sys.foreign_keys AS fk
+        INNER JOIN
+            sys.foreign_key_columns AS fkc ON fk.object_id = fkc.constraint_object_id
+        INNER JOIN
+            sys.tables AS fk_tab ON fk_tab.object_id = fk.parent_object_id
+        INNER JOIN
+            sys.tables AS pk_tab ON pk_tab.object_id = fk.referenced_object_id
+        INNER JOIN
+            sys.columns AS fk_col ON fkc.parent_object_id = fk_col.object_id AND fkc.parent_column_id = fk_col.column_id
+        INNER JOIN
+            sys.columns AS pk_col ON fkc.referenced_object_id = pk_col.object_id AND fkc.referenced_column_id = pk_col.column_id
+        ORDER BY
+            Entity, ForeignEntity;
+        """
+
 
 if __name__ == "__main__":
     data_dictionary_creator = SqlServerDataDictionaryCreator()