run black

Author: dongwonmoon

data_scribe/components/writers/postgres_comment_writer.py

@@ -78,7 +78,9 @@ def write(self, catalog_data: Dict[str, Any], **kwargs):
                     "'", "''"
                 )  # Basic SQL escaping
 
-                logger.info(f"  - Writing comment for VIEW: {schema_name}.{view_name}")
+                logger.info(
+                    f"  - Writing comment for VIEW: {schema_name}.{view_name}"
+                )
                 query = f'COMMENT ON VIEW "{schema_name}"."{view_name}" IS %s;'
                 cursor.execute(query, (description,))
 
@@ -103,6 +105,10 @@ def write(self, catalog_data: Dict[str, Any], **kwargs):
             logger.info("Successfully wrote all comments to PostgreSQL.")
 
         except Exception as e:
-            logger.error(f"Error writing comments to PostgreSQL: {e}", exc_info=True)
+            logger.error(
+                f"Error writing comments to PostgreSQL: {e}", exc_info=True
+            )
             db_connector.connection.rollback()  # Rollback on failure
-            raise WriterError(f"Error writing comments to PostgreSQL: {e}") from e
+            raise WriterError(
+                f"Error writing comments to PostgreSQL: {e}"
+            ) from e

data_scribe/core/db_workflow.py

@@ -61,12 +61,12 @@ def run(self):
         """
         # 1. Determine which database and LLM profiles to use.
         # Priority is given to CLI arguments, falling back to defaults in the config file.
-        db_profile_name = self.db_profile_name or self.config.get("default", {}).get(
-            "db"
-        )
-        llm_profile_name = self.llm_profile_name or self.config.get("default", {}).get(
-            "llm"
-        )
+        db_profile_name = self.db_profile_name or self.config.get(
+            "default", {}
+        ).get("db")
+        llm_profile_name = self.llm_profile_name or self.config.get(
+            "default", {}
+        ).get("llm")
 
         if not db_profile_name or not llm_profile_name:
             logger.error(
@@ -97,7 +97,9 @@ def run(self):
 
         try:
             # Initialize the writer based on the output profile.
-            writer_params = self.config["output_profiles"][self.output_profile_name]
+            writer_params = self.config["output_profiles"][
+                self.output_profile_name
+            ]
             writer_type = writer_params.pop("type")
             writer = get_writer(writer_type)
 

data_scribe/tests/unit/writers/test_postgres_comment_writer.py

@@ -1,6 +1,7 @@
 """
 Unit tests for the PostgresCommentWriter.
 """
+
 import pytest
 from unittest.mock import MagicMock
 
@@ -85,7 +86,9 @@ def test_postgres_comment_writer_write_success(
 def test_postgres_comment_writer_missing_db_connector():
     """Tests that ConfigError is raised if db_connector is missing."""
     writer = PostgresCommentWriter()
-    with pytest.raises(ConfigError, match="PostgresCommentWriter requires 'db_connector'"):
+    with pytest.raises(
+        ConfigError, match="PostgresCommentWriter requires 'db_connector'"
+    ):
         writer.write({}, some_other_arg="value")
 
 
@@ -94,7 +97,10 @@ def test_postgres_comment_writer_wrong_db_connector_type():
     writer = PostgresCommentWriter()
     mock_connector = MagicMock(spec=PostgresConnector)
     mock_connector.connection = None  # Simulate not connected
-    with pytest.raises(ConfigError, match="PostgresCommentWriter is only compatible with 'postgres'"):
+    with pytest.raises(
+        ConfigError,
+        match="PostgresCommentWriter is only compatible with 'postgres'",
+    ):
         writer.write({}, db_connector=MagicMock())  # Pass a generic mock
 
 
@@ -110,11 +116,15 @@ def test_postgres_comment_writer_db_error_rollback(
     mock_postgres_connector, mock_catalog_data
 ):
     """Tests that changes are rolled back on database error."""
-    mock_postgres_connector.cursor.execute.side_effect = Exception("DB write error")
+    mock_postgres_connector.cursor.execute.side_effect = Exception(
+        "DB write error"
+    )
 
     writer = PostgresCommentWriter()
-    with pytest.raises(WriterError, match="Error writing comments to PostgreSQL"):
+    with pytest.raises(
+        WriterError, match="Error writing comments to PostgreSQL"
+    ):
         writer.write(mock_catalog_data, db_connector=mock_postgres_connector)
 
     mock_postgres_connector.connection.commit.assert_not_called()
-    mock_postgres_connector.connection.rollback.assert_called_once()
+    mock_postgres_connector.connection.rollback.assert_called_once()