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 ]
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.
Inherited Members
- sqlmesh.core.engine_adapter.base.EngineAdapter
- EngineAdapter
- recycle
- close
- create_index
- create_table
- create_table_like
- drop_table
- create_view
- create_schema
- drop_schema
- drop_view
- columns
- table_exists
- delete_from
- insert_append
- insert_overwrite_by_time_partition
- update_table
- merge
- rename_table
- fetchone
- fetchall
- fetchdf
- fetch_pyspark_df
- transaction
- execute