fix: patch ruff and mypy issues

Fix formatting and type conflicts

Author: DanielBlei

workloads/ingest/ingest/core/config.py

@@ -10,20 +10,20 @@
 
 class DatabaseType(str, Enum):
     """Database type enumeration."""
-    
+
     MYSQL = "mysql"
     POSTGRES = "postgres"
 
 
 class TargetType(str, Enum):
     """Target type enumeration."""
-    
+
     BIGQUERY = "bigquery"
 
 
 class SourceConfig(BaseModel):
     """Source database configuration."""
-    
+
     name: str
     type: DatabaseType
     host: str
@@ -51,7 +51,7 @@ def build_connection_string(self, dialect: str, default_port: int = 0) -> str:
 
 class BigQueryTargetConfig(BaseModel):
     """BigQuery target configuration."""
-    
+
     name: str
     type: TargetType = TargetType.BIGQUERY
     project_id: str = Field(
@@ -77,23 +77,23 @@ class BigQueryTargetConfig(BaseModel):
 
 class RuntimeParams(BaseModel):
     """Runtime parameters for ingestion process."""
-    
+
     retry_attempts: int = Field(ge=1, le=10, default=3)
     retry_delay_seconds: int = Field(ge=1, le=3600, default=30)
     chunk_size: int = Field(default=10000)
 
 
 class SecretProvider(str, Enum):
     """Secret provider enumeration."""
-    
+
     GOOGLE_SECRET_MANAGER = "gcloud"
     # AWS_SECRET_MANAGER = "aws"
     # AZURE_KEY_VAULT = "azure"
 
 
 class SecretConfig(BaseModel):
     """Secret configuration."""
-    
+
     provider: SecretProvider = Field(
         default=SecretProvider.GOOGLE_SECRET_MANAGER, description="Secret Manager Provider"
     )
@@ -107,7 +107,7 @@ class SecretConfig(BaseModel):
 
 class Config(BaseModel):
     """Main configuration class."""
-    
+
     version: str
     params: RuntimeParams
     secrets: List[SecretConfig]

workloads/ingest/ingest/main.py

@@ -129,7 +129,7 @@ def process_source(
     source_config.password = config.get_gcloud_secret_value(source_config.password, env)
 
     source = create_source(
-        source_config, retry_attempts=config.params.retry_attempts, retry_delay=config.params.retry_delay
+        source_config, retry_attempts=config.params.retry_attempts, retry_delay=config.params.retry_delay_seconds
     )
     try:
         # Process each table from this source

workloads/ingest/ingest/sources/mysql_source.py

@@ -1,6 +1,6 @@
 """MySQL source implementation."""
 
-from typing import Optional, Iterator
+from typing import Optional, Iterator, List
 from .source import Source
 from ..extractors import BaseExtractor
 from ..core.config import SourceConfig
@@ -29,7 +29,7 @@ def connect(self) -> None:
         super().connect()
         logger.debug("MySQL source connected via BaseExtractor")
 
-    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[dict]:
+    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[List[dict]]:
         """Extract data from MySQL table."""
         if not self.extractor:
             raise RuntimeError("MySQL extractor not initialized")

workloads/ingest/ingest/sources/postgres_source.py

@@ -1,6 +1,6 @@
 """PostgreSQL source implementation."""
 
-from typing import Optional, Iterator
+from typing import Optional, Iterator, List
 from .source import Source
 from ..extractors import BaseExtractor
 from ..core.config import SourceConfig
@@ -29,7 +29,7 @@ def connect(self) -> None:
         super().connect()
         logger.debug("PostgreSQL source connected via BaseExtractor")
 
-    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[dict]:
+    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[List[dict]]:
         """Extract data from PostgreSQL table."""
         if not self.extractor:
             raise RuntimeError("PostgreSQL extractor not initialized")

workloads/ingest/ingest/sources/source.py

@@ -1,6 +1,6 @@
 """Source interface and base implementations."""
 
-from typing import Optional, Protocol, runtime_checkable, Iterator
+from typing import Optional, Protocol, runtime_checkable, Iterator, List
 from pydantic import BaseModel, Field, ConfigDict
 from ..core.config import SourceConfig
 from ..core.catalog import Table
@@ -20,12 +20,15 @@ class SourceInterface(Protocol):
     def connect(self) -> None:
         """Connect to the data source."""
         ...
-    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[dict]:
+
+    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[List[dict]]:
         """Extract data from the source."""
         ...
+
     def validate_connection(self) -> bool:
         """Validate the connection to the source."""
         ...
+
     def close(self) -> None:
         """Close the connection to the source."""
         ...
@@ -49,7 +52,7 @@ def connect(self) -> None:
         if self.extractor:
             self.extractor.connect()
 
-    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[dict]:
+    def extract(self, table: Table, chunk_size: int, limit: Optional[int] = None) -> Iterator[List[dict]]:
         """Extract data from the source using SQLAlchemy with streaming.
 
         Args:

workloads/ingest/ingest/targets/bigquery.py

@@ -38,7 +38,7 @@ def _build_client(self) -> bigquery.Client:
             raise e
         return client
 
-    def _get_write_disposition(self, write_disposition: ReplicationType) -> bigquery.WriteDisposition:
+    def _get_write_disposition(self, write_disposition: ReplicationType) -> str:
         """Get the write disposition for the data."""
         if write_disposition not in map_replication_type_to_write_disposition:
             raise ValueError(f"Invalid write disposition: {write_disposition}")

workloads/ingest/ingest/targets/target.py

@@ -15,9 +15,10 @@ class TargetInterface(Protocol):
 
     config: TargetTypes
 
-    def load(self, data: list[dict], target_table: str) -> None:
+    def load(self, data: list[dict], target_table: str, write_disposition: ReplicationType) -> None:
         """Load data into the target."""
         ...
+
     def validate_connection(self) -> bool:
         """Validate the connection to the target."""
         ...