Bug fixes

Author: BenConstable9

text_2_sql/data_dictionary/data_dictionary_creator.py

@@ -20,18 +20,24 @@
 
 
 class ForeignKeyRelationship(BaseModel):
-    source_column: str = Field(..., alias="Column")
+    column: str = Field(..., alias="Column")
     foreign_column: str = Field(..., alias="ForeignColumn")
 
+    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)
+
 
 class EntityRelationship(BaseModel):
+    entity: str = Field(..., alias="Entity")
     foreign_entity: str = Field(..., alias="ForeignEntity")
     foreign_keys: list[ForeignKeyRelationship] = Field(..., alias="ForeignKeys")
 
-    def pivot(self, entity: str):
+    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)
+
+    def pivot(self):
         """A method to pivot the entity relationship."""
         return EntityRelationship(
-            foreign_entity=entity,
+            entity=self.foreign_entity,
+            foreign_entity=self.entity,
             foreign_keys=[
                 ForeignKeyRelationship(
                     source_column=foreign_key.foreign_column,
@@ -53,7 +59,7 @@ def from_sql_row(cls, row, columns):
             foreign_entity=result["ForeignEntity"],
             foreign_keys=[
                 ForeignKeyRelationship(
-                    source_column=result["Column"],
+                    column=result["Column"],
                     foreign_column=result["ForeignColumn"],
                 )
             ],
@@ -100,7 +106,7 @@ class EntityItem(BaseModel):
         None, alias="EntityRelationships"
     )
 
-    complete_entity_relationship_graph = Optional[str] = Field(
+    complete_entity_relationship_graph: Optional[str] = Field(
         None, alias="CompleteEntityRelationshipGraph"
     )
 

text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py

@@ -25,10 +25,8 @@ def __init__(
         excluded_entities.extend(
             ["dbo.BuildVersion", "dbo.ErrorLog", "sys.database_firewall_rules"]
         )
-        cls = super().__init__(entities, excluded_entities, single_file)
-        cls.database = os.environ["Text2Sql__DatabaseName"]
-
-        return cls
+        super().__init__(entities, excluded_entities, single_file)
+        self.database = os.environ["Text2Sql__DatabaseName"]
 
     """A class to extract data dictionary information from a SQL Server database."""
 
@@ -85,12 +83,12 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
     AND c.TABLE_NAME = '{entity.name}';"""
 
     @property
-    def extract_entity_relationships_sql_query() -> str:
+    def extract_entity_relationships_sql_query(self) -> 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,
+            fk_col.name AS [Column],
             pk_col.name AS ForeignColumn
         FROM
             sys.foreign_keys AS fk