Fix!: Use target dbt adapter relation and column classes. (#1056)

Author: crericha

setup.cfg

@@ -75,6 +75,9 @@ ignore_missing_imports = True
 [mypy-pytest_lazyfixture.*]
 ignore_missing_imports = True
 
+[mypy-dbt.adapters.*]
+ignore_missing_imports = True
+
 [autoflake]
 in-place = True
 expand-star-imports = True

sqlmesh/dbt/builtin.py

@@ -8,12 +8,13 @@
 import agate
 import jinja2
 from dbt import version
-from dbt.adapters.base import BaseRelation
+from dbt.adapters.base import BaseRelation, Column
 from dbt.contracts.relation import Policy
 from ruamel.yaml import YAMLError
 
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.dbt.adapter import BaseAdapter, ParsetimeAdapter, RuntimeAdapter
+from sqlmesh.dbt.target import TargetConfig
 from sqlmesh.dbt.util import DBT_VERSION
 from sqlmesh.utils import AttributeDict, yaml
 from sqlmesh.utils.errors import ConfigError, MacroEvalError
@@ -37,12 +38,14 @@ def warn(self, msg: str) -> str:
 
 
 class Api:
-    def __init__(self) -> None:
-        from dbt.adapters.base.column import Column
-        from dbt.adapters.base.relation import BaseRelation
-
-        self.Relation = BaseRelation
-        self.Column = Column
+    def __init__(self, target: t.Optional[AttributeDict] = None) -> None:
+        if target:
+            config = TargetConfig.load(target)
+            self.Relation = config.relation_class
+            self.Column = config.column_class
+        else:
+            self.Relation = BaseRelation
+            self.Column = Column
 
 
 class Flags:
@@ -156,25 +159,25 @@ def var(name: str, default: t.Optional[t.Any] = None) -> t.Any:
     return var
 
 
-def generate_ref(refs: t.Dict[str, t.Any]) -> t.Callable:
+def generate_ref(refs: t.Dict[str, t.Any], api: Api) -> t.Callable:
     def ref(package: str, name: t.Optional[str] = None) -> t.Optional[BaseRelation]:
         ref_name = f"{package}.{name}" if name else package
         relation_info = refs.get(ref_name)
         if relation_info is None:
             return None
 
-        return BaseRelation.create(**relation_info)
+        return api.Relation.create(**relation_info)
 
     return ref
 
 
-def generate_source(sources: t.Dict[str, t.Any]) -> t.Callable:
+def generate_source(sources: t.Dict[str, t.Any], api: Api) -> t.Callable:
     def source(package: str, name: str) -> t.Optional[BaseRelation]:
         relation_info = sources.get(f"{package}.{name}")
         if relation_info is None:
             return None
 
-        return BaseRelation.create(**relation_info)
+        return api.Relation.create(**relation_info)
 
     return source
 
@@ -251,7 +254,6 @@ def _try_literal_eval(value: str) -> t.Any:
 
 
 BUILTIN_GLOBALS = {
-    "api": Api(),
     "dbt_version": version.__version__,
     "env_var": env_var,
     "exceptions": Exceptions(),
@@ -288,20 +290,25 @@ def create_builtin_globals(
     builtin_globals = BUILTIN_GLOBALS.copy()
     jinja_globals = jinja_globals.copy()
 
+    target: t.Optional[AttributeDict] = jinja_globals.get("target", None)
+    api = Api(target)
+
+    builtin_globals["api"] = api
+
     this = jinja_globals.pop("this", None)
     if this is not None:
-        if not isinstance(this, BaseRelation):
-            builtin_globals["this"] = BaseRelation.create(**this)
+        if not isinstance(this, api.Relation):
+            builtin_globals["this"] = api.Relation.create(**this)
         else:
             builtin_globals["this"] = this
 
     sources = jinja_globals.pop("sources", None)
     if sources is not None:
-        builtin_globals["source"] = generate_source(sources)
+        builtin_globals["source"] = generate_source(sources, api)
 
     refs = jinja_globals.pop("refs", None)
     if refs is not None:
-        builtin_globals["ref"] = generate_ref(refs)
+        builtin_globals["ref"] = generate_ref(refs, api)
 
     variables = jinja_globals.pop("vars", None)
     if variables is not None:
@@ -319,11 +326,10 @@ def create_builtin_globals(
         )
         builtin_globals.update({"log": log, "print": log})
     else:
-        target = jinja_globals.get("target")
         adapter = ParsetimeAdapter(
             jinja_macros,
             jinja_globals={**builtin_globals, **jinja_globals},
-            dialect=target.type if target else None,
+            dialect=target.type if target else None,  # type: ignore
         )
         builtin_globals.update({"log": no_log, "print": no_log})
 

sqlmesh/dbt/target.py

@@ -4,6 +4,7 @@
 import sys
 import typing as t
 
+from dbt.adapters.base import BaseRelation, Column
 from pydantic import Field, root_validator, validator
 
 from sqlmesh.core.config.connection import (
@@ -98,6 +99,14 @@ def quoting(self) -> QuotingConfig:
     def extra(self) -> t.Set[str]:
         return self.extra_fields(set(self.dict()))
 
+    @property
+    def relation_class(self) -> t.Type[BaseRelation]:
+        return BaseRelation
+
+    @property
+    def column_class(self) -> t.Type[Column]:
+        return Column
+
 
 class DuckDbConfig(TargetConfig):
     """
@@ -115,6 +124,12 @@ class DuckDbConfig(TargetConfig):
     def default_incremental_strategy(self, kind: IncrementalKind) -> str:
         return "delete+insert"
 
+    @property
+    def relation_class(self) -> t.Type[BaseRelation]:
+        from dbt.adapters.duckdb.relation import DuckDBRelation
+
+        return DuckDBRelation
+
     def to_sqlmesh(self) -> ConnectionConfig:
         return DuckDBConnectionConfig(database=self.path, concurrent_tasks=self.threads)
 
@@ -154,6 +169,18 @@ class SnowflakeConfig(TargetConfig):
     def default_incremental_strategy(self, kind: IncrementalKind) -> str:
         return "merge"
 
+    @property
+    def relation_class(self) -> t.Type[BaseRelation]:
+        from dbt.adapters.snowflake import SnowflakeRelation
+
+        return SnowflakeRelation
+
+    @property
+    def column_class(self) -> t.Type[Column]:
+        from dbt.adapters.snowflake import SnowflakeColumn
+
+        return SnowflakeColumn
+
     def to_sqlmesh(self) -> ConnectionConfig:
         return SnowflakeConnectionConfig(
             user=self.user,
@@ -274,6 +301,18 @@ def validate_database(
     def default_incremental_strategy(self, kind: IncrementalKind) -> str:
         return "append"
 
+    @property
+    def relation_class(self) -> t.Type[BaseRelation]:
+        from dbt.adapters.redshift import RedshiftRelation
+
+        return RedshiftRelation
+
+    @property
+    def column_class(self) -> t.Type[Column]:
+        from dbt.adapters.redshift import RedshiftColumn
+
+        return RedshiftColumn
+
     def to_sqlmesh(self) -> ConnectionConfig:
         return RedshiftConnectionConfig(
             user=self.user,
@@ -309,6 +348,18 @@ class DatabricksConfig(TargetConfig):
     def default_incremental_strategy(self, kind: IncrementalKind) -> str:
         return "merge"
 
+    @property
+    def relation_class(self) -> t.Type[BaseRelation]:
+        from dbt.adapters.databricks.relation import DatabricksRelation
+
+        return DatabricksRelation
+
+    @property
+    def column_class(self) -> t.Type[Column]:
+        from dbt.adapters.databricks.column import DatabricksColumn
+
+        return DatabricksColumn
+
     def to_sqlmesh(self) -> ConnectionConfig:
         return DatabricksConnectionConfig(
             server_hostname=self.host,
@@ -376,6 +427,18 @@ def validate_fields(
     def default_incremental_strategy(self, kind: IncrementalKind) -> str:
         return "merge"
 
+    @property
+    def relation_class(self) -> t.Type[BaseRelation]:
+        from dbt.adapters.bigquery.relation import BigQueryRelation
+
+        return BigQueryRelation
+
+    @property
+    def column_class(self) -> t.Type[Column]:
+        from dbt.adapters.bigquery import BigQueryColumn
+
+        return BigQueryColumn
+
     def to_sqlmesh(self) -> ConnectionConfig:
         return BigQueryConnectionConfig(
             method=self.method,

tests/dbt/test_transformation.py

@@ -315,6 +315,11 @@ def test_flags(sushi_test_project: Project):
 def test_relation(sushi_test_project: Project):
     context = sushi_test_project.context
 
+    assert (
+        context.render("{{ api.Relation }}")
+        == "<class 'dbt.adapters.duckdb.relation.DuckDBRelation'>"
+    )
+
     jinja = (
         "{% set relation = api.Relation.create(schema='sushi', identifier='waiters') %}"
         "{{ relation.schema }} {{ relation.identifier}}"
@@ -326,6 +331,8 @@ def test_relation(sushi_test_project: Project):
 def test_column(sushi_test_project: Project):
     context = sushi_test_project.context
 
+    assert context.render("{{ api.Column }}") == "<class 'dbt.adapters.base.column.Column'>"
+
     jinja = (
         "{% set col = api.Column('foo', 'integer') %}" "{{ col.is_integer() }} {{ col.is_string()}}"
     )