Add data parameter to _detect_relationships method

sdv/metadata/multi_table.py

@@ -497,8 +497,14 @@ def _validate_all_tables_connected(self, parent_map, child_map):
                 f'The relationships in the dataset are disjointed. {table_msg}'
             )
 
-    def _detect_relationships(self):
-        """Automatically detect relationships between tables."""
+    def _detect_relationships(self, data=None):
+        """Automatically detect relationships between tables.
+
+        Args:
+            data (dict):
+                Dictionary of table names to dataframes.
+                NOTE: this is only used in SDV-Enterprise.
+        """
         for parent_candidate in self.tables.keys():
             primary_key = self.tables[parent_candidate].primary_key
             for child_candidate in self.tables.keys() - {parent_candidate}:
@@ -552,7 +558,7 @@ def detect_from_dataframes(self, data):
         for table_name, dataframe in data.items():
             self.detect_table_from_dataframe(table_name, dataframe)
 
-        self._detect_relationships()
+        self._detect_relationships(data)
 
     def detect_table_from_csv(self, table_name, filepath, read_csv_parameters=None):
         """Detect the metadata for a table from a csv file.
@@ -579,7 +585,9 @@ def detect_from_csvs(self, folder_name, read_csv_parameters=None):
         Args:
             folder_name (str):
                 Name of the folder to detect the metadata from.
-
+            read_csv_parameters (dict):
+                A python dictionary of with string and value accepted by ``pandas.read_csv``
+                function. Defaults to ``None``.
         """
         folder_path = Path(folder_name)
 
@@ -595,7 +603,12 @@ def detect_from_csvs(self, folder_name, read_csv_parameters=None):
             table_name = csv_file.stem
             self.detect_table_from_csv(table_name, str(csv_file), read_csv_parameters)
 
-        self._detect_relationships()
+        read_csv_parameters = read_csv_parameters or {}
+        data = {
+            csv_file.stem: pd.read_csv(str(csv_file), **read_csv_parameters)
+            for csv_file in csv_files
+        }
+        self._detect_relationships(data)
 
     def set_primary_key(self, table_name, column_name):
         """Set the primary key of a table.

tests/unit/metadata/test_multi_table.py

@@ -2415,6 +2415,10 @@ def test_detect_from_csvs(self, tmp_path):
         assert instance.detect_table_from_csv.call_count == 2
 
         instance._detect_relationships.assert_called_once()
+        table1 = instance._detect_relationships.call_args[0][0]['table1']
+        table2 = instance._detect_relationships.call_args[0][0]['table2']
+        pd.testing.assert_frame_equal(table1, data1)
+        pd.testing.assert_frame_equal(table2, data2)
 
     def test_detect_from_csvs_no_csv(self, tmp_path):
         """Test the ``detect_from_csvs`` method with no csv file in the folder."""
@@ -2426,13 +2430,11 @@ def test_detect_from_csvs_no_csv(self, tmp_path):
             json_file.write('{"key": "value"}')
 
         # Run and Assert
-        expected_message = re.escape("No CSV files detected in the folder '{}'.".format(tmp_path))
+        expected_message = re.escape(f"No CSV files detected in the folder '{tmp_path}'.")
         with pytest.raises(ValueError, match=expected_message):
             instance.detect_from_csvs(tmp_path)
 
-        expected_message_folder = re.escape(
-            "The folder '{}' does not exist.".format('not_a_folder')
-        )
+        expected_message_folder = re.escape(f"The folder '{'not_a_folder'}' does not exist.")
         with pytest.raises(ValueError, match=expected_message_folder):
             instance.detect_from_csvs('not_a_folder')
 
@@ -2516,14 +2518,15 @@ def test_detect_from_dataframes(self):
 
         guests_table = pd.DataFrame()
         hotels_table = pd.DataFrame()
+        data = {'guests': guests_table, 'hotels': hotels_table}
 
         # Run
-        metadata.detect_from_dataframes(data={'guests': guests_table, 'hotels': hotels_table})
+        metadata.detect_from_dataframes(data)
 
         # Assert
         metadata.detect_table_from_dataframe.assert_any_call('guests', guests_table)
         metadata.detect_table_from_dataframe.assert_any_call('hotels', hotels_table)
-        metadata._detect_relationships.assert_called_once()
+        metadata._detect_relationships.assert_called_once_with(data)
 
     def test_detect_from_dataframes_no_dataframes(self):
         """Test ``detect_from_dataframes`` with no dataframes in the input.