add materializable sqlmodel (wip)

Author: sh-rp

dlt/destinations/dataset/dataset.py

@@ -90,8 +90,13 @@ def schema(self) -> Schema:
     def sqlglot_schema(self) -> SQLGlotSchema:
         # NOTE: no cache for now, it is probably more expensive to compute the current schema hash
         # to see wether this is stale than to compute a new sqlglot schema
-        dialect: str = self.sql_client.capabilities.sqlglot_dialect
-        return lineage.create_sqlglot_schema(self.schema, self.dataset_name, dialect=dialect)
+        return lineage.create_sqlglot_schema(
+            self.schema, self.dataset_name, dialect=self.sqlglot_dialect
+        )
+
+    @property
+    def sqlglot_dialect(self) -> str:
+        return self.sql_client.capabilities.sqlglot_dialect
 
     @property
     def dataset_name(self) -> str:

dlt/destinations/dataset/relation.py

@@ -79,9 +79,10 @@ def query(self, pretty: bool = False) -> str:
                 "Must be an SQL SELECT statement."
             )
 
-        return query.sql(
-            dialect=self._dataset.sql_client.capabilities.sqlglot_dialect, pretty=pretty
-        )
+        return query.sql(dialect=self._dataset.sqlglot_dialect, pretty=pretty)
+
+    def query_dialect(self) -> str:
+        return self._dataset.sqlglot_dialect
 
     def _query(self) -> sge.Query:
         """Returns a compliant with dlt schema in the relation.
@@ -249,16 +250,11 @@ def __init__(
         self._selected_columns = selected_columns
 
     def _query(self) -> sge.Query:
-        destination_dialect = self._dataset.sql_client.capabilities.sqlglot_dialect
-
         # TODO reimplement this using SQLGLot instead of passing strings
         if self._provided_query:
             return cast(
                 sge.Query,
-                sqlglot.parse_one(
-                    self._provided_query,
-                    dialect=self._provided_query_dialect or destination_dialect,
-                ),
+                sqlglot.parse_one(self._provided_query, dialect=self.query_dialect()),
             )
 
         dataset_schema = self._dataset.schema
@@ -285,11 +281,14 @@ def _query(self) -> sge.Query:
             ),
         )
 
+    def query_dialect(self) -> str:
+        return self._provided_query_dialect or self._dataset.sqlglot_dialect
+
     def __copy__(self) -> Self:
         return self.__class__(
             readable_dataset=self._dataset,
             provided_query=self._provided_query,
-            provided_query_dialect=self._provided_query_dialect,
+            provided_query_dialect=self.query_dialect(),
             table_name=self._table_name,
             limit=self._limit,
             selected_columns=self._selected_columns,

dlt/extract/hints.py

@@ -55,7 +55,6 @@
 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
 
 import sqlglot
 
@@ -99,9 +98,17 @@ def __init__(
         self.create_table_variant = create_table_variant
 
 
-class SqlModel(NamedTuple):
-    query: str
-    dialect: Optional[str] = None
+class SqlModel:
+    """
+    A SqlModel is a named tuple that contains a query and a dialect.
+    It is used to represent a SQL query and the dialect to use for parsing it.
+    """
+
+    __slots__ = ("query", "dialect")
+
+    def __init__(self, query: str, dialect: Optional[str] = None) -> None:
+        self.query = query
+        self.dialect = dialect
 
     @classmethod
     def from_query_string(cls, query: str, dialect: Optional[str] = None) -> "SqlModel":

dlt/transformations/transformation.py

@@ -29,12 +29,35 @@
     TTableFormat,
     TTableReferenceParam,
 )
-from dlt.common.destination.dataset import SupportsReadableRelation
 from dlt.transformations.configuration import TransformationConfiguration
 from dlt.common.utils import get_callable_name
 from dlt.extract.exceptions import CurrentSourceNotAvailable
 
 
+class MaterializableSqlModel(SqlModel):
+    # NOTE: we could forward all data access methods to this class
+    __slots__ = ("relation",)
+
+    def __init__(
+        self,
+        relation: Optional[BaseReadableDBAPIRelation] = None,
+    ) -> None:
+        super().__init__(relation.query(), relation.query_dialect())
+        self.relation = relation
+
+    @classmethod
+    def from_relation(cls, relation: BaseReadableDBAPIRelation) -> "MaterializableSqlModel":
+        return cls(relation=relation)
+
+    def compute_columns(self) -> TTableSchemaColumns:
+        computed_columns, _ = self.relation._compute_columns_schema(
+            infer_sqlglot_schema=True,
+            allow_anonymous_columns=True,
+            allow_partial=True,
+        )
+        return computed_columns
+
+
 class DltTransformationResource(DltResource):
     def __init__(self, *args: Any, **kwds: Any) -> None:
         super().__init__(*args, **kwds)
@@ -113,11 +136,12 @@ def transformation_function(*args: Any, **kwargs: Any) -> Iterator[TDataItems]:
         # Call the transformation function
         transformation_result: Any = func(*args, **kwargs)
 
-        # If a string is returned, treat it as a SQL query
-        if isinstance(transformation_result, str):
-            transformation_result = datasets[0](transformation_result)
-
-        if not isinstance(transformation_result, BaseReadableDBAPIRelation):
+        # If a string is returned, construct relation from first dataset from it
+        if isinstance(transformation_result, BaseReadableDBAPIRelation):
+            relation = transformation_result
+        elif isinstance(transformation_result, str):
+            relation = datasets[0](transformation_result)
+        else:
             raise TransformationInvalidReturnTypeException(
                 resource_name,
                 "Sql Transformation %s returned an invalid type: %s. Please either return a valid"
@@ -126,24 +150,14 @@ def transformation_function(*args: Any, **kwargs: Any) -> Iterator[TDataItems]:
                 % (name, type(transformation_result)),
             )
 
-        # Compute columns schema
-        computed_columns, _ = transformation_result._compute_columns_schema(
-            infer_sqlglot_schema=True,
-            allow_anonymous_columns=True,
-            allow_partial=True,
-        )
-        select_dialect = datasets[0].sql_client.capabilities.sqlglot_dialect
-        select_query = transformation_result.query()
-        all_columns = {**computed_columns, **(columns or {})}
+        sql_model = MaterializableSqlModel.from_relation(relation)
 
         if not should_materialize:
-            yield dlt.mark.with_hints(
-                SqlModel(select_query, dialect=select_dialect),
-                hints=make_hints(columns=all_columns),
-            )
+            yield sql_model
         else:
-            for chunk in datasets[0](select_query).iter_arrow(chunk_size=config.buffer_max_items):
-                yield dlt.mark.with_hints(chunk, hints=make_hints(columns=all_columns))
+            column_hints = make_hints(columns=sql_model.compute_columns())
+            for chunk in relation.iter_arrow(chunk_size=config.buffer_max_items):
+                yield dlt.mark.with_hints(chunk, hints=column_hints)
 
     return dlt.resource(  # type: ignore[return-value]
         name=name,