import os import shutil import tempfile import unittest from unittest.mock import patch, MagicMock, ANY from datetime import datetime, timezone import pyarrow as pa import pandas as pd try: # Try to import from the installed deltalake package import sys # Remove local deltalake from path temporarily if 'deltalake' in sys.modules: del sys.modules['deltalake'] from deltalake import DeltaTable, write_deltalake DELTALAKE_AVAILABLE = True except ImportError: DELTALAKE_AVAILABLE = False DeltaTable = None write_deltalake = None from o11ydiagnosis.subtasks.slowlogconverter.slowlog import SlowLog, parse from o11ydiagnosis.subtasks.slowlogconverter.write_deltalake import ( write_slowlogs_to_deltalake, write_slowlogs_to_deltalake_multitables, slowlogs_schema, plans_schema, queries_schema, ) class TestWriteDeltaLakeRealLogs(unittest.TestCase): @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") def setUp(self): # Create a temporary directory for test Delta tables self.temp_dir = tempfile.mkdtemp() self.base_path = os.path.join(self.temp_dir, "deltalake") # Read the real sample log file self.log_file_path = os.path.join( os.path.dirname(__file__), 'data', 'rawsampleslowlog.log' ) # Parse the log file to create SlowLog objects # Skip if file doesn't exist if os.path.exists(self.log_file_path): self.sample_logs = self._parse_log_file() else: # Create a minimal test log if file doesn't exist self.sample_logs = self._create_minimal_test_logs() def tearDown(self): # Clean up the temporary directory shutil.rmtree(self.temp_dir, ignore_errors=True) def _parse_log_file(self): """Parse the sample log file and return a list of SlowLog objects.""" logs = [] current_log_lines = [] with open(self.log_file_path, 'r') as f: for line in f: line = line.strip() if not line: continue # Check if this is the start of a new log entry if (line[0] == '2' and "# Time: " in line) or line.startswith("# Time: "): if current_log_lines: # Parse the completed log entry log = parse(current_log_lines) if log: logs.append(log) current_log_lines = [] current_log_lines.append(line) # Don't forget the last log entry if current_log_lines: log = parse(current_log_lines) if log: logs.append(log) return logs def _create_minimal_test_logs(self): """Create minimal test logs if sample file doesn't exist""" log = SlowLog() log.id = 1 log.time = 1625097600 log.digest = "test_digest" log.instance = "test-instance" log.query = "SELECT 1" log.db = "test_db" log.query_time = 1.0 log.binary_plan = "test_binary_plan" return [log] @patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.write_to_table') @patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.logger') @patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.parse_binary_plan') def test_write_slowlogs_to_deltalake_with_real_logs(self, mock_parse_binary_plan, mock_logger, mock_write_to_table): """Test writing real log data to Delta Lake.""" # Skip if no logs were parsed if not self.sample_logs: self.skipTest("No logs were parsed from the sample log file") # Ensure binary_plan is set for testing for log in self.sample_logs: if not hasattr(log, 'binary_plan') or not log.binary_plan: log.binary_plan = "test_binary_plan_data" # Configure the mock for parse_binary_plan mock_parse_binary_plan.return_value = "{\"test_plan\": true}" # Call the function with the parsed logs region = "us-west-2" cluster = "test-cluster" write_slowlogs_to_deltalake( added_logs=self.sample_logs, region=region, cluster=cluster, base_path=self.base_path, optimize=False # Skip optimization for faster tests ) # Verify binary_plan was processed for each log that has it binary_plan_logs = [log for log in self.sample_logs if hasattr(log, 'binary_plan') and log.binary_plan] if binary_plan_logs: # We can verify the binary_plan processing by checking the call to write_to_table pass # Verify write_to_table was called with the expected parameters self.assertTrue(mock_write_to_table.called) # Get the call arguments for write_to_table call_args = mock_write_to_table.call_args_list[0][1] # Verify the call was made with the expected parameters self.assertEqual(call_args['region'], region) self.assertTrue(call_args['s3_path'].startswith(self.base_path)) self.assertIn('slowlogs', call_args['s3_path']) # Verify the source data contains the expected fields source_table = call_args['source_pyarrow_table'] self.assertIsInstance(source_table, pa.Table) # Convert to pandas for easier assertions df = source_table.to_pandas() # Verify the DataFrame has the expected columns expected_columns = [field.name for field in slowlogs_schema] for col in expected_columns: self.assertIn(col, df.columns, f"Expected column {col} not found in DataFrame") # Verify the data was properly converted self.assertEqual(len(df), len(self.sample_logs)) # Verify some specific fields from the first log first_log = self.sample_logs[0] first_row = df.iloc[0] # Check some fields that should be in the log if hasattr(first_log, 'query_time'): self.assertAlmostEqual(first_row['query_time'], first_log.query_time) if hasattr(first_log, 'db'): self.assertEqual(first_row['db'], first_log.db) if hasattr(first_log, 'digest'): self.assertEqual(first_row['digest'], first_log.digest) @patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.write_to_table') @patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.logger') @patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.parse_binary_plan') def test_write_slowlogs_to_deltalake_multitables_with_real_logs(self, mock_parse_binary_plan, mock_logger, mock_write_to_table): """Test writing real log data to multiple Delta Lake tables.""" # Skip if no logs were parsed if not self.sample_logs: self.skipTest("No logs were parsed from the sample log file") # Ensure binary_plan is set for testing for log in self.sample_logs: if not hasattr(log, 'binary_plan') or not log.binary_plan: log.binary_plan = "test_binary_plan_data" # Configure the mock for parse_binary_plan mock_parse_binary_plan.return_value = "{\"test_plan\": true}" # Call the function with the parsed logs region = "us-west-2" cluster = "test-cluster" write_slowlogs_to_deltalake_multitables( added_logs=self.sample_logs, region=region, cluster=cluster, base_path=self.base_path, optimize=False # Skip optimization for faster tests ) # Verify binary_plan was processed for each log that has it binary_plan_logs = [log for log in self.sample_logs if hasattr(log, 'binary_plan') and log.binary_plan] if binary_plan_logs: # We can verify the binary_plan processing by checking the call to write_to_table pass # Verify write_to_table was called for each table type (slowlogs, plans, queries) self.assertGreaterEqual(len(mock_write_to_table.call_args_list), 1) # Get all the s3_paths that were written to written_paths = [call[1]['s3_path'] for call in mock_write_to_table.call_args_list] # Verify all expected tables were written to expected_tables = ['slowlogs', 'plans', 'queries'] for table in expected_tables: self.assertTrue(any(table in path for path in written_paths), f"Expected table '{table}' was not written to") # Verify the source data for the slowlogs table slowlogs_call = None for call in mock_write_to_table.call_args_list: if 'slowlogs' in call[1]['s3_path']: slowlogs_call = call[1] break self.assertIsNotNone(slowlogs_call, "No call to write to slowlogs table found") # Verify the source data contains the expected fields source_table = slowlogs_call['source_pyarrow_table'] self.assertIsInstance(source_table, pa.Table) # Convert to pandas for easier assertions df = source_table.to_pandas() # Verify the DataFrame has the expected columns expected_columns = [field.name for field in slowlogs_schema] for col in expected_columns: self.assertIn(col, df.columns, f"Expected column {col} not found in DataFrame") # Verify the data was properly converted self.assertEqual(len(df), len(self.sample_logs)) if __name__ == '__main__': unittest.main()