diff --git a/data-transfer/pontoon/pontoon/__init__.py b/data-transfer/pontoon/pontoon/__init__.py index 5713b9a..7fe89c6 100644 --- a/data-transfer/pontoon/pontoon/__init__.py +++ b/data-transfer/pontoon/pontoon/__init__.py @@ -1,5 +1,9 @@ from pontoon.logging_config import configure_logging, logger from pontoon.source.sql_source import SQLSource +from pontoon.source.postgresql_source import PostgreSQLSource +from pontoon.source.redshift_source import RedshiftSource +from pontoon.source.bigquery_source import BigQuerySource +from pontoon.source.snowflake_source import SnowflakeSource from pontoon.source.memory_source import MemorySource from pontoon.destination.glue_destination import GlueDestination from pontoon.destination.stdout_destination import StdoutDestination @@ -53,6 +57,10 @@ def get_destination_by_vendor(vendor_type:str) -> str: # register source and destination connectors __sources['source-sql'] = SQLSource +__sources['source-postgresql'] = PostgreSQLSource +__sources['source-redshift'] = RedshiftSource +__sources['source-bigquery'] = BigQuerySource +__sources['source-snowflake'] = SnowflakeSource __sources['source-memory'] = MemorySource __destinations['destination-sql'] = SQLDestination __destinations['destination-glue'] = GlueDestination @@ -75,10 +83,10 @@ def get_destination_by_vendor(vendor_type:str) -> str: # map from vendor types to source and destination connectors __vendor_source_map = { 'memory': 'source-memory', - 'redshift': 'source-sql', - 'snowflake': 'source-sql', - 'bigquery': 'source-sql', - 'postgresql': 'source-sql', + 'redshift': 'source-redshift', + 'snowflake': 'source-snowflake', + 'bigquery': 'source-bigquery', + 'postgresql': 'source-postgresql', 'mysql': 'source-sql', 'athena': 'source-sql' } diff --git a/data-transfer/pontoon/pontoon/source/bigquery_source.py b/data-transfer/pontoon/pontoon/source/bigquery_source.py new file mode 100644 index 0000000..65a2d09 --- /dev/null +++ b/data-transfer/pontoon/pontoon/source/bigquery_source.py @@ -0,0 +1,93 @@ +import json +from typing import List +from sqlalchemy import create_engine, inspect +from sqlalchemy.engine import Engine +from pontoon.base import Namespace +from pontoon.source.sql_source import SQLSource + + +class BigQuerySource(SQLSource): + """BigQuery-specific implementation of SQLSource""" + + def _create_engine(self, connect_config: dict) -> Engine: + """Create BigQuery-specific SQLAlchemy engine with service account authentication""" + project_id = connect_config.get('project_id') + service_account = connect_config.get('service_account') + + if not project_id: + raise ValueError("BigQuery connection config must include 'project_id' field") + + if not service_account: + raise ValueError("BigQuery connection config must include 'service_account' field") + + # Parse service account JSON + try: + credentials_info = json.loads(service_account) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid service account JSON: {e}") + + # Build BigQuery connection string and create engine + connection_string = f"bigquery://{project_id}" + + # Get chunk size from connect config, defaulting to 1024 if not specified + chunk_size = connect_config.get('chunk_size', 1024) + + return create_engine( + connection_string, + credentials_info=credentials_info, + arraysize=chunk_size + ) + + def _validate_auth_type(self, auth_type: str) -> None: + """Validate authentication type for BigQuery - only 'service_account' is supported""" + if auth_type != 'service_account': + raise ValueError(f"BigQuery source only supports 'service_account' authentication, got '{auth_type}'") + + def _get_namespace(self, connect_config: dict) -> Namespace: + """Extract namespace from BigQuery connection config using project_id""" + project_id = connect_config.get('project_id') + if not project_id: + raise ValueError("BigQuery connection config must include 'project_id' field") + + return Namespace(project_id) + + def _inspect_streams_impl(self) -> List[dict]: + """BigQuery-specific stream inspection logic""" + streams = [] + + with self._connect() as conn: + # Use the inspector to get schema and table information + inspector = inspect(conn) + + # Get all available schemas (datasets in BigQuery terminology) + schemas = inspector.get_schema_names() + + for schema in schemas: + # For each table in the schema + for table in inspector.get_table_names(schema=schema): + # BigQuery table names come in format "project.dataset.table" + # We need to extract just the table name + if '.' in table: + _, table_name = table.split('.', 1) + if '.' in table_name: + # Handle case where table is "dataset.table" + _, table_name = table_name.split('.', 1) + else: + table_name = table + + # Get column information using the full table reference + project_id = self._config['connect']['project_id'] + full_table_name = f"{project_id}.{schema}.{table_name}" + + try: + columns = inspector.get_columns(full_table_name) + streams.append({ + 'schema_name': schema, + 'stream_name': table_name, + 'fields': [{'name': col['name'], 'type': str(col['type'])} for col in columns] + }) + except Exception: + # Skip tables that can't be inspected (e.g., views, external tables) + continue + + return streams \ No newline at end of file diff --git a/data-transfer/pontoon/pontoon/source/postgresql_source.py b/data-transfer/pontoon/pontoon/source/postgresql_source.py new file mode 100644 index 0000000..ee64f74 --- /dev/null +++ b/data-transfer/pontoon/pontoon/source/postgresql_source.py @@ -0,0 +1,36 @@ +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from pontoon.base import Namespace +from pontoon.source.sql_source import SQLSource + + +class PostgreSQLSource(SQLSource): + """PostgreSQL-specific implementation of SQLSource""" + + def _create_engine(self, connect_config: dict) -> Engine: + """Create PostgreSQL-specific SQLAlchemy engine""" + host = connect_config.get('host') + port = connect_config.get('port', '5432') + user = connect_config.get('user') + password = connect_config.get('password') + database = connect_config.get('database') + + # Build PostgreSQL connection string using psycopg2 driver + connection_string = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}" + + return create_engine(connection_string) + + def _validate_auth_type(self, auth_type: str) -> None: + """Validate authentication type for PostgreSQL - only 'basic' is supported""" + if auth_type != 'basic': + raise ValueError(f"PostgreSQL source only supports 'basic' authentication, got '{auth_type}'") + + def _get_namespace(self, connect_config: dict) -> Namespace: + """Extract namespace from PostgreSQL connection config """ + + # Check for separate database field + database = connect_config.get('database') + if not database: + raise ValueError("PostgreSQL connection config must include 'database' field") + + return Namespace(database) \ No newline at end of file diff --git a/data-transfer/pontoon/pontoon/source/redshift_source.py b/data-transfer/pontoon/pontoon/source/redshift_source.py new file mode 100644 index 0000000..11a474d --- /dev/null +++ b/data-transfer/pontoon/pontoon/source/redshift_source.py @@ -0,0 +1,19 @@ +from pontoon.source.postgresql_source import PostgreSQLSource + + +class RedshiftSource(PostgreSQLSource): + """Redshift-specific implementation of SQLSource + + Inherits from PostgreSQLSource since Redshift uses the PostgreSQL protocol + and wire format. The connection string format, authentication method (basic), + and namespace extraction logic are identical to PostgreSQL. + + This class overrides the authentication validation to provide Redshift-specific + error messages and can be extended in the future if Redshift-specific behaviors + are needed (e.g., specific query optimizations, column type handling, etc.). + """ + + def _validate_auth_type(self, auth_type: str) -> None: + """Validate authentication type for Redshift - only 'basic' is supported""" + if auth_type != 'basic': + raise ValueError(f"Redshift source only supports 'basic' authentication, got '{auth_type}'") \ No newline at end of file diff --git a/data-transfer/pontoon/pontoon/source/snowflake_source.py b/data-transfer/pontoon/pontoon/source/snowflake_source.py new file mode 100644 index 0000000..ab58ff6 --- /dev/null +++ b/data-transfer/pontoon/pontoon/source/snowflake_source.py @@ -0,0 +1,49 @@ +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from pontoon.base import Namespace +from pontoon.source.sql_source import SQLSource + + +class SnowflakeSource(SQLSource): + """Snowflake-specific implementation of SQLSource""" + + def _create_engine(self, connect_config: dict) -> Engine: + """Create Snowflake-specific SQLAlchemy engine""" + user = connect_config.get('user') + access_token = connect_config.get('access_token') + account = connect_config.get('account') + database = connect_config.get('database') + warehouse = connect_config.get('warehouse') + + # Validate required fields + if not user: + raise ValueError("Snowflake connection config must include 'user' field") + if not access_token: + raise ValueError("Snowflake connection config must include 'access_token' field") + if not account: + raise ValueError("Snowflake connection config must include 'account' field") + if not database: + raise ValueError("Snowflake connection config must include 'database' field") + if not warehouse: + raise ValueError("Snowflake connection config must include 'warehouse' field") + + # Build Snowflake connection string + # Format: snowflake://user:access_token@account/database/schema?warehouse=warehouse + # Note: We'll use 'public' as default schema since it's commonly available + schema = connect_config.get('schema', 'public') + connection_string = f"snowflake://{user}:{access_token}@{account}/{database}/{schema}?warehouse={warehouse}" + + return create_engine(connection_string) + + def _validate_auth_type(self, auth_type: str) -> None: + """Validate authentication type for Snowflake - only 'access_token' is supported""" + if auth_type != 'access_token': + raise ValueError(f"Snowflake source only supports 'access_token' authentication, got '{auth_type}'") + + def _get_namespace(self, connect_config: dict) -> Namespace: + """Extract namespace from Snowflake connection config using database field""" + database = connect_config.get('database') + if not database: + raise ValueError("Snowflake connection config must include 'database' field") + + return Namespace(database) \ No newline at end of file diff --git a/data-transfer/pontoon/pontoon/source/sql_source.py b/data-transfer/pontoon/pontoon/source/sql_source.py index d6bc6e9..c6c853a 100644 --- a/data-transfer/pontoon/pontoon/source/sql_source.py +++ b/data-transfer/pontoon/pontoon/source/sql_source.py @@ -1,9 +1,11 @@ import json import re +from abc import ABC, abstractmethod from typing import List, Dict, Tuple, Generator, Any from datetime import datetime, timezone, date from decimal import Decimal from sqlalchemy import create_engine, inspect, MetaData, Table, text, select, func +from sqlalchemy.engine import Engine from sqlalchemy.exc import SQLAlchemyError, OperationalError, InterfaceError, DatabaseError, NoSuchTableError from pontoon import logger @@ -96,9 +98,8 @@ def build_select_query(stream:Stream, mode:Mode, count:bool=False) -> str: return select_query -class SQLSource(Source): - """ A Source implementation that can read from any database supported by SQLAlchemy """ - +class SQLSource(Source, ABC): + """ Abstract base class for SQL source implementations """ def __init__(self, config, cache_implementation, cache_config={}): self._config = config @@ -106,12 +107,8 @@ def __init__(self, config, cache_implementation, cache_config={}): connect = config.get('connect') - if connect.get('database'): - self._namespace = Namespace(connect.get('database')) - elif connect.get('project_id'): - self._namespace = Namespace(connect.get('project_id')) - else: - self._namespace = Namespace('unknown') + # Extract namespace using database-specific logic + self._namespace = self._get_namespace(connect) # our sync mode config self._mode = config.get('mode') @@ -129,48 +126,40 @@ def __init__(self, config, cache_implementation, cache_config={}): # additional fields to augment streams with self._with = config.get('with', {}) - # configure the SQLAlchemy engine + # Validate authentication type for this database auth_type = connect.get('auth_type') + self._validate_auth_type(auth_type) - if connect.get('dsn'): - self._engine = create_engine(connect.get('dsn')) - else: - vendor_type = connect['vendor_type'] - if vendor_type in ['redshift', 'postgresql']: - if auth_type != 'basic': - raise Exception(f'SQLSource (source-sql) does not support auth_type {auth_type} for {vendor_type}') - self._engine = create_engine( - f"postgresql+psycopg2://{connect['user']}:{connect['password']}@"\ - f"{connect['host']}:{connect['port']}/{connect['database']}" - ) - elif vendor_type == 'bigquery': - if auth_type != 'service_account': - raise Exception(f'SQLSource (source-sql) does not support auth_type {auth_type} for {vendor_type}') - self._engine = create_engine( - f"bigquery://{connect['project_id']}", - credentials_info=json.loads(connect['service_account']), - arraysize=self._chunk_size - ) - elif vendor_type == 'snowflake': - if auth_type != 'access_token': - raise Exception(f'SQLSource (source-sql) does not support auth_type {auth_type} for {vendor_type}') - self._engine = create_engine( - f"snowflake://{connect['user']}:{connect['access_token']}@"\ - f"{connect['account']}/{connect['database']}?warehouse={connect['warehouse']}" - ) - else: - raise Exception(f'SQLSource (source-sql) does not support vendor_type: {vendor_type}') + # Create database-specific engine + self._engine = self._create_engine(connect) + + @abstractmethod + def _create_engine(self, connect_config: dict) -> Engine: + """Create database-specific SQLAlchemy engine""" + pass + + @abstractmethod + def _validate_auth_type(self, auth_type: str) -> None: + """Validate authentication type for this database""" + pass + + @abstractmethod + def _get_namespace(self, connect_config: dict) -> Namespace: + """Extract namespace from connection config""" + pass + + def _inspect_streams_impl(self) -> List[dict]: + """Database-specific stream inspection - default implementation""" + return self.inspect_standard_streams() - def _connect(self): try: return self._engine.connect() except (OperationalError, InterfaceError, DatabaseError) as e: raise SourceConnectionFailed("Failed to connect to source database") from e - def test_connect(self): - # connect test + """Test database connection using template method pattern""" with self._connect() as conn: return True @@ -205,51 +194,14 @@ def inspect_standard_streams(self): return streams - def inspect_bigquery_streams(self): - - streams = [] - - with self._connect() as conn: - # Use the inspector to get schema and table information - inspector = inspect(conn) - - # Get all available schemas - schemas = inspector.get_schema_names() - - for schema in schemas: - - # For each table in the schema - for table in inspector.get_table_names(schema=schema): - _, table_name = table.split('.') - - columns = inspector.get_columns(f"{self._config['connect']['project_id']}.{schema}.{table_name}") - streams.append({ - 'schema_name': schema, - 'stream_name': table_name, - 'fields': [{'name': col['name'], 'type': str(col['type'])} for col in columns] - }) - - return streams - - - def inspect_streams(self): - # schema info reflection - vendor_type = self._config['connect']['vendor_type'] - - if vendor_type in ['postgresql', 'redshift', 'snowflake']: - return self.inspect_standard_streams() - - if vendor_type == 'bigquery': - return self.inspect_bigquery_streams() - - return [] + """Inspect available streams using database-specific implementation""" + return self._inspect_streams_impl() - def read(self, progress_callback=None) -> Dataset: - # Read from source and write to a cached Dataset + """Read from source and write to a cached Dataset using template method pattern""" with self._connect() as conn: @@ -359,4 +311,5 @@ def read(self, progress_callback=None) -> Dataset: def close(self): + """Close the cache - common cleanup logic""" self._cache.close() \ No newline at end of file diff --git a/data-transfer/pontoon/tests/unit/test_postgresql_source.py b/data-transfer/pontoon/tests/unit/test_postgresql_source.py new file mode 100644 index 0000000..22170ef --- /dev/null +++ b/data-transfer/pontoon/tests/unit/test_postgresql_source.py @@ -0,0 +1,299 @@ +import pytest +from unittest.mock import Mock, patch, MagicMock +from pontoon.source.postgresql_source import PostgreSQLSource +from pontoon.base import Namespace, Mode +from datetime import datetime, timezone + + +class TestPostgreSQLSource: + """Unit tests for the refactored PostgreSQLSource implementation""" + + def test_create_engine(self): + """Test PostgreSQL engine creation with connection string generation""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'port': '5432', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine') as mock_create_engine: + mock_engine = Mock() + mock_create_engine.return_value = mock_engine + + source = PostgreSQLSource(config, Mock(), {}) + + # Verify create_engine was called with correct connection string + mock_create_engine.assert_called_with( + "postgresql+psycopg2://testuser:testpass@localhost:5432/testdb" + ) + + def test_create_engine_custom_port(self): + """Test PostgreSQL engine creation with custom port""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'port': '5433', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine') as mock_create_engine: + mock_engine = Mock() + mock_create_engine.return_value = mock_engine + + source = PostgreSQLSource(config, Mock(), {}) + + # Verify create_engine was called with custom port + mock_create_engine.assert_called_with( + "postgresql+psycopg2://testuser:testpass@localhost:5433/testdb" + ) + + def test_validate_auth_type_basic(self): + """Test that 'basic' authentication type is accepted""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + # Should not raise an exception + source = PostgreSQLSource(config, Mock(), {}) + + def test_validate_auth_type_invalid(self): + """Test that non-'basic' authentication types are rejected""" + connect_config = { + 'auth_type': 'service_account', + 'host': 'localhost', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + with pytest.raises(ValueError, match="PostgreSQL source only supports 'basic' authentication, got 'service_account'"): + PostgreSQLSource(config, Mock(), {}) + + def test_get_namespace(self): + """Test namespace extraction from 'database' field""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'mydb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + source = PostgreSQLSource(config, Mock(), {}) + + # Verify namespace is extracted from database field + assert source._namespace.name == 'mydb' + + def test_get_namespace_missing_database(self): + """Test that missing 'database' field raises error""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'user': 'testuser', + 'password': 'testpass' + # Missing 'database' field + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + with pytest.raises(ValueError, match="PostgreSQL connection config must include 'database' field"): + PostgreSQLSource(config, Mock(), {}) + + def test_inspect_streams_impl_uses_default(self): + """Test that PostgreSQL uses default stream inspection implementation""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + source = PostgreSQLSource(config, Mock(), {}) + + # Mock the inspect_standard_streams method + with patch.object(source, 'inspect_standard_streams') as mock_inspect: + mock_inspect.return_value = [{'test': 'stream'}] + + result = source._inspect_streams_impl() + + # Verify it calls the default implementation + mock_inspect.assert_called_once() + assert result == [{'test': 'stream'}] + + +class TestPostgreSQLSourceIntegration: + """Integration tests for PostgreSQLSource to verify it works with the full system""" + + def test_postgresql_source_instantiation(self): + """Test that PostgreSQLSource can be instantiated with proper config""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'port': '5432', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + mode_config = { + 'type': Mode.FULL_REFRESH + } + + config = { + 'connect': connect_config, + 'mode': Mode(mode_config), + 'streams': [{ + 'schema': 'public', + 'table': 'test_table', + 'primary_field': 'id' + }] + } + + # Mock the SQLAlchemy engine creation + with patch('pontoon.source.postgresql_source.create_engine') as mock_create_engine: + mock_engine = MagicMock() + mock_create_engine.return_value = mock_engine + + # Mock cache implementation + mock_cache = Mock() + + # Create the PostgreSQL source + source = PostgreSQLSource(config, mock_cache, {}) + + # Verify the source was created correctly + assert source._namespace.name == 'testdb' + assert source._engine == mock_engine + + # Verify the connection string was built correctly + mock_create_engine.assert_called_with( + "postgresql+psycopg2://testuser:testpass@localhost:5432/testdb" + ) + + def test_postgresql_source_test_connect(self): + """Test that test_connect method works""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'port': '5432', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine') as mock_create_engine: + # Mock engine and connection + mock_connection = MagicMock() + mock_engine = MagicMock() + mock_engine.connect.return_value.__enter__.return_value = mock_connection + mock_create_engine.return_value = mock_engine + + source = PostgreSQLSource(config, Mock(), {}) + + # Test the connection + result = source.test_connect() + + # Verify connection was attempted + assert result == True + mock_engine.connect.assert_called_once() + + def test_postgresql_source_inspect_streams(self): + """Test that inspect_streams method works and uses default implementation""" + connect_config = { + 'auth_type': 'basic', + 'host': 'localhost', + 'port': '5432', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb' + } + + config = { + 'connect': connect_config, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine') as mock_create_engine: + mock_engine = MagicMock() + mock_create_engine.return_value = mock_engine + + source = PostgreSQLSource(config, Mock(), {}) + + # Mock the inspect_standard_streams method + expected_streams = [ + { + 'schema_name': 'public', + 'stream_name': 'users', + 'fields': [ + {'name': 'id', 'type': 'INTEGER'}, + {'name': 'name', 'type': 'VARCHAR'} + ] + } + ] + + with patch.object(source, 'inspect_standard_streams', return_value=expected_streams): + result = source.inspect_streams() + + # Verify it returns the expected streams + assert result == expected_streams \ No newline at end of file diff --git a/data-transfer/pontoon/tests/unit/test_postgres_source.py b/data-transfer/pontoon/tests/unit/test_postgresql_types.py similarity index 93% rename from data-transfer/pontoon/tests/unit/test_postgres_source.py rename to data-transfer/pontoon/tests/unit/test_postgresql_types.py index 1e222d6..189e1c5 100644 --- a/data-transfer/pontoon/tests/unit/test_postgres_source.py +++ b/data-transfer/pontoon/tests/unit/test_postgresql_types.py @@ -1,13 +1,16 @@ import pytest +import uuid +import os from datetime import datetime, timedelta, timezone, date import pyarrow as pa -from pontoon import Stream, Mode, Namespace -from pontoon.source.sql_source import SQLSource, SQLUtil +from pontoon import Stream, Mode, Namespace, Dataset +from pontoon.source.sql_source import SQLUtil +from pontoon.cache.arrow_ipc_cache import ArrowIpcCache from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String, DateTime, Boolean, Float, Text, Date, Numeric, inspect, text from sqlalchemy.dialects.postgresql import UUID, JSONB, TIMESTAMP -class TestPostgresSource: +class TestPostgresTypes: """ Tests specific to PostgreSQL schema handling and type conversion """ @@ -211,19 +214,11 @@ def test_sqlite_date_handling_fixed(self): cursor_field='created_at' ) - # Test that the SQLite cache can handle date objects correctly - from pontoon.cache.sqlite_cache import SqliteCache - import tempfile - import os - - # Create a temporary SQLite cache - temp_db = tempfile.NamedTemporaryFile(delete=False, suffix='.db') - temp_db.close() try: - cache = SqliteCache( + cache = ArrowIpcCache( namespace=Namespace("test"), - config={'db': temp_db.name, 'chunk_size': 1000} + config={'cache_dir': f"./cache-{uuid.uuid4()}"} ) # Create sample data with date objects @@ -245,8 +240,7 @@ def test_sqlite_date_handling_fixed(self): finally: # Clean up - if os.path.exists(temp_db.name): - os.unlink(temp_db.name) + pass def test_postgres_boolean_type_conversion_error(self): """ @@ -283,19 +277,12 @@ def test_postgres_boolean_type_conversion_error(self): cursor_field='created_at' ) - # Test that the SQLite cache can handle boolean objects correctly - from pontoon.cache.sqlite_cache import SqliteCache - import tempfile - import os - # Create a temporary SQLite cache - temp_db = tempfile.NamedTemporaryFile(delete=False, suffix='.db') - temp_db.close() try: - cache = SqliteCache( + cache = ArrowIpcCache( namespace=Namespace("test"), - config={'db': temp_db.name, 'chunk_size': 1000} + config={'cache_dir': f"./cache-{uuid.uuid4()}"} ) # Create sample data with boolean objects @@ -344,8 +331,7 @@ def test_postgres_boolean_type_conversion_error(self): finally: # Clean up cache.close() - if os.path.exists(temp_db.name): - os.unlink(temp_db.name) + def test_sqlite_boolean_handling_fixed(self): """ @@ -382,19 +368,12 @@ def test_sqlite_boolean_handling_fixed(self): cursor_field='created_at' ) - # Test that the SQLite cache can handle boolean objects correctly - from pontoon.cache.sqlite_cache import SqliteCache - import tempfile - import os - # Create a temporary SQLite cache - temp_db = tempfile.NamedTemporaryFile(delete=False, suffix='.db') - temp_db.close() try: - cache = SqliteCache( + cache = ArrowIpcCache( namespace=Namespace("test"), - config={'db': temp_db.name, 'chunk_size': 1000} + config={'cache_dir': f"./cache-{uuid.uuid4()}"} ) # Create sample data with boolean objects @@ -422,8 +401,7 @@ def test_sqlite_boolean_handling_fixed(self): finally: # Clean up cache.close() - if os.path.exists(temp_db.name): - os.unlink(temp_db.name) + def test_schema_compatibility_fix(self): """ @@ -539,21 +517,13 @@ def test_empty_stream_handling_fixed(self): cursor_field='created_at' ) - # Test that the destination handles empty streams correctly - from pontoon.cache.sqlite_cache import SqliteCache - from pontoon.base import Dataset, Mode - import tempfile - import os - # Create a temporary SQLite cache - temp_db = tempfile.NamedTemporaryFile(delete=False, suffix='.db') - temp_db.close() cache = None try: - cache = SqliteCache( + cache = ArrowIpcCache( namespace=Namespace("test"), - config={'db': temp_db.name, 'chunk_size': 1000} + config={'cache_dir': f"./cache-{uuid.uuid4()}"} ) # Create a dataset with an empty stream (no records written) @@ -582,12 +552,7 @@ def test_empty_stream_handling_fixed(self): # Clean up if cache: cache.close() - # Don't try to delete the file if it's still in use - try: - if os.path.exists(temp_db.name): - os.unlink(temp_db.name) - except PermissionError: - pass # File is still in use, that's okay + def test_investigate_missing_records_scenarios(self): """ diff --git a/data-transfer/pontoon/tests/unit/test_redshift_source.py b/data-transfer/pontoon/tests/unit/test_redshift_source.py new file mode 100644 index 0000000..4164fb1 --- /dev/null +++ b/data-transfer/pontoon/tests/unit/test_redshift_source.py @@ -0,0 +1,127 @@ +import pytest +from unittest.mock import Mock, patch +from pontoon.source.redshift_source import RedshiftSource +from pontoon.base import Namespace + + +class TestRedshiftSource: + """Test RedshiftSource implementation""" + + def test_inheritance_from_postgresql_source(self): + """Test that RedshiftSource properly inherits from PostgreSQLSource""" + from pontoon.source.postgresql_source import PostgreSQLSource + + # Verify inheritance + assert issubclass(RedshiftSource, PostgreSQLSource) + + def test_validate_auth_type_basic_success(self): + """Test that 'basic' authentication type is accepted""" + config = { + 'connect': { + 'host': 'test-cluster.redshift.amazonaws.com', + 'port': '5439', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb', + 'auth_type': 'basic' + }, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + source = RedshiftSource(config, Mock(), {}) + # If we get here without exception, auth validation passed + assert source is not None + + def test_validate_auth_type_invalid_fails(self): + """Test that invalid authentication types are rejected""" + config = { + 'connect': { + 'host': 'test-cluster.redshift.amazonaws.com', + 'port': '5439', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb', + 'auth_type': 'service_account' + }, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + with pytest.raises(ValueError, match="Redshift source only supports 'basic' authentication"): + RedshiftSource(config, Mock(), {}) + + @patch('pontoon.source.postgresql_source.create_engine') + def test_create_engine_inherits_postgresql_behavior(self, mock_create_engine): + """Test that engine creation inherits PostgreSQL behavior""" + config = { + 'connect': { + 'host': 'test-cluster.redshift.amazonaws.com', + 'port': '5439', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb', + 'auth_type': 'basic' + }, + 'mode': Mock(), + 'streams': [] + } + + mock_engine = Mock() + mock_create_engine.return_value = mock_engine + + source = RedshiftSource(config, Mock(), {}) + + # Verify create_engine was called with PostgreSQL connection string + mock_create_engine.assert_called_once_with( + "postgresql+psycopg2://testuser:testpass@test-cluster.redshift.amazonaws.com:5439/testdb" + ) + + @patch('pontoon.source.postgresql_source.create_engine') + def test_get_namespace_inherits_postgresql_behavior(self, mock_create_engine): + """Test that namespace extraction inherits PostgreSQL behavior""" + config = { + 'connect': { + 'host': 'test-cluster.redshift.amazonaws.com', + 'port': '5439', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb', + 'auth_type': 'basic' + }, + 'mode': Mock(), + 'streams': [] + } + + mock_engine = Mock() + mock_create_engine.return_value = mock_engine + + source = RedshiftSource(config, Mock(), {}) + + # Verify namespace is extracted from database field + assert str(source._namespace) == str(Namespace('testdb')) + + def test_redshift_specific_error_message(self): + """Test that error messages are Redshift-specific""" + config = { + 'connect': { + 'host': 'test-cluster.redshift.amazonaws.com', + 'port': '5439', + 'user': 'testuser', + 'password': 'testpass', + 'database': 'testdb', + 'auth_type': 'oauth' + }, + 'mode': Mock(), + 'streams': [] + } + + with patch('pontoon.source.postgresql_source.create_engine'): + with pytest.raises(ValueError) as exc_info: + RedshiftSource(config, Mock(), {}) + + # Verify the error message mentions Redshift specifically + assert "Redshift source only supports 'basic' authentication" in str(exc_info.value) + assert "got 'oauth'" in str(exc_info.value) \ No newline at end of file