"""Tests for src/connectors/duckdb/utils.py - DuckDB table operations.""" from pathlib import Path from unittest.mock import MagicMock, patch import pytest from src.connectors.duckdb.utils import ( add_snowflake_secret, create_table_from_file, get_file_columns, normalize_table_rows, ) from src.connectors.snowflake import SnowflakeConfig from src.enums import FileType from src.errors import FileParsingError from src.schemas.models import FileMetadata class TestTableNormalization: """Test table normalization operations.""" def test_normalize_table_rows_control_chars(self): """Test removal of control characters from string columns.""" mock_conn = MagicMock() with patch( 'src.connectors.duckdb.utils.get_table_string_columns', return_value=['name', 'description'], ): normalize_table_rows(mock_conn, 'test_table') mock_conn.execute.assert_called_once() sql = mock_conn.execute.call_args[0][0] # Should update the table assert 'UPDATE "test_table"' in sql # Should remove control characters using REGEXP_REPLACE assert 'REGEXP_REPLACE' in sql assert r'[\x00-\x08\x0B-\x0C\x0E-\x1F\x7F]+' in sql # Should trim whitespace assert 'TRIM' in sql # Should update both columns assert '"name"' in sql assert '"description"' in sql def test_normalize_table_rows_null_preservation(self): """Test that NULL values are preserved and not converted to empty strings.""" mock_conn = MagicMock() with patch( 'src.connectors.duckdb.utils.get_table_string_columns', return_value=['status'], ): normalize_table_rows(mock_conn, 'test_table') sql = mock_conn.execute.call_args[0][0] # Should use null-safe functions (TRIM, REGEXP_REPLACE) # Explicit IF check was removed as functions are null-safe assert 'TRIM(REGEXP_REPLACE("status"' in sql def test_normalize_table_rows_max_length(self): """Test enforcement of maximum length limits.""" mock_conn = MagicMock() with patch( 'src.connectors.duckdb.utils.get_table_string_columns', return_value=['title'], ): normalize_table_rows(mock_conn, 'test_table', max_len=255) sql = mock_conn.execute.call_args[0][0] # Should use LEFT to truncate assert 'LEFT(' in sql assert ', 255)' in sql def test_normalize_table_rows_whitespace(self): """Test trimming of leading/trailing whitespace.""" mock_conn = MagicMock() with patch( 'src.connectors.duckdb.utils.get_table_string_columns', return_value=['address'], ): normalize_table_rows(mock_conn, 'test_table') sql = mock_conn.execute.call_args[0][0] # Should trim whitespace assert 'TRIM' in sql # Should preserve legitimate whitespace (TAB, LF, CR) by not removing them # The regex removes 0x00-0x08, 0x0B-0x0C, 0x0E-0x1F, 0x7F # But keeps 0x09 (TAB), 0x0A (LF), 0x0D (CR) assert r'[\x00-\x08\x0B-\x0C\x0E-\x1F\x7F]+' in sql def test_normalize_table_rows_multiple_columns(self): """Test normalization handles multiple columns.""" mock_conn = MagicMock() # Test with 5 columns columns = ['col1', 'col2', 'col3', 'col4', 'col5'] with patch( 'src.connectors.duckdb.utils.get_table_string_columns', return_value=columns ): normalize_table_rows(mock_conn, 'test_table') sql = mock_conn.execute.call_args[0][0] # Should update all columns for col in columns: assert f'"{col}"' in sql # Should use SET with comma-separated updates assert 'SET' in sql # Count number of column updates (each column appears 2 times: once in SET, once in REGEXP_REPLACE) assert sql.count('"col1"') == 2 def test_normalize_table_rows_no_string_columns(self): """Test that normalization is skipped if there are no string columns.""" mock_conn = MagicMock() with patch( 'src.connectors.duckdb.utils.get_table_string_columns', return_value=[] ): normalize_table_rows(mock_conn, 'test_table') # Should not execute any SQL if there are no string columns mock_conn.execute.assert_not_called() def test_normalize_table_rows_negative_max_length(self): """Test that negative max_len raises ValueError.""" mock_conn = MagicMock() with ( patch( 'src.connectors.duckdb.utils.get_table_string_columns', return_value=['col1'], ), pytest.raises(ValueError, match='max_len must be nonnegative'), ): normalize_table_rows(mock_conn, 'test_table', max_len=-1) class TestAddSnowflakeSecret: """Test add_snowflake_secret function.""" def test_add_snowflake_secret_with_private_key(self): """Test creating secret with private key string.""" mock_cursor = MagicMock() config = SnowflakeConfig( account='test_account', user='test_user', role='test_role', database='test_db', warehouse='test_wh', schema='test_schema', private_key='test_key_string', ) result = add_snowflake_secret(mock_cursor, config, 'test_secret', replace=True) assert result == 'test_secret' mock_cursor.execute.assert_called_once() sql = mock_cursor.execute.call_args[0][0] assert 'CREATE OR REPLACE SECRET test_secret' in sql assert 'TYPE snowflake' in sql assert 'AUTH_TYPE key_pair' in sql assert "ACCOUNT 'test_account'" in sql assert "USER 'test_user'" in sql assert "ROLE 'test_role'" in sql assert "DATABASE 'test_db'" in sql assert "WAREHOUSE 'test_wh'" in sql assert "PRIVATE_KEY 'test_key_string'" in sql def test_add_snowflake_secret_with_private_key_path(self): """Test creating secret with private key path.""" mock_cursor = MagicMock() config = SnowflakeConfig( account='acct', user='user', role='role', database='db', warehouse='wh', schema='test_schema', private_key_path='/path/to/key.pem', ) result = add_snowflake_secret(mock_cursor, config, replace=False) assert result == 'my_snowflake_secret' # default name sql = mock_cursor.execute.call_args[0][0] assert 'CREATE SECRET' in sql assert 'CREATE OR REPLACE' not in sql assert "PRIVATE_KEY_PATH '/path/to/key.pem'" in sql def test_add_snowflake_secret_with_passphrase_string(self): """Test creating secret with passphrase as string.""" mock_cursor = MagicMock() config = SnowflakeConfig( account='acct', user='user', role='role', database='db', warehouse='wh', schema='test_schema', private_key='key', private_key_passphrase='my_passphrase', ) add_snowflake_secret(mock_cursor, config) sql = mock_cursor.execute.call_args[0][0] assert "PRIVATE_KEY_PASSPHRASE 'my_passphrase'" in sql class TestGetFileColumns: """Test get_file_columns function.""" def test_get_file_columns_csv(self): """Test getting columns from CSV file.""" mock_cursor = MagicMock() mock_result = MagicMock() mock_result.fetchall.return_value = [ ('id', 'INTEGER'), ('name', 'VARCHAR'), ('email', 'VARCHAR'), ] mock_cursor.execute.return_value = mock_result meta = FileMetadata( file_path=Path('/tmp/test.csv'), file_type=FileType.CSV, encoding='utf-8', gzipped=False, ) result = get_file_columns(mock_cursor, meta) assert result == {'id', 'name', 'email'} mock_cursor.execute.assert_called_once() call_args = mock_cursor.execute.call_args assert 'DESCRIBE SELECT * FROM' in call_args[0][0] def test_get_file_columns_parquet(self): """Test getting columns from Parquet file.""" mock_cursor = MagicMock() mock_result = MagicMock() mock_result.fetchall.return_value = [('col1', 'BIGINT'), ('col2', 'DOUBLE')] mock_cursor.execute.return_value = mock_result meta = FileMetadata( file_path=Path('/data/file.parquet'), file_type=FileType.PQT, encoding='utf-8', gzipped=False, ) result = get_file_columns(mock_cursor, meta) assert result == {'col1', 'col2'} def test_get_file_columns_error_handling(self): """Test error handling when file parsing fails.""" mock_cursor = MagicMock() mock_cursor.execute.side_effect = Exception('Parse error') meta = FileMetadata( file_path=Path('/bad/file.csv'), file_type=FileType.CSV, encoding='utf-8', gzipped=False, ) with pytest.raises(FileParsingError, match='Failed to parse csv columns'): get_file_columns(mock_cursor, meta) class TestCreateTableFromFile: """Test create_table_from_file function.""" def test_create_table_from_file_csv_no_alias(self): """Test creating table from CSV without column aliasing.""" mock_cursor = MagicMock() meta = FileMetadata( file_path=Path('/tmp/data.csv'), file_type=FileType.CSV, encoding='utf-8', gzipped=False, ) with patch('src.connectors.duckdb.utils.get_select_clause', return_value='*'): create_table_from_file(mock_cursor, 'my_table', meta) mock_cursor.execute.assert_called_once() sql = mock_cursor.execute.call_args[0][0] assert 'CREATE OR REPLACE TABLE "my_table"' in sql assert 'SELECT *' in sql def test_create_table_from_file_with_alias_map(self): """Test creating table with column aliasing.""" mock_cursor = MagicMock() meta = FileMetadata( file_path=Path('/data/input.csv'), file_type=FileType.CSV, encoding='utf-8', gzipped=False, ) alias_map = {'new_id': 'id', 'new_name': 'name'} with patch( 'src.connectors.duckdb.utils.get_select_clause', return_value='"id" AS "new_id", "name" AS "new_name"', ): create_table_from_file(mock_cursor, 'aliased_table', meta, alias_map) sql = mock_cursor.execute.call_args[0][0] assert 'CREATE OR REPLACE TABLE "aliased_table"' in sql assert '"id" AS "new_id", "name" AS "new_name"' in sql def test_create_table_from_file_parquet(self): """Test creating table from Parquet file.""" mock_cursor = MagicMock() meta = FileMetadata( file_path=Path('/storage/data.parquet'), file_type=FileType.PQT, encoding='utf-8', gzipped=False, ) with patch('src.connectors.duckdb.utils.get_select_clause', return_value='*'): create_table_from_file(mock_cursor, 'parquet_table', meta) mock_cursor.execute.assert_called_once() sql = mock_cursor.execute.call_args[0][0] assert 'CREATE OR REPLACE TABLE "parquet_table"' in sql def test_create_table_from_file_error_handling(self): """Test error handling when table creation fails.""" mock_cursor = MagicMock() mock_cursor.execute.side_effect = Exception('Creation failed') meta = FileMetadata( file_path=Path('/bad/file.csv'), file_type=FileType.CSV, encoding='utf-8', gzipped=False, ) with ( patch('src.connectors.duckdb.utils.get_select_clause', return_value='*'), pytest.raises(FileParsingError, match='Failed to load csv file'), ): create_table_from_file(mock_cursor, 'failed_table', meta)