Add support for continuously starting load jobs as slots free up in the loader

dlt/common/destination/reference.py

@@ -27,6 +27,7 @@
 from dlt.common import logger
 from dlt.common.configuration.specs.base_configuration import extract_inner_hint
 from dlt.common.destination.utils import verify_schema_capabilities
+from dlt.common.exceptions import TerminalValueError
 from dlt.common.normalizers.naming import NamingConvention
 from dlt.common.schema import Schema, TTableSchema, TSchemaTables
 from dlt.common.schema.utils import (
@@ -42,6 +43,8 @@
     InvalidDestinationReference,
     UnknownDestinationModule,
     DestinationSchemaTampered,
+    DestinationTransientException,
+    DestinationTerminalException,
 )
 from dlt.common.schema.exceptions import UnknownTableException
 from dlt.common.storages import FileStorage
@@ -258,11 +261,45 @@ class DestinationClientDwhWithStagingConfiguration(DestinationClientDwhConfigura
     """configuration of the staging, if present, injected at runtime"""
 
 
-TLoadJobState = Literal["running", "failed", "retry", "completed"]
+TLoadJobState = Literal["ready", "running", "failed", "retry", "completed"]
 
 
-class LoadJob:
-    """Represents a job that loads a single file
+class LoadJob(ABC):
+    """
+    A stateful load job, represents one job file
+    """
+
+    def __init__(self, file_path: str) -> None:
+        self._file_path = file_path
+        self._file_name = FileStorage.get_file_name_from_file_path(file_path)
+        # NOTE: we only accept a full filepath in the constructor
+        assert self._file_name != self._file_path
+        self._parsed_file_name = ParsedLoadJobFileName.parse(self._file_name)
+
+    def job_id(self) -> str:
+        """The job id that is derived from the file name and does not changes during job lifecycle"""
+        return self._parsed_file_name.job_id()
+
+    def file_name(self) -> str:
+        """A name of the job file"""
+        return self._file_name
+
+    def job_file_info(self) -> ParsedLoadJobFileName:
+        return self._parsed_file_name
+
+    @abstractmethod
+    def state(self) -> TLoadJobState:
+        """Returns current state. Should poll external resource if necessary."""
+        pass
+
+    @abstractmethod
+    def exception(self) -> str:
+        """The exception associated with failed or retry states"""
+        pass
+
+
+class RunnableLoadJob(LoadJob, ABC):
+    """Represents a runnable job that loads a single file
 
     Each job starts in "running" state and ends in one of terminal states: "retry", "failed" or "completed".
     Each job is uniquely identified by a file name. The file is guaranteed to exist in "running" state. In terminal state, the file may not be present.
@@ -273,75 +310,96 @@ class LoadJob:
     immediately transition job into "failed" or "retry" state respectively.
     """
 
-    def __init__(self, file_name: str) -> None:
+    def __init__(self, file_path: str) -> None:
         """
         File name is also a job id (or job id is deterministically derived) so it must be globally unique
         """
         # ensure file name
-        assert file_name == FileStorage.get_file_name_from_file_path(file_name)
-        self._file_name = file_name
-        self._parsed_file_name = ParsedLoadJobFileName.parse(file_name)
+        super().__init__(file_path)
+        self._state: TLoadJobState = "ready"
+        self._exception: Exception = None
 
-    @abstractmethod
-    def state(self) -> TLoadJobState:
-        """Returns current state. Should poll external resource if necessary."""
-        pass
+        # variables needed by most jobs, set by the loader in set_run_vars
+        self._schema: Schema = None
+        self._load_table: TTableSchema = None
+        self._load_id: str = None
+        self._job_client: "JobClientBase" = None
 
-    def file_name(self) -> str:
-        """A name of the job file"""
-        return self._file_name
+    def set_run_vars(self, load_id: str, schema: Schema, load_table: TTableSchema) -> None:
+        """
+        called by the loader right before the job is run
+        """
+        self._load_id = load_id
+        self._schema = schema
+        self._load_table = load_table
 
-    def job_id(self) -> str:
-        """The job id that is derived from the file name and does not changes during job lifecycle"""
-        return self._parsed_file_name.job_id()
+    @property
+    def load_table_name(self) -> str:
+        return self._load_table["name"]
 
-    def job_file_info(self) -> ParsedLoadJobFileName:
-        return self._parsed_file_name
+    def run_managed(
+        self,
+        job_client: "JobClientBase",
+    ) -> None:
+        """
+        wrapper around the user implemented run method
+        """
+        # only jobs that are not running or have not reached a final state
+        # may be started
+        assert self._state in ("ready", "retry")
+        self._job_client = job_client
+
+        # filepath is now moved to running
+        try:
+            self._state = "running"
+            self.run()
+            self._state = "completed"
+        except (DestinationTerminalException, TerminalValueError) as e:
+            logger.exception(f"Terminal problem when starting job {self.file_name}")
+            self._state = "failed"
+            self._exception = e
+        except (DestinationTransientException, Exception) as e:
+            logger.exception(f"Temporary problem when starting job {self.file_name}")
+            self._state = "retry"
+            self._exception = e
+        finally:
+            # sanity check
+            assert self._state in ("completed", "retry", "failed")
 
     @abstractmethod
+    def run(self) -> None:
+        """
+        run the actual job, this will be executed on a thread and should be implemented by the user
+        exception will be handled outside of this function
+        """
+        raise NotImplementedError()
+
+    def state(self) -> TLoadJobState:
+        """Returns current state. Should poll external resource if necessary."""
+        return self._state
+
     def exception(self) -> str:
         """The exception associated with failed or retry states"""
-        pass
+        return str(self._exception)
 
 
-class NewLoadJob(LoadJob):
-    """Adds a trait that allows to save new job file"""
+class FollowupJob:
+    """Base class for follow up jobs that should be created"""
 
     @abstractmethod
     def new_file_path(self) -> str:
         """Path to a newly created temporary job file. If empty, no followup job should be created"""
         pass
 
 
-class FollowupJob:
-    """Adds a trait that allows to create a followup job"""
+class HasFollowupJobs:
+    """Adds a trait that allows to create single or table chain followup jobs"""
 
-    def create_followup_jobs(self, final_state: TLoadJobState) -> List[NewLoadJob]:
+    def create_followup_jobs(self, final_state: TLoadJobState) -> List[FollowupJob]:
         """Return list of new jobs. `final_state` is state to which this job transits"""
         return []
 
 
-class DoNothingJob(LoadJob):
-    """The most lazy class of dlt"""
-
-    def __init__(self, file_path: str) -> None:
-        super().__init__(FileStorage.get_file_name_from_file_path(file_path))
-
-    def state(self) -> TLoadJobState:
-        # this job is always done
-        return "completed"
-
-    def exception(self) -> str:
-        # this part of code should be never reached
-        raise NotImplementedError()
-
-
-class DoNothingFollowupJob(DoNothingJob, FollowupJob):
-    """The second most lazy class of dlt"""
-
-    pass
-
-
 class JobClientBase(ABC):
     def __init__(
         self,
@@ -394,13 +452,10 @@ def update_stored_schema(
         return expected_update
 
     @abstractmethod
-    def start_file_load(self, table: TTableSchema, file_path: str, load_id: str) -> LoadJob:
-        """Creates and starts a load job for a particular `table` with content in `file_path`"""
-        pass
-
-    @abstractmethod
-    def restore_file_load(self, file_path: str) -> LoadJob:
-        """Finds and restores already started loading job identified by `file_path` if destination supports it."""
+    def create_load_job(
+        self, table: TTableSchema, file_path: str, load_id: str, restore: bool = False
+    ) -> LoadJob:
+        """Creates a load job for a particular `table` with content in `file_path`"""
         pass
 
     def should_truncate_table_before_load(self, table: TTableSchema) -> bool:
@@ -410,7 +465,7 @@ def create_table_chain_completed_followup_jobs(
         self,
         table_chain: Sequence[TTableSchema],
         completed_table_chain_jobs: Optional[Sequence[LoadJobInfo]] = None,
-    ) -> List[NewLoadJob]:
+    ) -> List[FollowupJob]:
         """Creates a list of followup jobs that should be executed after a table chain is completed"""
         return []
 

dlt/common/runtime/signals.py

@@ -32,6 +32,11 @@ def raise_if_signalled() -> None:
         raise SignalReceivedException(_received_signal)
 
 
+def signal_received() -> bool:
+    """check if a signal was received"""
+    return True if _received_signal else False
+
+
 def sleep(sleep_seconds: float) -> None:
     """A signal-aware version of sleep function. Will raise SignalReceivedException if signal was received during sleep period."""
     # do not allow sleeping if signal was received

dlt/common/typing.py

@@ -106,7 +106,7 @@
 VARIANT_FIELD_FORMAT = "v_%s"
 TFileOrPath = Union[str, PathLike, IO[Any]]
 TSortOrder = Literal["asc", "desc"]
-TLoaderFileFormat = Literal["jsonl", "typed-jsonl", "insert_values", "parquet", "csv"]
+TLoaderFileFormat = Literal["jsonl", "typed-jsonl", "insert_values", "parquet", "csv", "reference"]
 """known loader file formats"""
 
 

dlt/destinations/impl/athena/athena.py

@@ -45,10 +45,10 @@
 )
 from dlt.common.schema.utils import table_schema_has_type
 from dlt.common.destination import DestinationCapabilitiesContext
-from dlt.common.destination.reference import LoadJob, DoNothingFollowupJob, DoNothingJob
-from dlt.common.destination.reference import NewLoadJob, SupportsStagingDestination
+from dlt.common.destination.reference import LoadJob
+from dlt.common.destination.reference import FollowupJob, SupportsStagingDestination
 from dlt.common.data_writers.escape import escape_hive_identifier
-from dlt.destinations.sql_jobs import SqlStagingCopyJob, SqlMergeJob
+from dlt.destinations.sql_jobs import SqlStagingCopyFollowupJob, SqlMergeFollowupJob
 
 from dlt.destinations.typing import DBApi, DBTransaction
 from dlt.destinations.exceptions import (
@@ -65,6 +65,7 @@
 )
 from dlt.destinations.typing import DBApiCursor
 from dlt.destinations.job_client_impl import SqlJobClientWithStaging
+from dlt.destinations.job_impl import FinalizedLoadJobWithFollowupJobs, FinalizedLoadJob
 from dlt.destinations.impl.athena.configuration import AthenaClientConfiguration
 from dlt.destinations.type_mapping import TypeMapper
 from dlt.destinations import path_utils
@@ -160,7 +161,7 @@ def __init__(self) -> None:
         DLTAthenaFormatter._INSTANCE = self
 
 
-class AthenaMergeJob(SqlMergeJob):
+class AthenaMergeJob(SqlMergeFollowupJob):
     @classmethod
     def _new_temp_table_name(cls, name_prefix: str, sql_client: SqlClientBase[Any]) -> str:
         # reproducible name so we know which table to drop
@@ -468,40 +469,48 @@ def _get_table_update_sql(
                         LOCATION '{location}';""")
         return sql
 
-    def start_file_load(self, table: TTableSchema, file_path: str, load_id: str) -> LoadJob:
+    def create_load_job(
+        self, table: TTableSchema, file_path: str, load_id: str, restore: bool = False
+    ) -> LoadJob:
         """Starts SqlLoadJob for files ending with .sql or returns None to let derived classes to handle their specific jobs"""
         if table_schema_has_type(table, "time"):
             raise LoadJobTerminalException(
                 file_path,
                 "Athena cannot load TIME columns from parquet tables. Please convert"
                 " `datetime.time` objects in your data to `str` or `datetime.datetime`.",
             )
-        job = super().start_file_load(table, file_path, load_id)
+        job = super().create_load_job(table, file_path, load_id, restore)
         if not job:
             job = (
-                DoNothingFollowupJob(file_path)
+                FinalizedLoadJobWithFollowupJobs(file_path)
                 if self._is_iceberg_table(self.prepare_load_table(table["name"]))
-                else DoNothingJob(file_path)
+                else FinalizedLoadJob(file_path)
             )
         return job
 
-    def _create_append_followup_jobs(self, table_chain: Sequence[TTableSchema]) -> List[NewLoadJob]:
+    def _create_append_followup_jobs(
+        self, table_chain: Sequence[TTableSchema]
+    ) -> List[FollowupJob]:
         if self._is_iceberg_table(self.prepare_load_table(table_chain[0]["name"])):
             return [
-                SqlStagingCopyJob.from_table_chain(table_chain, self.sql_client, {"replace": False})
+                SqlStagingCopyFollowupJob.from_table_chain(
+                    table_chain, self.sql_client, {"replace": False}
+                )
             ]
         return super()._create_append_followup_jobs(table_chain)
 
     def _create_replace_followup_jobs(
         self, table_chain: Sequence[TTableSchema]
-    ) -> List[NewLoadJob]:
+    ) -> List[FollowupJob]:
         if self._is_iceberg_table(self.prepare_load_table(table_chain[0]["name"])):
             return [
-                SqlStagingCopyJob.from_table_chain(table_chain, self.sql_client, {"replace": True})
+                SqlStagingCopyFollowupJob.from_table_chain(
+                    table_chain, self.sql_client, {"replace": True}
+                )
             ]
         return super()._create_replace_followup_jobs(table_chain)
 
-    def _create_merge_followup_jobs(self, table_chain: Sequence[TTableSchema]) -> List[NewLoadJob]:
+    def _create_merge_followup_jobs(self, table_chain: Sequence[TTableSchema]) -> List[FollowupJob]:
         return [AthenaMergeJob.from_table_chain(table_chain, self.sql_client)]
 
     def _is_iceberg_table(self, table: TTableSchema) -> bool: