fdc-shared-kernel 0.0.197__tar.gz → 0.0.199__tar.gz

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fdc_shared_kernel
-Version: 0.0.197
+Version: 0.0.199
 Summary: Shared library for microservice
 Author-email: Shikhil S <shikhil.s@dbizsolution.com>, Ahammed Akdham N <ahammedakdham.n@dbizsolution.com>
 Classifier: Programming Language :: Python :: 3

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/fdc_shared_kernel.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fdc_shared_kernel
-Version: 0.0.197
+Version: 0.0.199
 Summary: Shared library for microservice
 Author-email: Shikhil S <shikhil.s@dbizsolution.com>, Ahammed Akdham N <ahammedakdham.n@dbizsolution.com>
 Classifier: Programming Language :: Python :: 3

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/fdc_shared_kernel.egg-info/SOURCES.txt

@@ -21,6 +21,8 @@ shared_kernel/auth/__init__.py
 shared_kernel/auth/jwt_helper.py
 shared_kernel/auth/token_handler.py
 shared_kernel/auth/workbook_permission_handler.py
+shared_kernel/aws/__init__.py
+shared_kernel/aws/athena_client.py
 shared_kernel/batch_job/__init__.py
 shared_kernel/batch_job/batch_job_status_tracker/container_job_status_tracker.py
 shared_kernel/batch_job/job_handlers/__init__.py
@@ -36,6 +38,7 @@ shared_kernel/config/__init__.py
 shared_kernel/constants/constants.py
 shared_kernel/data_warehouse_handlers/__init__.py
 shared_kernel/data_warehouse_handlers/base_executor.py
+shared_kernel/data_warehouse_handlers/csv_handler.py
 shared_kernel/data_warehouse_handlers/databricks_handler.py
 shared_kernel/data_warehouse_handlers/fabric_sql_handler.py
 shared_kernel/data_warehouse_handlers/factory.py
@@ -49,6 +52,7 @@ shared_kernel/data_warehouse_handlers/redshift_handler.py
 shared_kernel/data_warehouse_handlers/snowflake_handler.py
 shared_kernel/data_warehouse_handlers/utils.py
 shared_kernel/data_warehouse_handlers/warehouse_handlers_map.py
+shared_kernel/data_warehouse_handlers/connections/csv_warehouse_connection.py
 shared_kernel/data_warehouse_handlers/connections/databricks_warehouse_connection.py
 shared_kernel/data_warehouse_handlers/connections/fabric_sql_warehouse_connection.py
 shared_kernel/data_warehouse_handlers/connections/jdbc_warehouse_connection.py
@@ -58,6 +62,7 @@ shared_kernel/data_warehouse_handlers/connections/postgresql_warehouse_connectio
 shared_kernel/data_warehouse_handlers/connections/redshift_warehouse_connection.py
 shared_kernel/data_warehouse_handlers/connections/snowflake_warehouse_connection.py
 shared_kernel/data_warehouse_handlers/custom_query_builders/databricks_query.py
+shared_kernel/data_warehouse_handlers/query_executors/csv_query_executor.py
 shared_kernel/data_warehouse_handlers/query_executors/databricks_query_executor.py
 shared_kernel/data_warehouse_handlers/query_executors/fabric_sql_query_executor.py
 shared_kernel/data_warehouse_handlers/query_executors/jdbc_query_executor.py

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fdc_shared_kernel"
-version = "0.0.197"
+version = "0.0.199"
 requires-python = ">=3.7"
 readme = "README_pypi.md"
 description = "Shared library for microservice"

fdc_shared_kernel-0.0.199/shared_kernel/aws/athena_client.py

@@ -0,0 +1,27 @@
+import boto3
+
+
+def get_assumed_role_credentials(role_arn: str, service_name: str) -> dict:
+    sts_client = boto3.client("sts")
+    assumed_role = sts_client.assume_role(
+        RoleArn=role_arn, RoleSessionName=f"{str(service_name).capitalize()}Session"
+    )
+    return assumed_role["Credentials"]
+
+
+def build_aws_client(service_name: str, role_arn: str = None, region: str = None):
+    """Build a boto3 client for `service_name`, optionally assuming `role_arn`."""
+    if role_arn:
+        credentials = get_assumed_role_credentials(role_arn, service_name)
+        return boto3.client(
+            service_name,
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+            region_name=region,
+        )
+    return boto3.client(service_name, region_name=region) if region else boto3.client(service_name)
+
+
+def build_athena_client(role_arn: str = None, region: str = None):
+    return build_aws_client("athena", role_arn=role_arn, region=region)

fdc_shared_kernel-0.0.199/shared_kernel/data_warehouse_handlers/connections/csv_warehouse_connection.py

@@ -0,0 +1,161 @@
+import time
+from contextlib import contextmanager
+
+from shared_kernel.aws.athena_client import build_athena_client, build_aws_client
+from shared_kernel.config import Config
+from shared_kernel.dataclasses.warehouse_configs import CSVConfig
+from shared_kernel.exceptions.http_exceptions import BadRequest
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection, cancel_check_connection
+from shared_kernel.logger import Logger
+
+config = Config()
+logger = Logger(config.get("APP_NAME"))
+
+TERMINAL_STATES = ("SUCCEEDED", "FAILED", "CANCELLED")
+
+
+class AthenaCursor:
+    """Cursor-like adapter over boto3 Athena start_query_execution/poll/get_query_results,
+    so Athena's async query model can be used behind the same cursor-shaped contract.
+    """
+
+    def __init__(
+        self,
+        athena_client,
+        database: str,
+        output_location: str,
+        role_arn: str = None,
+        region: str = None,
+        poll_interval: float = 1.0,
+    ):
+        self._client = athena_client
+        self._database = database
+        self._output_location = output_location
+        self._role_arn = role_arn
+        self._region = region
+        self._poll_interval = poll_interval
+        self._columns = []
+        self._column_types = []
+        self._rows = []
+        self._exec_id = None
+
+    @property
+    def description(self):
+        # DB-API-shaped: (name, type_code, ...) — type_code holds the Athena/
+        # Presto native type string (e.g. "varchar", "bigint")
+        return [
+            (name, type_code, None, None, None, None, None)
+            for name, type_code in zip(self._columns, self._column_types)
+        ]
+
+    def execute(self, query: str, params: tuple = None):
+        exec_response = self._client.start_query_execution(
+            QueryString=query,
+            QueryExecutionContext={"Database": self._database},
+            ResultConfiguration={"OutputLocation": self._output_location},
+            WorkGroup="primary",
+        )
+        self._exec_id = exec_response["QueryExecutionId"]
+
+        while True:
+            stats = self._client.get_query_execution(QueryExecutionId=self._exec_id)
+            state = stats["QueryExecution"]["Status"]["State"]
+            if state in TERMINAL_STATES:
+                break
+            time.sleep(self._poll_interval)
+
+        if state != "SUCCEEDED":
+            reason = stats["QueryExecution"]["Status"].get("StateChangeReason", "Unknown error")
+            raise Exception(f"Athena query {state.lower()}: {reason}")
+
+        response = self._client.get_query_results(QueryExecutionId=self._exec_id, MaxResults=1000)
+        column_info = response["ResultSet"]["ResultSetMetadata"]["ColumnInfo"]
+        self._columns = [c.get("Name", "") for c in column_info]
+        self._column_types = [c.get("Type", "varchar") for c in column_info]
+
+        # Athena's first result row echoes the column headers (as VarCharValue),
+        # skip it — real column names/order already come from ColumnInfo above.
+        rows_data = response["ResultSet"]["Rows"]
+        data_rows = rows_data[1:] if rows_data else []
+        self._rows = [
+            tuple(c.get("VarCharValue", "") for c in r["Data"]) for r in data_rows
+        ]
+
+        self._delete_output()
+        return self
+
+    def _delete_output(self):
+        try:
+            s3_client = build_aws_client("s3", role_arn=self._role_arn, region=self._region)
+            bucket, prefix = _parse_s3_location(self._output_location)
+            s3_client.delete_object(Bucket=bucket, Key=f"{prefix}{self._exec_id}.csv")
+            s3_client.delete_object(Bucket=bucket, Key=f"{prefix}{self._exec_id}.csv.metadata")
+        except Exception as e:
+            logger.warning(f"Failed to delete Athena output for {self._exec_id}: {e}")
+
+    def fetchall(self):
+        return self._rows
+
+    def fetchone(self):
+        return self._rows[0] if self._rows else None
+
+    def cancel(self):
+        if self._exec_id:
+            self._client.stop_query_execution(QueryExecutionId=self._exec_id)
+
+
+def _parse_s3_location(s3_uri: str):
+    without_scheme = s3_uri.replace("s3://", "", 1)
+    bucket, _, prefix = without_scheme.partition("/")
+    if prefix and not prefix.endswith("/"):
+        prefix += "/"
+    return bucket, prefix
+
+
+class CSVWarehouseConnection(DataWarehouseConnection):
+    """Athena/Glue-backed connection for the CSV data source.
+
+    Athena is stateless and per-query (no persistent session), so unlike the
+    SQLAlchemy-pooled connections this class has no engine/connection to keep
+    alive between queries.
+    """
+
+    def __init__(self, source_config: CSVConfig):
+        if source_config is None:
+            raise ValueError("Missing required parameter: source_config (CSVConfig)")
+        self._config = source_config
+        self._client = None
+
+    def _get_client(self):
+        if self._client is None:
+            self._client = build_athena_client(
+                role_arn=self._config.athena_cross_account_role_arn,
+                region=self._config.region,
+            )
+        return self._client
+
+    @contextmanager
+    @cancel_check_connection
+    def get_connection(self, retries: int = 3, delay: int = 5, query_timeout: int = 30):
+        cursor = AthenaCursor(
+            athena_client=self._get_client(),
+            database=self._config.glue_database,
+            output_location=self._config.athena_output_location,
+            role_arn=self._config.athena_cross_account_role_arn,
+            region=self._config.region,
+        )
+        yield cursor
+
+    def is_valid_connection(self) -> bool:
+        return True
+
+    def close_connection(self):
+        self._client = None
+
+    def test_connection(self) -> bool:
+        try:
+            with self.get_connection(retries=1) as cursor:
+                cursor.execute("SELECT 1")
+            return True
+        except Exception as e:
+            raise BadRequest(f"Athena/Glue connection test failed: {str(e)}") from e

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/data_warehouse_handlers/connections/fabric_sql_warehouse_connection.py

@@ -69,6 +69,7 @@ class FabricSQLWarehouseConnection(DataWarehouseConnection):
             f"DATABASE={cfg.database};"
             "Encrypt=yes;"
             "TrustServerCertificate=no;"
+            "Connection Timeout=10;"
         )
         # We use odbc_connect query param to prevent SQLAlchemy dialect from auto-injecting
         # conflicting credentials/parameters (avoiding error FA005).

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/data_warehouse_handlers/connections/mssql_warehouse_connection.py

@@ -60,7 +60,7 @@ class MSSQLWarehouseConnection(DataWarehouseConnection):
                     connection_string = get_constant("mssql", "CONNECTION_STRING").format(
                         cfg.username, password, cfg.host, int(cfg.port), cfg.database
                     )
-                    connect_args = {}
+                    connect_args = {"timeout": 10}
 
                     pool_size = int(Config.get("MSSQL_POOL_SIZE", 10))
 

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/data_warehouse_handlers/connections/mysql_warehouse_connection.py

@@ -31,7 +31,8 @@ class MySQLWarehouseConnection(DataWarehouseConnection):
     _CONNECT_ARGS = {
         "conv": {
             pymysql.FIELD_TYPE.BIT: lambda data: data == b'\x01'
-        }
+        },
+        "connect_timeout": 10,
     }
 
     def __init__(

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/data_warehouse_handlers/connections/postgresql_warehouse_connection.py

@@ -64,6 +64,7 @@ class PostgreSQLWarehouseConnection(DataWarehouseConnection):
                     pool_size = int(Config.get("POSTGRES_POOL_SIZE", 10))
                     self._engines[key] = create_engine(
                         url,
+                        connect_args={"connect_timeout": 10},
                         pool_size=pool_size,
                         max_overflow=5,
                         pool_timeout=30,

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/data_warehouse_handlers/connections/redshift_warehouse_connection.py

@@ -197,6 +197,7 @@ class RedshiftWarehouseConnection(DataWarehouseConnection):
                     pool_size = int(Config.get("REDSHIFT_POOL_SIZE", 10))
                     self._engines[key] = create_engine(
                         url,
+                        connect_args={"timeout": 10},
                         pool_size=pool_size,
                         max_overflow=5,
                         pool_timeout=30,

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/data_warehouse_handlers/connections/snowflake_warehouse_connection.py

@@ -63,6 +63,7 @@ class SnowflakeWarehouseConnection(DataWarehouseConnection):
                 schema=cfg.schema,
                 private_key=pkb,
                 authenticator="snowflake_jwt",
+                validate_default_parameters=True,
             )
         else:
             self._conn = snowflake.connector.connect(
@@ -71,6 +72,7 @@ class SnowflakeWarehouseConnection(DataWarehouseConnection):
                 account=cfg.account,
                 database=cfg.database,
                 schema=cfg.schema,
+                validate_default_parameters=True,
             )
         self._cursor = self._conn.cursor()
 

fdc_shared_kernel-0.0.199/shared_kernel/data_warehouse_handlers/csv_handler.py

@@ -0,0 +1,53 @@
+from pypika import Query
+
+from shared_kernel.data_warehouse_handlers.connections.csv_warehouse_connection import CSVWarehouseConnection
+from shared_kernel.data_warehouse_handlers.query_executors.csv_query_executor import CSVQueryExecutor
+from shared_kernel.dataclasses.warehouse_configs import CSVConfig
+from shared_kernel.interfaces.data_warehouse_handler import WarehouseHandler
+
+
+class CSVWarehouseHandler(WarehouseHandler):
+
+    def get_config(self) -> CSVConfig:
+        required = ("glue_database", "region", "s3_bucket", "s3_prefix", "athena_output_location")
+        missing = [f for f in required if not self.payload.get(f)]
+        if missing:
+            raise ValueError(f"CSV payload is missing required fields: {missing}.")
+
+        return CSVConfig(
+            glue_database=str(self.payload.get("glue_database")).strip(),
+            region=str(self.payload.get("region")).strip(),
+            s3_bucket=str(self.payload.get("s3_bucket")).strip(),
+            s3_prefix=str(self.payload.get("s3_prefix")).strip(),
+            athena_output_location=str(self.payload.get("athena_output_location")).strip(),
+            athena_cross_account_role_arn=self.payload.get("athena_cross_account_role_arn"),
+            schema=self.payload.get("schema"),
+            schema_list=self.payload.get("schema_list"),
+            pagination=self.pagination,
+        )
+
+    def _build_connection(self) -> CSVWarehouseConnection:
+        """Internal — not exposed to callers."""
+        return CSVWarehouseConnection(source_config=self.get_config())
+
+    def get_executor(self) -> CSVQueryExecutor:
+        return CSVQueryExecutor(warehouse_connection=self._build_connection())
+
+    def get_query_object(self) -> Query:
+        return Query
+
+    def create_records_and_count_query(
+        self,
+        query: str,
+        limit: str,
+        page: str,
+    ) -> tuple[str, str]:
+        count_query = f"SELECT COUNT(*) FROM ({query}) AS _count"
+
+        if limit:
+            query = f"{query} LIMIT {int(limit)}"
+            if page:
+                offset = (int(page) - 1) * int(limit)
+                query = f"{query} OFFSET {offset}"
+
+        return count_query, query

fdc_shared_kernel-0.0.199/shared_kernel/data_warehouse_handlers/query_executors/csv_query_executor.py

@@ -0,0 +1,105 @@
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Dict, List, Optional, Tuple
+
+from shared_kernel.config import Config
+from shared_kernel.data_warehouse_handlers.base_executor import BaseQueryExecutor, cancel_check
+from shared_kernel.interfaces.warehouse_connection import DataWarehouseConnection
+from shared_kernel.logger import Logger
+
+config = Config()
+logger = Logger(config.get("APP_NAME"))
+MAX_QUERY_CONCURRENCY = 5
+
+
+class CSVQueryExecutor(BaseQueryExecutor):
+    """Query executor for the Athena/Glue-backed CSV data source."""
+
+    def __init__(self, warehouse_connection: DataWarehouseConnection) -> None:
+        self.warehouse_connection = warehouse_connection
+
+    @cancel_check
+    def execute_query(
+        self,
+        query: str,
+        params: Optional[tuple] = None,
+    ) -> List[Dict[str, Any]]:
+        with self.warehouse_connection.get_connection() as cursor:
+            self.active_connections.add(cursor)
+            try:
+                logger.info(f"Executing Athena query: {query}")
+                start = time.perf_counter()
+                cursor.execute(query, params)
+                elapsed = time.perf_counter() - start
+                logger.info(f"Athena query executed in {elapsed:.6f}s")
+
+                columns = [col[0] for col in cursor.description]
+                return [dict(zip(columns, row)) for row in cursor.fetchall()]
+            except Exception as e:
+                logger.error(f"Athena query failed: {query} — {e}", exc_info=True)
+                raise Exception(f"Query execution failed: {str(e)}") from e
+            finally:
+                self.active_connections.discard(cursor)
+
+    def execute_queries(
+        self,
+        queries: List[str],
+    ) -> List[List[Dict[str, Any]]]:
+        if not queries:
+            return []
+
+        indexed: List[Tuple[int, str]] = list(enumerate(queries))
+        results: List[Optional[List[Dict[str, Any]]]] = [None] * len(queries)
+
+        with ThreadPoolExecutor(max_workers=MAX_QUERY_CONCURRENCY) as pool:
+            future_to_index = {
+                pool.submit(self.execute_query, query): idx
+                for idx, query in indexed
+            }
+
+            for future in as_completed(future_to_index):
+                idx = future_to_index[future]
+                try:
+                    results[idx] = future.result()
+                except Exception as e:
+                    for pending in future_to_index:
+                        if not pending.done():
+                            pending.cancel()
+                    raise Exception(
+                        f"Parallel query execution failed on query index {idx}: {e}"
+                    ) from e
+
+        return results
+
+    @cancel_check
+    def execute_query_with_column_types(
+        self,
+        query: str,
+        params: Optional[tuple] = None,
+    ) -> Tuple[List[Dict[str, Any]], List[Tuple]]:
+        with self.warehouse_connection.get_connection() as cursor:
+            self.active_connections.add(cursor)
+            try:
+                logger.info(f"Executing Athena query with column types: {query}")
+                start = time.perf_counter()
+                cursor.execute(query, params)
+                elapsed = time.perf_counter() - start
+                logger.info(f"Athena query executed in {elapsed:.6f}s")
+
+                description = cursor.description or []
+                if not description:
+                    return [], []
+
+                columns = [col[0] for col in description]
+                rows = [dict(zip(columns, row)) for row in cursor.fetchall()]
+                return rows, description
+            except Exception as e:
+                logger.error(f"Athena query failed: {query} — {e}", exc_info=True)
+                raise Exception(f"Query execution failed: {str(e)}") from e
+            finally:
+                self.active_connections.discard(cursor)
+
+    def cancel(self) -> None:
+        logger.warning(f"Cancelling {len(self.active_connections)} active Athena query cursor(s).")
+        for cursor in list(self.active_connections):
+            self._cancel_connection(cursor)

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/data_warehouse_handlers/warehouse_handlers_map.py

@@ -8,6 +8,7 @@ from shared_kernel.data_warehouse_handlers.redshift_handler import RedshiftWareh
 from shared_kernel.data_warehouse_handlers.snowflake_handler import SnowflakeWarehouseHandler
 from shared_kernel.data_warehouse_handlers.jdbc_handler import JDBCWarehouseHandler
 from shared_kernel.data_warehouse_handlers.fabric_sql_handler import FabricSQLWarehouseHandler
+from shared_kernel.data_warehouse_handlers.csv_handler import CSVWarehouseHandler
 from shared_kernel.enums.connection_engine import ConnectionEngine
 
 
@@ -22,4 +23,5 @@ warehouse_handlers = {
     ConnectionEngine.DATABRICKS.value: DatabricksWarehouseHandler,
     ConnectionEngine.JDBC.value: JDBCWarehouseHandler,
     ConnectionEngine.FABRIC_SQL.value: FabricSQLWarehouseHandler,
+    ConnectionEngine.CSV.value: CSVWarehouseHandler,
 }

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/dataclasses/warehouse_configs.py

@@ -112,6 +112,19 @@ class SnowflakeConfig(BaseConfig):
     default_schema: str = "PUBLIC"
 
 
+@dataclass
+class CSVConfig(BaseConfig):
+    """Configuration for the AWS Glue/Athena-backed CSV data source."""
+
+    glue_database: str
+    region: str
+    s3_bucket: str
+    s3_prefix: str
+    athena_output_location: str
+    athena_cross_account_role_arn: Optional[str] = None
+    default_schema: Optional[str] = None
+
+
 @dataclass
 class JDBCConfig(BaseConfig):
     """Configuration for JDBC-based connections."""

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/datatype_mappings/connectors_to_system/__init__.py

@@ -5,4 +5,5 @@ from shared_kernel.datatype_mappings.connectors_to_system.databricks import data
 from shared_kernel.datatype_mappings.connectors_to_system.mysql import mysql_to_system
 from shared_kernel.datatype_mappings.connectors_to_system.snowflake import snowflake_to_system
 from shared_kernel.datatype_mappings.connectors_to_system.sqlalchemy import sqlalchemy_to_system
+from shared_kernel.datatype_mappings.connectors_to_system.csv import csv_to_system
 

fdc_shared_kernel-0.0.199/shared_kernel/datatype_mappings/connectors_to_system/csv.py

@@ -0,0 +1,23 @@
+csv_to_system = {
+    "VARCHAR": "string",
+    "CHAR": "string",
+    "STRING": "string",
+    "BOOLEAN": "boolean",
+    "TINYINT": "integer",
+    "SMALLINT": "integer",
+    "INTEGER": "integer",
+    "INT": "integer",
+    "BIGINT": "integer",
+    "REAL": "double",
+    "FLOAT": "double",
+    "DOUBLE": "double",
+    "DECIMAL": "numeric",
+    "DATE": "date",
+    "TIMESTAMP": "timestamp",
+    "TIMESTAMP WITH TIME ZONE": "timestamp",
+    "VARBINARY": "binary",
+    "JSON": "json",
+    "ARRAY": "array",
+    "MAP": "json",
+    "ROW": "json",
+}

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/enums/connection_engine.py

@@ -12,3 +12,4 @@ class ConnectionEngine(Enum):
     FACTWEAVERS_DW = "factweavers-dw"
     JDBC = "jdbc"
     FABRIC_SQL = "fabric_sql"
+    CSV = "csv"

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/schema_onboarding_tracker/schema_onboarding_tracker.py

@@ -157,6 +157,17 @@ class SchemaOnboardingTracker(BaseSchemaOnboardingTracker):
             task_id, status, output, error, error_message, metadata
         )
 
+    def claim_task(
+        self, task_id: str, expected_status: TaskStatus, claimed_status: TaskStatus = TaskStatus.RUNNING
+    ) -> bool:
+        """
+        Atomically transitions a task from expected_status to claimed_status.
+        Returns True only if this call won the transition — see
+        TaskManager.claim_task for why this exists (fan-in double-dispatch
+        prevention).
+        """
+        return self._task_manager.claim_task(task_id, expected_status, claimed_status)
+
     def increment_task_attempts(self, task_id: str) -> Dict[str, Any]:
         """
         Increment attempts for retries.

{fdc_shared_kernel-0.0.197 → fdc_shared_kernel-0.0.199}/shared_kernel/schema_onboarding_tracker/task_manager.py

@@ -87,6 +87,42 @@ class TaskManager(BaseSchemaOnboardingTracker):
         
         return self._make_request("UPDATE_TASK_STATUS", payload)
 
+    def claim_task(
+        self, task_id: str, expected_status: TaskStatus, claimed_status: TaskStatus = TaskStatus.RUNNING
+    ) -> bool:
+        """
+        Atomically transitions a task from expected_status to claimed_status,
+        returning True only if this call won the transition.
+
+        Used to prevent a task from being dispatched twice when multiple
+        upstream branches of a fan-in join (e.g. two entities' chains both
+        feeding a shared downstream task) complete concurrently and each
+        independently decides the joined task is now eligible — a plain
+        update_task_status() call is not conditional and both callers would
+        otherwise "succeed".
+
+        Args:
+            task_id (str): Task ID
+            expected_status (TaskStatus): Status the task must currently be in
+            claimed_status (TaskStatus): Status to set if the claim succeeds
+
+        Returns:
+            bool: True if this call won the claim, False if a concurrent
+                caller already claimed it first (or the task was already past
+                expected_status for any other reason).
+        """
+        self._validate_task_id(task_id)
+
+        payload = {
+            "task_id": task_id,
+            "status": claimed_status.value,
+            "expected_status": expected_status.value,
+        }
+
+        response = self._make_request("UPDATE_TASK_STATUS", payload)
+        data = self._get_data_from_response(response) or {}
+        return bool(data.get("claimed"))
+
     def increment_task_attempts(self, task_id: str) -> Dict[str, Any]:
         """
         Increment attempts for retries.

fdc_shared_kernel-0.0.197/shared_kernel/datatype_mappings/connectors_to_system/csv.py

@@ -1 +0,0 @@
-csv_to_system = {}