Edit on GitHub

sqlmesh.dbt.adapter

  1from __future__ import annotations
  2
  3import abc
  4import typing as t
  5
  6import pandas as pd
  7from sqlglot import exp
  8from sqlglot.helper import seq_get
  9
 10from sqlmesh.core.engine_adapter import EngineAdapter, TransactionType
 11from sqlmesh.utils.errors import ConfigError
 12from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroReference
 13
 14if t.TYPE_CHECKING:
 15    import agate
 16    from dbt.adapters.base import BaseRelation
 17    from dbt.adapters.base.column import Column
 18    from dbt.adapters.base.impl import AdapterResponse
 19
 20
 21class BaseAdapter(abc.ABC):
 22    def __init__(
 23        self,
 24        jinja_macros: JinjaMacroRegistry,
 25        jinja_globals: t.Optional[t.Dict[str, t.Any]] = None,
 26        dialect: str = "",
 27    ):
 28        self.jinja_macros = jinja_macros
 29        self.jinja_globals = jinja_globals or {}
 30        self.dialect = dialect
 31
 32    @abc.abstractmethod
 33    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
 34        """Returns a single relation that matches the provided path."""
 35
 36    @abc.abstractmethod
 37    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
 38        """Gets all relations in a given schema and optionally database.
 39
 40        TODO: Add caching functionality to avoid repeat visits to DB
 41        """
 42
 43    @abc.abstractmethod
 44    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
 45        """Using the engine adapter, gets all the relations that match the given schema grain relation."""
 46
 47    @abc.abstractmethod
 48    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
 49        """Returns the columns for a given table grained relation."""
 50
 51    @abc.abstractmethod
 52    def get_missing_columns(
 53        self, from_relation: BaseRelation, to_relation: BaseRelation
 54    ) -> t.List[Column]:
 55        """Returns the columns in from_relation missing from to_relation."""
 56
 57    @abc.abstractmethod
 58    def create_schema(self, relation: BaseRelation) -> None:
 59        """Creates a schema in the target database."""
 60
 61    @abc.abstractmethod
 62    def drop_schema(self, relation: BaseRelation) -> None:
 63        """Drops a schema in the target database."""
 64
 65    @abc.abstractmethod
 66    def drop_relation(self, relation: BaseRelation) -> None:
 67        """Drops a relation (table) in the target database."""
 68
 69    @abc.abstractmethod
 70    def execute(
 71        self, sql: str, auto_begin: bool = False, fetch: bool = False
 72    ) -> t.Optional[t.Tuple[AdapterResponse, agate.Table]]:
 73        """Executes the given SQL statement and returns the results as an agate table."""
 74
 75    @abc.abstractmethod
 76    def quote(self, identifier: str) -> str:
 77        """Returns a quoted identifeir."""
 78
 79    def dispatch(self, name: str, package: t.Optional[str] = None) -> t.Callable:
 80        """Returns a dialect-specific version of a macro with the given name."""
 81        dialect_name = f"{self.dialect}__{name}"
 82        default_name = f"default__{name}"
 83
 84        references_to_try = [
 85            MacroReference(package=package, name=dialect_name),
 86            MacroReference(package=package, name=default_name),
 87        ]
 88
 89        for reference in references_to_try:
 90            macro_callable = self.jinja_macros.build_macro(
 91                reference, **{**self.jinja_globals, "adapter": self}
 92            )
 93            if macro_callable is not None:
 94                return macro_callable
 95
 96        raise ConfigError(f"Macro '{name}', package '{package}' was not found.")
 97
 98
 99class ParsetimeAdapter(BaseAdapter):
100    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
101        return None
102
103    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
104        return []
105
106    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
107        return []
108
109    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
110        return []
111
112    def get_missing_columns(
113        self, from_relation: BaseRelation, to_relation: BaseRelation
114    ) -> t.List[Column]:
115        return []
116
117    def create_schema(self, relation: BaseRelation) -> None:
118        pass
119
120    def drop_schema(self, relation: BaseRelation) -> None:
121        pass
122
123    def drop_relation(self, relation: BaseRelation) -> None:
124        pass
125
126    def execute(
127        self, sql: str, auto_begin: bool = False, fetch: bool = False
128    ) -> t.Optional[t.Tuple[AdapterResponse, agate.Table]]:
129        return None
130
131    def quote(self, identifier: str) -> str:
132        return identifier
133
134
135class RuntimeAdapter(BaseAdapter):
136    def __init__(
137        self,
138        engine_adapter: EngineAdapter,
139        jinja_macros: JinjaMacroRegistry,
140        jinja_globals: t.Optional[t.Dict[str, t.Any]] = None,
141    ):
142        from dbt.adapters.base.relation import Policy
143
144        super().__init__(jinja_macros, jinja_globals=jinja_globals, dialect=engine_adapter.dialect)
145
146        self.engine_adapter = engine_adapter
147        # All engines quote by default except Snowflake
148        quote_param = engine_adapter.DIALECT != "snowflake"
149        self.quote_policy = Policy(
150            database=quote_param,
151            schema=quote_param,
152            identifier=quote_param,
153        )
154
155    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
156        relations_list = self.list_relations(database, schema)
157        matching_relations = [
158            r
159            for r in relations_list
160            if r.identifier == identifier and r.schema == schema and r.database == database
161        ]
162        return seq_get(matching_relations, 0)
163
164    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
165        from dbt.adapters.base import BaseRelation
166
167        reference_relation = BaseRelation.create(
168            database=database,
169            schema=schema,
170        )
171        return self.list_relations_without_caching(reference_relation)
172
173    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
174        from dbt.adapters.base import BaseRelation
175        from dbt.contracts.relation import RelationType
176
177        assert schema_relation.schema is not None
178        data_objects = self.engine_adapter._get_data_objects(
179            schema_name=schema_relation.schema, catalog_name=schema_relation.database
180        )
181        relations = [
182            BaseRelation.create(
183                database=do.catalog,
184                schema=do.schema_name,
185                identifier=do.name,
186                quote_policy=self.quote_policy,
187                type=RelationType.External if do.type.is_unknown else RelationType(do.type.lower()),
188            )
189            for do in data_objects
190        ]
191        return relations
192
193    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
194        from dbt.adapters.base.column import Column
195
196        return [
197            Column.from_description(name=name, raw_data_type=dtype)
198            for name, dtype in self.engine_adapter.columns(table_name=relation.render()).items()
199        ]
200
201    def get_missing_columns(
202        self, from_relation: BaseRelation, to_relation: BaseRelation
203    ) -> t.List[Column]:
204        target_columns = {col.name for col in self.get_columns_in_relation(to_relation)}
205
206        return [
207            col
208            for col in self.get_columns_in_relation(from_relation)
209            if col.name not in target_columns
210        ]
211
212    def create_schema(self, relation: BaseRelation) -> None:
213        if relation.schema is not None:
214            self.engine_adapter.create_schema(relation.schema)
215
216    def drop_schema(self, relation: BaseRelation) -> None:
217        if relation.schema is not None:
218            self.engine_adapter.drop_schema(relation.schema)
219
220    def drop_relation(self, relation: BaseRelation) -> None:
221        if relation.schema is not None and relation.identifier is not None:
222            self.engine_adapter.drop_table(f"{relation.schema}.{relation.identifier}")
223
224    def execute(
225        self, sql: str, auto_begin: bool = False, fetch: bool = False
226    ) -> t.Tuple[AdapterResponse, agate.Table]:
227        from dbt.adapters.base.impl import AdapterResponse
228        from dbt.clients.agate_helper import empty_table
229
230        from sqlmesh.dbt.util import pandas_to_agate
231
232        # mypy bug: https://github.com/python/mypy/issues/10740
233        exec_func: t.Callable[[str], None | pd.DataFrame] = (
234            self.engine_adapter.fetchdf if fetch else self.engine_adapter.execute  # type: ignore
235        )
236
237        if auto_begin:
238            # TODO: This could be a bug. I think dbt leaves the transaction open while we close immediately.
239            with self.engine_adapter.transaction(TransactionType.DML):
240                resp = exec_func(sql)
241        else:
242            resp = exec_func(sql)
243
244        # TODO: Properly fill in adapter response
245        if fetch:
246            assert isinstance(resp, pd.DataFrame)
247            return AdapterResponse("Success"), pandas_to_agate(resp)
248        return AdapterResponse("Success"), empty_table()
249
250    def quote(self, identifier: str) -> str:
251        return exp.to_column(identifier).sql(dialect=self.engine_adapter.dialect, identify=True)
class BaseAdapter(abc.ABC):
22class BaseAdapter(abc.ABC):
23    def __init__(
24        self,
25        jinja_macros: JinjaMacroRegistry,
26        jinja_globals: t.Optional[t.Dict[str, t.Any]] = None,
27        dialect: str = "",
28    ):
29        self.jinja_macros = jinja_macros
30        self.jinja_globals = jinja_globals or {}
31        self.dialect = dialect
32
33    @abc.abstractmethod
34    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
35        """Returns a single relation that matches the provided path."""
36
37    @abc.abstractmethod
38    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
39        """Gets all relations in a given schema and optionally database.
40
41        TODO: Add caching functionality to avoid repeat visits to DB
42        """
43
44    @abc.abstractmethod
45    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
46        """Using the engine adapter, gets all the relations that match the given schema grain relation."""
47
48    @abc.abstractmethod
49    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
50        """Returns the columns for a given table grained relation."""
51
52    @abc.abstractmethod
53    def get_missing_columns(
54        self, from_relation: BaseRelation, to_relation: BaseRelation
55    ) -> t.List[Column]:
56        """Returns the columns in from_relation missing from to_relation."""
57
58    @abc.abstractmethod
59    def create_schema(self, relation: BaseRelation) -> None:
60        """Creates a schema in the target database."""
61
62    @abc.abstractmethod
63    def drop_schema(self, relation: BaseRelation) -> None:
64        """Drops a schema in the target database."""
65
66    @abc.abstractmethod
67    def drop_relation(self, relation: BaseRelation) -> None:
68        """Drops a relation (table) in the target database."""
69
70    @abc.abstractmethod
71    def execute(
72        self, sql: str, auto_begin: bool = False, fetch: bool = False
73    ) -> t.Optional[t.Tuple[AdapterResponse, agate.Table]]:
74        """Executes the given SQL statement and returns the results as an agate table."""
75
76    @abc.abstractmethod
77    def quote(self, identifier: str) -> str:
78        """Returns a quoted identifeir."""
79
80    def dispatch(self, name: str, package: t.Optional[str] = None) -> t.Callable:
81        """Returns a dialect-specific version of a macro with the given name."""
82        dialect_name = f"{self.dialect}__{name}"
83        default_name = f"default__{name}"
84
85        references_to_try = [
86            MacroReference(package=package, name=dialect_name),
87            MacroReference(package=package, name=default_name),
88        ]
89
90        for reference in references_to_try:
91            macro_callable = self.jinja_macros.build_macro(
92                reference, **{**self.jinja_globals, "adapter": self}
93            )
94            if macro_callable is not None:
95                return macro_callable
96
97        raise ConfigError(f"Macro '{name}', package '{package}' was not found.")

Helper class that provides a standard way to create an ABC using inheritance.

@abc.abstractmethod
def get_relation( self, database: str, schema: str, identifier: str) -> Optional[dbt.adapters.base.relation.BaseRelation]:
33    @abc.abstractmethod
34    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
35        """Returns a single relation that matches the provided path."""

Returns a single relation that matches the provided path.

@abc.abstractmethod
def list_relations( self, database: Optional[str], schema: str) -> List[dbt.adapters.base.relation.BaseRelation]:
37    @abc.abstractmethod
38    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
39        """Gets all relations in a given schema and optionally database.
40
41        TODO: Add caching functionality to avoid repeat visits to DB
42        """

Gets all relations in a given schema and optionally database.

TODO: Add caching functionality to avoid repeat visits to DB

@abc.abstractmethod
def list_relations_without_caching( self, schema_relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.relation.BaseRelation]:
44    @abc.abstractmethod
45    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
46        """Using the engine adapter, gets all the relations that match the given schema grain relation."""

Using the engine adapter, gets all the relations that match the given schema grain relation.

@abc.abstractmethod
def get_columns_in_relation( self, relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.column.Column]:
48    @abc.abstractmethod
49    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
50        """Returns the columns for a given table grained relation."""

Returns the columns for a given table grained relation.

@abc.abstractmethod
def get_missing_columns( self, from_relation: dbt.adapters.base.relation.BaseRelation, to_relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.column.Column]:
52    @abc.abstractmethod
53    def get_missing_columns(
54        self, from_relation: BaseRelation, to_relation: BaseRelation
55    ) -> t.List[Column]:
56        """Returns the columns in from_relation missing from to_relation."""

Returns the columns in from_relation missing from to_relation.

@abc.abstractmethod
def create_schema(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
58    @abc.abstractmethod
59    def create_schema(self, relation: BaseRelation) -> None:
60        """Creates a schema in the target database."""

Creates a schema in the target database.

@abc.abstractmethod
def drop_schema(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
62    @abc.abstractmethod
63    def drop_schema(self, relation: BaseRelation) -> None:
64        """Drops a schema in the target database."""

Drops a schema in the target database.

@abc.abstractmethod
def drop_relation(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
66    @abc.abstractmethod
67    def drop_relation(self, relation: BaseRelation) -> None:
68        """Drops a relation (table) in the target database."""

Drops a relation (table) in the target database.

@abc.abstractmethod
def execute( self, sql: str, auto_begin: bool = False, fetch: bool = False) -> Optional[Tuple[dbt.contracts.connection.AdapterResponse, agate.table.Table]]:
70    @abc.abstractmethod
71    def execute(
72        self, sql: str, auto_begin: bool = False, fetch: bool = False
73    ) -> t.Optional[t.Tuple[AdapterResponse, agate.Table]]:
74        """Executes the given SQL statement and returns the results as an agate table."""

Executes the given SQL statement and returns the results as an agate table.

@abc.abstractmethod
def quote(self, identifier: str) -> str:
76    @abc.abstractmethod
77    def quote(self, identifier: str) -> str:
78        """Returns a quoted identifeir."""

Returns a quoted identifeir.

def dispatch(self, name: str, package: Optional[str] = None) -> Callable:
80    def dispatch(self, name: str, package: t.Optional[str] = None) -> t.Callable:
81        """Returns a dialect-specific version of a macro with the given name."""
82        dialect_name = f"{self.dialect}__{name}"
83        default_name = f"default__{name}"
84
85        references_to_try = [
86            MacroReference(package=package, name=dialect_name),
87            MacroReference(package=package, name=default_name),
88        ]
89
90        for reference in references_to_try:
91            macro_callable = self.jinja_macros.build_macro(
92                reference, **{**self.jinja_globals, "adapter": self}
93            )
94            if macro_callable is not None:
95                return macro_callable
96
97        raise ConfigError(f"Macro '{name}', package '{package}' was not found.")

Returns a dialect-specific version of a macro with the given name.

class ParsetimeAdapter(BaseAdapter):
100class ParsetimeAdapter(BaseAdapter):
101    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
102        return None
103
104    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
105        return []
106
107    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
108        return []
109
110    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
111        return []
112
113    def get_missing_columns(
114        self, from_relation: BaseRelation, to_relation: BaseRelation
115    ) -> t.List[Column]:
116        return []
117
118    def create_schema(self, relation: BaseRelation) -> None:
119        pass
120
121    def drop_schema(self, relation: BaseRelation) -> None:
122        pass
123
124    def drop_relation(self, relation: BaseRelation) -> None:
125        pass
126
127    def execute(
128        self, sql: str, auto_begin: bool = False, fetch: bool = False
129    ) -> t.Optional[t.Tuple[AdapterResponse, agate.Table]]:
130        return None
131
132    def quote(self, identifier: str) -> str:
133        return identifier

Helper class that provides a standard way to create an ABC using inheritance.

def get_relation( self, database: str, schema: str, identifier: str) -> Optional[dbt.adapters.base.relation.BaseRelation]:
101    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
102        return None

Returns a single relation that matches the provided path.

def list_relations( self, database: Optional[str], schema: str) -> List[dbt.adapters.base.relation.BaseRelation]:
104    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
105        return []

Gets all relations in a given schema and optionally database.

TODO: Add caching functionality to avoid repeat visits to DB

def list_relations_without_caching( self, schema_relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.relation.BaseRelation]:
107    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
108        return []

Using the engine adapter, gets all the relations that match the given schema grain relation.

def get_columns_in_relation( self, relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.column.Column]:
110    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
111        return []

Returns the columns for a given table grained relation.

def get_missing_columns( self, from_relation: dbt.adapters.base.relation.BaseRelation, to_relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.column.Column]:
113    def get_missing_columns(
114        self, from_relation: BaseRelation, to_relation: BaseRelation
115    ) -> t.List[Column]:
116        return []

Returns the columns in from_relation missing from to_relation.

def create_schema(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
118    def create_schema(self, relation: BaseRelation) -> None:
119        pass

Creates a schema in the target database.

def drop_schema(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
121    def drop_schema(self, relation: BaseRelation) -> None:
122        pass

Drops a schema in the target database.

def drop_relation(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
124    def drop_relation(self, relation: BaseRelation) -> None:
125        pass

Drops a relation (table) in the target database.

def execute( self, sql: str, auto_begin: bool = False, fetch: bool = False) -> Optional[Tuple[dbt.contracts.connection.AdapterResponse, agate.table.Table]]:
127    def execute(
128        self, sql: str, auto_begin: bool = False, fetch: bool = False
129    ) -> t.Optional[t.Tuple[AdapterResponse, agate.Table]]:
130        return None

Executes the given SQL statement and returns the results as an agate table.

def quote(self, identifier: str) -> str:
132    def quote(self, identifier: str) -> str:
133        return identifier

Returns a quoted identifeir.

Inherited Members
BaseAdapter
BaseAdapter
dispatch
class RuntimeAdapter(BaseAdapter):
136class RuntimeAdapter(BaseAdapter):
137    def __init__(
138        self,
139        engine_adapter: EngineAdapter,
140        jinja_macros: JinjaMacroRegistry,
141        jinja_globals: t.Optional[t.Dict[str, t.Any]] = None,
142    ):
143        from dbt.adapters.base.relation import Policy
144
145        super().__init__(jinja_macros, jinja_globals=jinja_globals, dialect=engine_adapter.dialect)
146
147        self.engine_adapter = engine_adapter
148        # All engines quote by default except Snowflake
149        quote_param = engine_adapter.DIALECT != "snowflake"
150        self.quote_policy = Policy(
151            database=quote_param,
152            schema=quote_param,
153            identifier=quote_param,
154        )
155
156    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
157        relations_list = self.list_relations(database, schema)
158        matching_relations = [
159            r
160            for r in relations_list
161            if r.identifier == identifier and r.schema == schema and r.database == database
162        ]
163        return seq_get(matching_relations, 0)
164
165    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
166        from dbt.adapters.base import BaseRelation
167
168        reference_relation = BaseRelation.create(
169            database=database,
170            schema=schema,
171        )
172        return self.list_relations_without_caching(reference_relation)
173
174    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
175        from dbt.adapters.base import BaseRelation
176        from dbt.contracts.relation import RelationType
177
178        assert schema_relation.schema is not None
179        data_objects = self.engine_adapter._get_data_objects(
180            schema_name=schema_relation.schema, catalog_name=schema_relation.database
181        )
182        relations = [
183            BaseRelation.create(
184                database=do.catalog,
185                schema=do.schema_name,
186                identifier=do.name,
187                quote_policy=self.quote_policy,
188                type=RelationType.External if do.type.is_unknown else RelationType(do.type.lower()),
189            )
190            for do in data_objects
191        ]
192        return relations
193
194    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
195        from dbt.adapters.base.column import Column
196
197        return [
198            Column.from_description(name=name, raw_data_type=dtype)
199            for name, dtype in self.engine_adapter.columns(table_name=relation.render()).items()
200        ]
201
202    def get_missing_columns(
203        self, from_relation: BaseRelation, to_relation: BaseRelation
204    ) -> t.List[Column]:
205        target_columns = {col.name for col in self.get_columns_in_relation(to_relation)}
206
207        return [
208            col
209            for col in self.get_columns_in_relation(from_relation)
210            if col.name not in target_columns
211        ]
212
213    def create_schema(self, relation: BaseRelation) -> None:
214        if relation.schema is not None:
215            self.engine_adapter.create_schema(relation.schema)
216
217    def drop_schema(self, relation: BaseRelation) -> None:
218        if relation.schema is not None:
219            self.engine_adapter.drop_schema(relation.schema)
220
221    def drop_relation(self, relation: BaseRelation) -> None:
222        if relation.schema is not None and relation.identifier is not None:
223            self.engine_adapter.drop_table(f"{relation.schema}.{relation.identifier}")
224
225    def execute(
226        self, sql: str, auto_begin: bool = False, fetch: bool = False
227    ) -> t.Tuple[AdapterResponse, agate.Table]:
228        from dbt.adapters.base.impl import AdapterResponse
229        from dbt.clients.agate_helper import empty_table
230
231        from sqlmesh.dbt.util import pandas_to_agate
232
233        # mypy bug: https://github.com/python/mypy/issues/10740
234        exec_func: t.Callable[[str], None | pd.DataFrame] = (
235            self.engine_adapter.fetchdf if fetch else self.engine_adapter.execute  # type: ignore
236        )
237
238        if auto_begin:
239            # TODO: This could be a bug. I think dbt leaves the transaction open while we close immediately.
240            with self.engine_adapter.transaction(TransactionType.DML):
241                resp = exec_func(sql)
242        else:
243            resp = exec_func(sql)
244
245        # TODO: Properly fill in adapter response
246        if fetch:
247            assert isinstance(resp, pd.DataFrame)
248            return AdapterResponse("Success"), pandas_to_agate(resp)
249        return AdapterResponse("Success"), empty_table()
250
251    def quote(self, identifier: str) -> str:
252        return exp.to_column(identifier).sql(dialect=self.engine_adapter.dialect, identify=True)

Helper class that provides a standard way to create an ABC using inheritance.

RuntimeAdapter( engine_adapter: sqlmesh.core.engine_adapter.base.EngineAdapter, jinja_macros: sqlmesh.utils.jinja.JinjaMacroRegistry, jinja_globals: Optional[Dict[str, Any]] = None)
137    def __init__(
138        self,
139        engine_adapter: EngineAdapter,
140        jinja_macros: JinjaMacroRegistry,
141        jinja_globals: t.Optional[t.Dict[str, t.Any]] = None,
142    ):
143        from dbt.adapters.base.relation import Policy
144
145        super().__init__(jinja_macros, jinja_globals=jinja_globals, dialect=engine_adapter.dialect)
146
147        self.engine_adapter = engine_adapter
148        # All engines quote by default except Snowflake
149        quote_param = engine_adapter.DIALECT != "snowflake"
150        self.quote_policy = Policy(
151            database=quote_param,
152            schema=quote_param,
153            identifier=quote_param,
154        )
def get_relation( self, database: str, schema: str, identifier: str) -> Optional[dbt.adapters.base.relation.BaseRelation]:
156    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
157        relations_list = self.list_relations(database, schema)
158        matching_relations = [
159            r
160            for r in relations_list
161            if r.identifier == identifier and r.schema == schema and r.database == database
162        ]
163        return seq_get(matching_relations, 0)

Returns a single relation that matches the provided path.

def list_relations( self, database: Optional[str], schema: str) -> List[dbt.adapters.base.relation.BaseRelation]:
165    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
166        from dbt.adapters.base import BaseRelation
167
168        reference_relation = BaseRelation.create(
169            database=database,
170            schema=schema,
171        )
172        return self.list_relations_without_caching(reference_relation)

Gets all relations in a given schema and optionally database.

TODO: Add caching functionality to avoid repeat visits to DB

def list_relations_without_caching( self, schema_relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.relation.BaseRelation]:
174    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
175        from dbt.adapters.base import BaseRelation
176        from dbt.contracts.relation import RelationType
177
178        assert schema_relation.schema is not None
179        data_objects = self.engine_adapter._get_data_objects(
180            schema_name=schema_relation.schema, catalog_name=schema_relation.database
181        )
182        relations = [
183            BaseRelation.create(
184                database=do.catalog,
185                schema=do.schema_name,
186                identifier=do.name,
187                quote_policy=self.quote_policy,
188                type=RelationType.External if do.type.is_unknown else RelationType(do.type.lower()),
189            )
190            for do in data_objects
191        ]
192        return relations

Using the engine adapter, gets all the relations that match the given schema grain relation.

def get_columns_in_relation( self, relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.column.Column]:
194    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
195        from dbt.adapters.base.column import Column
196
197        return [
198            Column.from_description(name=name, raw_data_type=dtype)
199            for name, dtype in self.engine_adapter.columns(table_name=relation.render()).items()
200        ]

Returns the columns for a given table grained relation.

def get_missing_columns( self, from_relation: dbt.adapters.base.relation.BaseRelation, to_relation: dbt.adapters.base.relation.BaseRelation) -> List[dbt.adapters.base.column.Column]:
202    def get_missing_columns(
203        self, from_relation: BaseRelation, to_relation: BaseRelation
204    ) -> t.List[Column]:
205        target_columns = {col.name for col in self.get_columns_in_relation(to_relation)}
206
207        return [
208            col
209            for col in self.get_columns_in_relation(from_relation)
210            if col.name not in target_columns
211        ]

Returns the columns in from_relation missing from to_relation.

def create_schema(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
213    def create_schema(self, relation: BaseRelation) -> None:
214        if relation.schema is not None:
215            self.engine_adapter.create_schema(relation.schema)

Creates a schema in the target database.

def drop_schema(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
217    def drop_schema(self, relation: BaseRelation) -> None:
218        if relation.schema is not None:
219            self.engine_adapter.drop_schema(relation.schema)

Drops a schema in the target database.

def drop_relation(self, relation: dbt.adapters.base.relation.BaseRelation) -> None:
221    def drop_relation(self, relation: BaseRelation) -> None:
222        if relation.schema is not None and relation.identifier is not None:
223            self.engine_adapter.drop_table(f"{relation.schema}.{relation.identifier}")

Drops a relation (table) in the target database.

def execute( self, sql: str, auto_begin: bool = False, fetch: bool = False) -> Tuple[dbt.contracts.connection.AdapterResponse, agate.table.Table]:
225    def execute(
226        self, sql: str, auto_begin: bool = False, fetch: bool = False
227    ) -> t.Tuple[AdapterResponse, agate.Table]:
228        from dbt.adapters.base.impl import AdapterResponse
229        from dbt.clients.agate_helper import empty_table
230
231        from sqlmesh.dbt.util import pandas_to_agate
232
233        # mypy bug: https://github.com/python/mypy/issues/10740
234        exec_func: t.Callable[[str], None | pd.DataFrame] = (
235            self.engine_adapter.fetchdf if fetch else self.engine_adapter.execute  # type: ignore
236        )
237
238        if auto_begin:
239            # TODO: This could be a bug. I think dbt leaves the transaction open while we close immediately.
240            with self.engine_adapter.transaction(TransactionType.DML):
241                resp = exec_func(sql)
242        else:
243            resp = exec_func(sql)
244
245        # TODO: Properly fill in adapter response
246        if fetch:
247            assert isinstance(resp, pd.DataFrame)
248            return AdapterResponse("Success"), pandas_to_agate(resp)
249        return AdapterResponse("Success"), empty_table()

Executes the given SQL statement and returns the results as an agate table.

def quote(self, identifier: str) -> str:
251    def quote(self, identifier: str) -> str:
252        return exp.to_column(identifier).sql(dialect=self.engine_adapter.dialect, identify=True)

Returns a quoted identifeir.

Inherited Members
BaseAdapter
dispatch