Store warehouse and entity info

Author: BenConstable9

text_2_sql/data_dictionary/data_dictionary_creator.py

@@ -152,6 +152,9 @@ def __init__(
         self.entity_relationships = {}
         self.relationship_graph = nx.DiGraph()
 
+        self.warehouse = None
+        self.database = None
+
         load_dotenv(find_dotenv())
 
     @property
@@ -321,6 +324,11 @@ async def extract_entities_with_descriptions(self) -> list[EntityItem]:
                 if entity.entity not in self.excluded_entities
             ]
 
+        # Add warehouse and database to entities
+        for entity in all_entities:
+            entity.warehouse = self.warehouse
+            entity.database = self.database
+
         return all_entities
 
     async def extract_column_distinct_values(

text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 from data_dictionary_creator import DataDictionaryCreator, EntityItem
 import asyncio
+import os
 
 
 class SqlServerDataDictionaryCreator(DataDictionaryCreator):
@@ -24,7 +25,10 @@ def __init__(
         excluded_entities.extend(
             ["dbo.BuildVersion", "dbo.ErrorLog", "sys.database_firewall_rules"]
         )
-        return super().__init__(entities, excluded_entities, single_file)
+        cls = super().__init__(entities, excluded_entities, single_file)
+        cls.database = os.environ["Text2Sql__DatabaseName"]
+
+        return cls
 
     """A class to extract data dictionary information from a SQL Server database."""