Add support for model jobs

Author: sh-rp

dlt/common/data_writers/writers.py

@@ -44,7 +44,7 @@
     from dlt.common.libs.pyarrow import pyarrow as pa
 
 
-TDataItemFormat = Literal["arrow", "object", "file"]
+TDataItemFormat = Literal["arrow", "object", "file", "model"]
 TWriter = TypeVar("TWriter", bound="DataWriter")
 
 
@@ -59,6 +59,8 @@ class FileWriterSpec(NamedTuple):
     """File format supports changes of schema: True - at any moment, Buffer - in memory buffer before opening file,  False - not at all"""
     requires_destination_capabilities: bool = False
     supports_compression: bool = False
+    file_max_items: Optional[int] = None
+    """Set an upper limit on the number of items in one file"""
 
 
 EMPTY_DATA_WRITER_METRICS = DataWriterMetrics("", 0, 0, 2**32, 0.0)
@@ -115,6 +117,8 @@ def item_format_from_file_extension(cls, extension: str) -> TDataItemFormat:
             return "object"
         elif extension == "parquet":
             return "arrow"
+        elif extension == "model":
+            return "model"
         # those files may be imported by normalizer as is
         elif extension in LOADER_FILE_FORMATS:
             return "file"
@@ -175,6 +179,32 @@ def writer_spec(cls) -> FileWriterSpec:
         )
 
 
+class ModelWriter(DataWriter):
+    """Writes incoming items row by row into a text file and ensures a trailing ;"""
+
+    def write_header(self, columns_schema: TTableSchemaColumns) -> None:
+        pass
+
+    def write_data(self, items: Sequence[TDataItem]) -> None:
+        super().write_data(items)
+        self.items_count += len(items)
+        for item in items:
+            self._f.write(item + "\n")
+
+    @classmethod
+    def writer_spec(cls) -> FileWriterSpec:
+        return FileWriterSpec(
+            "model",
+            "model",
+            file_extension="model",
+            is_binary_format=False,
+            supports_schema_changes="True",
+            supports_compression=False,
+            # NOTE: we create a new model file for each sql row
+            file_max_items=1,
+        )
+
+
 class TypedJsonlListWriter(JsonlWriter):
     def write_data(self, items: Sequence[TDataItem]) -> None:
         # skip JsonlWriter when calling super
@@ -670,6 +700,7 @@ def is_native_writer(writer_type: Type[DataWriter]) -> bool:
     ArrowToJsonlWriter,
     ArrowToTypedJsonlListWriter,
     ArrowToCsvWriter,
+    ModelWriter,
 ]
 
 WRITER_SPECS: Dict[FileWriterSpec, Type[DataWriter]] = {
@@ -689,6 +720,11 @@ def is_native_writer(writer_type: Type[DataWriter]) -> bool:
         for writer in ALL_WRITERS
         if writer.writer_spec().data_item_format == "arrow" and is_native_writer(writer)
     ),
+    "model": tuple(
+        writer
+        for writer in ALL_WRITERS
+        if writer.writer_spec().data_item_format == "model" and is_native_writer(writer)
+    ),
 }
 
 

dlt/common/destination/utils.py

@@ -152,7 +152,7 @@ def verify_supported_data_types(
     for parsed_file in new_jobs:
         formats = table_file_formats.setdefault(parsed_file.table_name, set())
         if parsed_file.file_format in LOADER_FILE_FORMATS:
-            formats.add(parsed_file.file_format)  # type: ignore[arg-type]
+            formats.add(parsed_file.file_format)
     # all file formats
     all_file_formats = set(capabilities.supported_loader_file_formats or []) | set(
         capabilities.supported_staging_file_formats or []

dlt/common/storages/data_item_storage.py

@@ -28,7 +28,9 @@ def _get_writer(
         if not writer:
             # assign a writer for each table
             path = self._get_data_item_path_template(load_id, schema_name, table_name)
-            writer = BufferedDataWriter(self.writer_spec, path)
+            writer = BufferedDataWriter(
+                self.writer_spec, path, file_max_items=self.writer_spec.file_max_items
+            )
             self.buffered_writers[writer_id] = writer
         return writer
 

dlt/common/typing.py

@@ -127,7 +127,9 @@ class SecretSentinel:
 VARIANT_FIELD_FORMAT = "v_%s"
 TFileOrPath = Union[str, PathLike, IO[Any]]
 TSortOrder = Literal["asc", "desc"]
-TLoaderFileFormat = Literal["jsonl", "typed-jsonl", "insert_values", "parquet", "csv", "reference"]
+TLoaderFileFormat = Literal[
+    "jsonl", "typed-jsonl", "insert_values", "parquet", "csv", "reference", "sql", "model"
+]
 """known loader file formats"""
 
 TDynHintType = TypeVar("TDynHintType")

dlt/destinations/impl/duckdb/factory.py

@@ -129,7 +129,7 @@ class duckdb(Destination[DuckDbClientConfiguration, "DuckDbClient"]):
     def _raw_capabilities(self) -> DestinationCapabilitiesContext:
         caps = DestinationCapabilitiesContext()
         caps.preferred_loader_file_format = "insert_values"
-        caps.supported_loader_file_formats = ["insert_values", "parquet", "jsonl"]
+        caps.supported_loader_file_formats = ["insert_values", "parquet", "jsonl", "model"]
         caps.preferred_staging_file_format = None
         caps.supported_staging_file_formats = []
         caps.type_mapper = DuckDbTypeMapper

dlt/destinations/job_client_impl.py

@@ -113,6 +113,26 @@ def is_sql_job(file_path: str) -> bool:
         return os.path.splitext(file_path)[1][1:] == "sql"
 
 
+class ModelLoadJob(RunnableLoadJob):
+    """
+    A job to insert rows into a table from a model file which contains a list of select statements
+    """
+
+    def __init__(self, file_path: str) -> None:
+        super().__init__(file_path)
+        self._job_client: "SqlJobClientBase" = None
+
+    def run(self) -> None:
+        with FileStorage.open_zipsafe_ro(self._file_path, "r", encoding="utf-8") as f:
+            sql = f.read()
+        self._sql_client = self._job_client.sql_client
+        self._sql_client.execute_sql(sql)
+
+    @staticmethod
+    def is_model_job(file_path: str) -> bool:
+        return os.path.splitext(file_path)[1][1:] == "model"
+
+
 class CopyRemoteFileLoadJob(RunnableLoadJob, HasFollowupJobs):
     def __init__(
         self,
@@ -279,6 +299,9 @@ def create_load_job(
         if SqlLoadJob.is_sql_job(file_path):
             # create sql load job
             return SqlLoadJob(file_path)
+        if ModelLoadJob.is_model_job(file_path):
+            # create model load job
+            return ModelLoadJob(file_path)
         return None
 
     def complete_load(self, load_id: str) -> None:

dlt/extract/extract.py

@@ -45,7 +45,7 @@
 from dlt.extract.reference import SourceReference
 from dlt.extract.resource import DltResource
 from dlt.extract.storage import ExtractStorage
-from dlt.extract.extractors import ObjectExtractor, ArrowExtractor, Extractor
+from dlt.extract.extractors import ObjectExtractor, ArrowExtractor, Extractor, TextExtractor
 from dlt.extract.utils import get_data_item_format
 
 
@@ -343,6 +343,9 @@ def _extract_single_source(
             "arrow": ArrowExtractor(
                 load_id, self.extract_storage.item_storages["arrow"], schema, collector=collector
             ),
+            "model": TextExtractor(
+                load_id, self.extract_storage.item_storages["model"], schema, collector=collector
+            ),
         }
         # make sure we close storage on exception
         with collector(f"Extract {source.name}"):
@@ -363,7 +366,7 @@ def _extract_single_source(
                             collector.update("Resources", delta)
                         signals.raise_if_signalled()
                         resource = source.resources[pipe_item.pipe.name]
-                        item_format = get_data_item_format(pipe_item.item)
+                        item_format = get_data_item_format(pipe_item.item, pipe_item.meta)
                         extractors[item_format].write_items(
                             resource, pipe_item.item, pipe_item.meta
                         )

dlt/extract/extractors.py

@@ -62,7 +62,7 @@ def __init__(
         hints: TResourceHints = None,
         create_table_variant: bool = None,
     ) -> None:
-        super().__init__(hints, create_table_variant)
+        super().__init__(hints=hints, create_table_variant=create_table_variant)
         self.file_path = file_path
         self.metrics = metrics
         self.file_format = file_format
@@ -292,6 +292,12 @@ class ObjectExtractor(Extractor):
     pass
 
 
+class TextExtractor(Extractor):
+    """Extracts text items and writes them row by row into a text file"""
+
+    pass
+
+
 class ArrowExtractor(Extractor):
     """Extracts arrow data items into parquet. Normalizes arrow items column names.
     Compares the arrow schema to actual dlt table schema to reorder the columns and to

dlt/extract/hints.py

@@ -43,6 +43,7 @@
 from dlt.extract.items_transform import ValidateItem
 from dlt.extract.utils import ensure_table_schema_columns, ensure_table_schema_columns_hint
 from dlt.extract.validation import create_item_validator
+from dlt.common.data_writers import TDataItemFormat
 
 
 class TResourceHintsBase(TypedDict, total=False):
@@ -68,11 +69,17 @@ class TResourceHints(TResourceHintsBase, total=False):
 
 
 class HintsMeta:
-    __slots__ = ("hints", "create_table_variant")
+    __slots__ = ("hints", "create_table_variant", "data_item_format")
 
-    def __init__(self, hints: TResourceHints, create_table_variant: bool) -> None:
+    def __init__(
+        self,
+        hints: TResourceHints,
+        create_table_variant: bool,
+        data_item_format: TDataItemFormat = None,
+    ) -> None:
         self.hints = hints
         self.create_table_variant = create_table_variant
+        self.data_item_format = data_item_format
 
 
 NATURAL_CALLABLES = ["incremental", "validator", "original_columns"]

dlt/extract/resource.py

@@ -29,6 +29,8 @@
     pipeline_state,
 )
 from dlt.common.utils import flatten_list_or_items, get_callable_name, uniq_id
+from dlt.common.data_writers import TDataItemFormat
+
 from dlt.common.schema.typing import TTableSchema
 from dlt.extract.utils import wrap_async_iterator, wrap_parallel_iterator
 
@@ -72,7 +74,10 @@ def with_table_name(item: TDataItems, table_name: str) -> DataItemWithMeta:
 
 
 def with_hints(
-    item: TDataItems, hints: TResourceHints, create_table_variant: bool = False
+    item: TDataItems,
+    hints: TResourceHints = None,
+    create_table_variant: bool = False,
+    data_item_format: TDataItemFormat = None,
 ) -> DataItemWithMeta:
     """Marks `item` to update the resource with specified `hints`.
 
@@ -81,7 +86,7 @@ def with_hints(
     Create `TResourceHints` with `make_hints`.
     Setting `table_name` will dispatch the `item` to a specified table, like `with_table_name`
     """
-    return DataItemWithMeta(HintsMeta(hints, create_table_variant), item)
+    return DataItemWithMeta(HintsMeta(hints or {}, create_table_variant, data_item_format), item)
 
 
 TDltResourceImpl = TypeVar("TDltResourceImpl", bound="DltResource", default="DltResource")

dlt/extract/storage.py

@@ -50,6 +50,9 @@ def __init__(self, config: NormalizeStorageConfiguration) -> None:
             "arrow": ExtractorItemStorage(
                 self.new_packages, DataWriter.writer_spec_from_file_format("parquet", "arrow")
             ),
+            "model": ExtractorItemStorage(
+                self.new_packages, DataWriter.writer_spec_from_file_format("model", "model")
+            ),
         }
 
     def create_load_package(self, schema: Schema, reuse_exiting_package: bool = True) -> str:

dlt/extract/utils.py

@@ -43,6 +43,8 @@
     SupportsPipe,
 )
 
+from dlt.common.schema.typing import TFileFormat
+
 try:
     from dlt.common.libs import pydantic
 except MissingDependencyException:
@@ -60,14 +62,21 @@
     pandas = None
 
 
-def get_data_item_format(items: TDataItems) -> TDataItemFormat:
+def get_data_item_format(items: TDataItems, meta: Any = None) -> TDataItemFormat:
     """Detect the format of the data item from `items`.
 
     Reverts to `object` for empty lists
 
     Returns:
         The data file format.
     """
+
+    # if incoming item is hints meta, check if item format is forced
+    from dlt.extract.hints import HintsMeta
+
+    if isinstance(meta, HintsMeta) and meta.data_item_format:
+        return meta.data_item_format
+
     if not pyarrow and not pandas:
         return "object"
 

dlt/normalize/worker.py

@@ -111,7 +111,7 @@ def _get_items_normalizer(
             if item_format == "file":
                 # if we want to import file, create a spec that may be used only for importing
                 best_writer_spec = create_import_spec(
-                    parsed_file_name.file_format, items_supported_file_formats  # type: ignore[arg-type]
+                    parsed_file_name.file_format, items_supported_file_formats
                 )
 
             config_loader_file_format = config.loader_file_format

tests/load/test_sql_resource.py

@@ -0,0 +1,50 @@
+# test the sql insert job loader, works only on duckdb for now
+
+from typing import Any
+
+import dlt
+
+from dlt.common.destination.dataset import SupportsReadableDataset
+
+from tests.pipeline.utils import load_table_counts
+
+from dlt.extract.hints import make_hints
+
+
+def test_sql_job() -> None:
+    # populate a table with 10 items and retrieve dataset
+    pipeline = dlt.pipeline(
+        pipeline_name="example_pipeline", destination="duckdb", dataset_name="example_dataset"
+    )
+    pipeline.run([{"a": i} for i in range(10)], table_name="example_table")
+    dataset = pipeline.dataset()
+
+    # create a resource that generates sql statements to create 2 new tables
+    @dlt.resource()
+    def copied_table() -> Any:
+        query = dataset["example_table"].limit(5).query()
+        yield dlt.mark.with_hints(
+            f"CREATE OR REPLACE TABLE copied_table AS {query}",
+            make_hints(file_format="sql"),
+        )
+
+        query = dataset["example_table"].limit(7).query()
+        yield dlt.mark.with_hints(
+            f"CREATE OR REPLACE TABLE copied_table2 AS {query}",
+            make_hints(file_format="sql"),
+        )
+
+    # run sql jobs
+    pipeline.run(copied_table())
+
+    # the two tables where created
+    assert load_table_counts(pipeline, "example_table", "copied_table", "copied_table2") == {
+        "example_table": 10,
+        "copied_table": 5,
+        "copied_table2": 7,
+    }
+
+    # we have a table entry for the main table "copied_table"
+    assert "copied_table" in pipeline.default_schema.tables
+    # but no columns, it's up to the user to provide a schema
+    assert len(pipeline.default_schema.tables["copied_table"]["columns"]) == 0