Refactor SQLSource from a monolith into a proper class hierarchy (#36)

Author: kmacrow

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'
 }

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

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)

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}'")

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)