Add mem alloc handling

terrawrap/utils/cli.py

@@ -190,19 +190,53 @@ def _execute_command(
         # pylint: disable=consider-using-with
         process = subprocess.Popen(args, *pargs, **kwargs)
 
+        # Read in larger chunks to avoid memory allocation issues
+        buffer_size = 8192  # 8KB buffer
         while True:
-            output = stdout_read.read(1).decode(errors="replace")
-
-            if output == "" and process.poll() is not None:
-                break
-
-            if print_output and output:
-                print(output, end="", flush=True)
+            try:
+                chunk = stdout_read.read(buffer_size)
+                if not chunk and process.poll() is not None:
+                    break
+
+                if chunk and print_output:
+                    try:
+                        decoded_chunk = chunk.decode(errors="replace")
+                        print(decoded_chunk, end="", flush=True)
+                    except UnicodeDecodeError:
+                        # Handle any decoding issues gracefully
+                        print(chunk.decode(errors="ignore"), end="", flush=True)
+
+            except OSError as e:
+                if e.errno == 12:  # Cannot allocate memory
+                    logger.warning("Memory allocation issue while reading output, reducing buffer size")
+                    buffer_size = max(1024, buffer_size // 2)  # Reduce buffer size but keep minimum
+                    continue
+                else:
+                    raise
 
         exit_code = process.poll()
 
+        # Read the complete output in a more memory-efficient way
         stdout_read.seek(0)
-        stdout = [line.decode(errors="replace") for line in stdout_read.readlines()]
+        stdout = []
+        try:
+            while True:
+                chunk = stdout_read.read(buffer_size)
+                if not chunk:
+                    break
+                try:
+                    decoded_lines = chunk.decode(errors="replace").splitlines(keepends=True)
+                    stdout.extend(decoded_lines)
+                except UnicodeDecodeError:
+                    # Handle decoding issues
+                    decoded_lines = chunk.decode(errors="ignore").splitlines(keepends=True)
+                    stdout.extend(decoded_lines)
+        except OSError as e:
+            if e.errno == 12:  # Cannot allocate memory
+                logger.warning("Memory allocation issue while reading final output, truncating")
+                stdout.append("...[Output truncated due to memory constraints]...\n")
+            else:
+                raise
 
         # ignoring mypy error below because it thinks exit_code can sometimes be None
         # we know that will never be the case because the above While loop will keep looping forever

terrawrap/version.py

@@ -1,4 +1,4 @@
 """Place of record for the package version"""
 
-__version__ = "0.10.8"
+__version__ = "0.10.9"
 __git_hash__ = "GIT_HASH"

test/unit/test_cli.py

@@ -2,10 +2,10 @@
 import os
 from logging import Logger
 from unittest import TestCase
-from unittest.mock import patch, ANY, call
+from unittest.mock import patch, ANY, call, mock_open, MagicMock
 from requests.exceptions import HTTPError
 
-from terrawrap.utils.cli import execute_command, MAX_RETRIES, Status, _post_audit_info
+from terrawrap.utils.cli import execute_command, MAX_RETRIES, Status, _post_audit_info, _execute_command
 
 
 MOCK_ERROR = HTTPError()
@@ -124,3 +124,95 @@ def test_post_audit_info_statuses(self, mock_post, _):
                 },
                 timeout=30,
             )
+
+    @patch("tempfile.mkstemp")
+    @patch.object(Logger, "warning")
+    def test_execute_command_memory_allocation_error(self, mock_logger, mock_mkstemp):
+        """Test handling of OSError errno 12 (Cannot allocate memory) during command execution"""
+        # Setup mock file objects
+        mock_stdout_fd = 3
+        mock_stdout_path = "/tmp/mock_stdout"
+        mock_mkstemp.return_value = (mock_stdout_fd, mock_stdout_path)
+
+        # Create mock file object that raises OSError errno 12 on first read, then succeeds
+        mock_file = MagicMock()
+        memory_error = OSError()
+        memory_error.errno = 12  # Cannot allocate memory
+
+        # Configure read to fail first, then succeed
+        mock_file.read.side_effect = [
+            memory_error,  # First read fails with memory error
+            b"test output",  # Second read succeeds with reduced buffer
+            b"",  # Third read returns empty (process finished)
+            b"",  # Additional reads for final output collection
+            b"",  # More reads to handle any additional calls
+        ] + [b""] * 10  # Ensure we have enough empty responses
+
+        # Mock process
+        mock_process = MagicMock()
+        mock_process.poll.return_value = 0  # Process finished successfully
+
+        with patch("builtins.open", mock_open()) as mock_file_open:
+            mock_file_open.return_value.__enter__.return_value = mock_file
+
+            with patch("subprocess.Popen", return_value=mock_process):
+                # Test that the function handles memory error gracefully
+                exit_code, stdout = _execute_command(
+                    ["test", "command"],
+                    print_output=False,
+                    capture_stderr=True,
+                    print_command=False
+                )
+
+                # Verify the function completed successfully
+                self.assertEqual(exit_code, 0)
+
+                # Verify that warning was logged about memory allocation issue
+                mock_logger.assert_called_with(
+                    "Memory allocation issue while reading output, reducing buffer size"
+                )
+
+                # Verify that read was called multiple times (initial failure, then retry)
+                self.assertTrue(mock_file.read.call_count >= 2)
+
+    @patch("tempfile.mkstemp")
+    @patch.object(Logger, "warning")
+    def test_execute_command_memory_allocation_error_final_read(self, mock_logger, mock_mkstemp):
+        """Test handling of OSError errno 12 during final output reading"""
+        # Setup mock file objects
+        mock_stdout_fd = 3
+        mock_stdout_path = "/tmp/mock_stdout"
+        mock_mkstemp.return_value = (mock_stdout_fd, mock_stdout_path)
+
+        # Create mock file object that works for live reading but fails on final read
+        mock_file = MagicMock()
+        memory_error = OSError()
+        memory_error.errno = 12  # Cannot allocate memory
+
+        # Setup side effects: normal read during live output, then memory error on seek+read
+        mock_file.read.side_effect = [b"", memory_error]  # Empty for live, error for final
+        mock_process = MagicMock()
+        mock_process.poll.return_value = 0
+
+        with patch("builtins.open", mock_open()) as mock_file_open:
+            mock_file_open.return_value.__enter__.return_value = mock_file
+
+            with patch("subprocess.Popen", return_value=mock_process):
+                # Test that the function handles memory error gracefully during final read
+                exit_code, stdout = _execute_command(
+                    ["test", "command"],
+                    print_output=False,
+                    capture_stderr=True,
+                    print_command=False
+                )
+
+                # Verify the function completed successfully
+                self.assertEqual(exit_code, 0)
+
+                # Verify that warning was logged about memory allocation issue during final read
+                mock_logger.assert_called_with(
+                    "Memory allocation issue while reading final output, truncating"
+                )
+
+                # Verify that stdout contains the truncation message
+                self.assertIn("...[Output truncated due to memory constraints]...\n", stdout)