Edit on GitHub

sqlmesh.core.engine_adapter.databricks_api

 1from __future__ import annotations
 2
 3import typing as t
 4
 5from sqlglot import exp
 6
 7from sqlmesh.core.engine_adapter.base_spark import BaseSparkEngineAdapter
 8from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 9
10if t.TYPE_CHECKING:
11    from sqlmesh.core.engine_adapter._typing import DF
12
13
14class DatabricksSQLEngineAdapter(BaseSparkEngineAdapter):
15    DIALECT = "databricks"
16
17    def _fetch_native_df(self, query: t.Union[exp.Expression, str]) -> DF:
18        """
19        Returns a Pandas DataFrame from a query or expression.
20        """
21        self.execute(query)
22        return self.cursor.fetchall_arrow().to_pandas()
23
24    def _get_data_objects(
25        self, schema_name: str, catalog_name: t.Optional[str] = None
26    ) -> t.List[DataObject]:
27        """
28        Returns all the data objects that exist in the given schema and optionally catalog.
29        """
30        tables = [x.asDict() for x in self.cursor.tables(catalog_name, schema_name)]
31        if any(not row["TABLE_TYPE"] for row in tables):
32            if catalog_name:
33                self.execute(f"USE CATALOG {catalog_name}")
34            view_names = [
35                row.viewName for row in self.fetchdf(f"SHOW VIEWS FROM {schema_name}").itertuples()  # type: ignore
36            ]
37            for table in tables:
38                if not table["TABLE_TYPE"]:
39                    table["TABLE_TYPE"] = "VIEW" if table["TABLE_NAME"] in view_names else "TABLE"
40        return [
41            DataObject(
42                catalog=row["TABLE_CAT"],
43                schema=row["TABLE_SCHEM"],
44                name=row["TABLE_NAME"],
45                type=DataObjectType.from_str(row["TABLE_TYPE"]),
46            )
47            for row in tables
48        ]
class DatabricksSQLEngineAdapter(sqlmesh.core.engine_adapter.base_spark.BaseSparkEngineAdapter):
15class DatabricksSQLEngineAdapter(BaseSparkEngineAdapter):
16    DIALECT = "databricks"
17
18    def _fetch_native_df(self, query: t.Union[exp.Expression, str]) -> DF:
19        """
20        Returns a Pandas DataFrame from a query or expression.
21        """
22        self.execute(query)
23        return self.cursor.fetchall_arrow().to_pandas()
24
25    def _get_data_objects(
26        self, schema_name: str, catalog_name: t.Optional[str] = None
27    ) -> t.List[DataObject]:
28        """
29        Returns all the data objects that exist in the given schema and optionally catalog.
30        """
31        tables = [x.asDict() for x in self.cursor.tables(catalog_name, schema_name)]
32        if any(not row["TABLE_TYPE"] for row in tables):
33            if catalog_name:
34                self.execute(f"USE CATALOG {catalog_name}")
35            view_names = [
36                row.viewName for row in self.fetchdf(f"SHOW VIEWS FROM {schema_name}").itertuples()  # type: ignore
37            ]
38            for table in tables:
39                if not table["TABLE_TYPE"]:
40                    table["TABLE_TYPE"] = "VIEW" if table["TABLE_NAME"] in view_names else "TABLE"
41        return [
42            DataObject(
43                catalog=row["TABLE_CAT"],
44                schema=row["TABLE_SCHEM"],
45                name=row["TABLE_NAME"],
46                type=DataObjectType.from_str(row["TABLE_TYPE"]),
47            )
48            for row in tables
49        ]

Base class wrapping a Database API compliant connection.

The EngineAdapter is an easily-subclassable interface that interacts with the underlying engine and data store.

Arguments:
  • connection_factory: a callable which produces a new Database API-compliant connection on every call.
  • dialect: The dialect with which this adapter is associated.
  • multithreaded: Indicates whether this adapter will be used by more than one thread.