Merge pull request #17 from Kaligo/fix/data-24484-array-to-json

fix: convert array types to JSON in fast-sync flow [DATA-24484]

Author: kidlab

tap_postgres/sync_strategies/fast_sync_rds.py

@@ -94,6 +94,34 @@ def _get_metadata_column_sql(self, column_name: str) -> str:
 
         raise ValueError(f"Unknown metadata column: {column_name}")
 
+    def _is_array_column(self, column_name: str, md_map: Dict) -> bool:
+        """
+        Check if a column is an array type based on metadata map.
+        """
+        column_metadata = md_map.get(("properties", column_name), {})
+        sql_datatype = column_metadata.get("sql-datatype", "")
+        # Array types in PostgreSQL have '[]' suffix in sql-datatype
+        return sql_datatype.endswith("[]")
+
+    def _convert_array_column_to_json(self, column_name: str) -> str:
+        """
+        Convert an array column to JSON format using PostgreSQL's array_to_json function.
+
+        This converts PostgreSQL array format (e.g., {fashion}) to JSON format (e.g., ["fashion"]).
+        The conversion happens before CSV export, so the exported CSV will contain JSON arrays
+        instead of PostgreSQL array format.
+
+        Args:
+            column_name: Column name to convert
+
+        Returns:
+            SQL expression that converts array to JSON format
+        """
+        column_identifier = post_db.prepare_columns_sql(column_name)
+        # array_to_json converts PostgreSQL arrays to JSON arrays
+        # Returns NULL for NULL arrays (which is correct)
+        return f"array_to_json({column_identifier}) AS {column_identifier}"
+
     def _build_sorted_column_expressions(
         self, desired_columns: List[str], md_map: Dict
     ) -> List[str]:
@@ -104,6 +132,9 @@ def _build_sorted_column_expressions(
         This ensures the exported data column order matches the table column
         order exactly.
 
+        Array columns are automatically converted from PostgreSQL array format
+        (e.g., {fashion}) to JSON format (e.g., ["fashion"]) using array_to_json().
+
         Args:
             desired_columns: List of desired column names from the source table
             md_map: Metadata map for column transformations
@@ -116,11 +147,18 @@ def _build_sorted_column_expressions(
         # Sort columns to ensure the output CSV headers match target's schema order.
         all_column_names.sort()
 
-        return [
-            self._get_metadata_column_sql(name) if name in metadata_column_names
-            else post_db.prepare_columns_for_select_sql(name, md_map=md_map)
-            for name in all_column_names
-        ]
+        column_expressions = []
+        for name in all_column_names:
+            if name in metadata_column_names:
+                column_expressions.append(self._get_metadata_column_sql(name))
+            elif self._is_array_column(name, md_map):
+                column_expressions.append(self._convert_array_column_to_json(name))
+            else:
+                column_expressions.append(
+                    post_db.prepare_columns_for_select_sql(name, md_map=md_map)
+                )
+
+        return column_expressions
 
     def _build_select_query(  # pylint: disable=too-many-arguments,too-many-positional-arguments
         self,

tests/unit/test_fast_sync_rds.py

@@ -218,7 +218,6 @@ def test_sync_table_full_column_ordering(self, mock_open_conn):
             f"Columns should be in alphabetical order. Found positions: {dict(zip(expected_columns, column_positions))}",
         )
 
-
     @patch("tap_postgres.sync_strategies.fast_sync_rds.post_db.open_connection")
     def test_sync_table_full_no_prefix(self, mock_open_conn):
         """Test sync_table_full with empty prefix - verifies S3 path generation"""
@@ -533,3 +532,141 @@ def test_do_sync_incremental_backward_compatibility_with_fast_sync_s3_info(
         self.assertIn(
             "fast_sync_s3_info", result["bookmarks"]["test_schema-test_table"]
         )
+
+    @patch("tap_postgres.sync_strategies.fast_sync_rds.post_db.open_connection")
+    def test_sync_table_full_with_array_columns(self, mock_open_conn):
+        """Test sync_table_full converts array columns to JSON format in export query"""
+        # Test with mixed array types (text, integer, date, timestamp) and non-array columns
+        desired_columns = ["id", "_text", "name", "_date", "_int", "_timestamp"]
+        md_map_with_arrays = {
+            (): {"schema-name": "test_schema"},
+            ("properties", "id"): {"sql-datatype": "integer"},
+            ("properties", "_text"): {"sql-datatype": "text[]"},
+            ("properties", "name"): {"sql-datatype": "varchar"},
+            ("properties", "_date"): {"sql-datatype": "date[]"},
+            ("properties", "_int"): {"sql-datatype": "integer[]"},
+            ("properties", "_timestamp"): {
+                "sql-datatype": "timestamp without time zone[]"
+            },
+        }
+
+        mock_conn, mock_cursor = self._setup_mock_connection()
+        mock_open_conn.return_value.__enter__.return_value = mock_conn
+
+        self.fast_sync_rds_strategy.sync_table_full(
+            stream=self.stream,
+            state=self.state,
+            desired_columns=desired_columns,
+            md_map=md_map_with_arrays,
+        )
+
+        # Verify export query contains array_to_json for array columns
+        export_query = self._extract_export_query(mock_cursor)
+        self.assertIsNotNone(export_query)
+
+        # Array columns should use array_to_json conversion
+        self.assertIn("array_to_json", export_query)
+
+        # Count array_to_json occurrences - should match number of array columns
+        array_columns = ["_text", "_date", "_int", "_timestamp"]
+        array_to_json_count = export_query.count("array_to_json")
+        self.assertEqual(
+            array_to_json_count,
+            len(array_columns),
+            "All array columns should use array_to_json",
+        )
+
+        # Verify each array column is present and wrapped with array_to_json
+        # (accounting for possible whitespace variations in SQL formatting)
+        for col in array_columns:
+            self.assertIn(col, export_query)
+            self.assertRegex(
+                export_query,
+                rf'array_to_json\s*\(\s*["\s]*{col}["\s]*\)',
+                f"Array column '{col}' should be wrapped with array_to_json",
+            )
+
+        # Non-array columns should not be wrapped with array_to_json
+        non_array_columns = ["id", "name"]
+        for col in non_array_columns:
+            self.assertIn(f'"{col}"', export_query)
+            # Verify they are not wrapped
+            self.assertNotRegex(
+                export_query,
+                rf'array_to_json\s*\(\s*["\s]*{col}["\s]*\)',
+                f"Non-array column '{col}' should not use array_to_json",
+            )
+
+    @patch("tap_postgres.sync_strategies.fast_sync_rds.post_db.open_connection")
+    @patch("tap_postgres.sync_strategies.fast_sync_rds.singer.get_bookmark")
+    def test_sync_table_incremental_with_array_columns(
+        self, mock_get_bookmark, mock_open_conn
+    ):
+        """Test sync_table_incremental converts array columns to JSON format"""
+        mock_get_bookmark.side_effect = lambda state, stream_id, key: {
+            "version": 1234567890,
+            "replication_key_value": "100",
+        }.get(key)
+
+        desired_columns = ["id", "_text", "name", "_int"]
+        md_map_with_arrays = {
+            (): {"schema-name": "test_schema", "replication-key": "id"},
+            ("properties", "id"): {"sql-datatype": "integer"},
+            ("properties", "_text"): {"sql-datatype": "text[]"},
+            ("properties", "name"): {"sql-datatype": "varchar"},
+            ("properties", "_int"): {"sql-datatype": "integer[]"},
+        }
+
+        mock_replication_key_result = MagicMock()
+        mock_replication_key_result.__getitem__.return_value = "200"
+
+        mock_conn, mock_cursor = self._setup_mock_connection(
+            export_result=self._create_mock_export_result(rows=25, bytes_uploaded=1250),
+            replication_key_result=mock_replication_key_result,
+        )
+        mock_open_conn.return_value.__enter__.return_value = mock_conn
+
+        self.fast_sync_rds_strategy.sync_table_incremental(
+            stream=self.stream,
+            state=self.state,
+            desired_columns=desired_columns,
+            md_map=md_map_with_arrays,
+            replication_key="id",
+            replication_key_value="100",
+        )
+
+        # Verify export query contains array_to_json for array columns
+        export_query = self._extract_export_query(mock_cursor)
+        self.assertIsNotNone(export_query)
+
+        # Array columns should use array_to_json
+        self.assertIn("array_to_json", export_query)
+        self.assertIn("_text", export_query)
+        self.assertIn("_int", export_query)
+
+        # Verify array columns are wrapped correctly (with flexible whitespace)
+        self.assertRegex(
+            export_query,
+            r'array_to_json\s*\(\s*["\s]*_text["\s]*\)',
+            "Array column '_text' should be wrapped with array_to_json",
+        )
+        self.assertRegex(
+            export_query,
+            r'array_to_json\s*\(\s*["\s]*_int["\s]*\)',
+            "Array column '_int' should be wrapped with array_to_json",
+        )
+
+        # Non-array columns should not be wrapped
+        self.assertIn('"id"', export_query)
+        self.assertIn('"name"', export_query)
+        # Verify they are not wrapped with array_to_json
+        self.assertNotRegex(
+            export_query,
+            r'array_to_json\s*\(\s*["\s]*id["\s]*\)',
+            "Non-array column 'id' should not use array_to_json",
+        )
+        self.assertNotRegex(
+            export_query,
+            r'array_to_json\s*\(\s*["\s]*name["\s]*\)',
+            "Non-array column 'name' should not use array_to_json",
+        )