sqlmesh.core.engine_adapter.databricks_api API documentation

Edit on GitHub

# sqlmesh.core.engine_adapter.databricks_api

View Source

     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): View Source

    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
sqlmesh.core.engine_adapter.base_spark.BaseSparkEngineAdapter
  replace_query
  create_state_table
  alter_table
  supports_transactions