Edit on GitHub

sqlmesh.core.engine_adapter

View Source

 1from __future__ import annotations
 2
 3import typing as t
 4
 5from sqlmesh.core.engine_adapter.base import (
 6    EngineAdapter,
 7    EngineAdapterWithIndexSupport,
 8)
 9from sqlmesh.core.engine_adapter.bigquery import BigQueryEngineAdapter
10from sqlmesh.core.engine_adapter.clickhouse import ClickhouseEngineAdapter
11from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
12from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
13from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter
14from sqlmesh.core.engine_adapter.mysql import MySQLEngineAdapter
15from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter
16from sqlmesh.core.engine_adapter.redshift import RedshiftEngineAdapter
17from sqlmesh.core.engine_adapter.snowflake import SnowflakeEngineAdapter
18from sqlmesh.core.engine_adapter.spark import SparkEngineAdapter
19from sqlmesh.core.engine_adapter.starrocks import StarRocksEngineAdapter
20from sqlmesh.core.engine_adapter.trino import TrinoEngineAdapter
21from sqlmesh.core.engine_adapter.athena import AthenaEngineAdapter
22from sqlmesh.core.engine_adapter.risingwave import RisingwaveEngineAdapter
23from sqlmesh.core.engine_adapter.fabric import FabricEngineAdapter
24
25DIALECT_TO_ENGINE_ADAPTER = {
26    "hive": SparkEngineAdapter,
27    "spark": SparkEngineAdapter,
28    "bigquery": BigQueryEngineAdapter,
29    "clickhouse": ClickhouseEngineAdapter,
30    "duckdb": DuckDBEngineAdapter,
31    "snowflake": SnowflakeEngineAdapter,
32    "databricks": DatabricksEngineAdapter,
33    "redshift": RedshiftEngineAdapter,
34    "postgres": PostgresEngineAdapter,
35    "mysql": MySQLEngineAdapter,
36    "mssql": MSSQLEngineAdapter,
37    "trino": TrinoEngineAdapter,
38    "athena": AthenaEngineAdapter,
39    "risingwave": RisingwaveEngineAdapter,
40    "fabric": FabricEngineAdapter,
41    "starrocks": StarRocksEngineAdapter,
42}
43
44DIALECT_ALIASES = {
45    "postgresql": "postgres",
46}
47
48
49def create_engine_adapter(
50    connection_factory: t.Callable[[], t.Any], dialect: str, **kwargs: t.Any
51) -> EngineAdapter:
52    dialect = dialect.lower()
53    dialect = DIALECT_ALIASES.get(dialect, dialect)
54    engine_adapter = DIALECT_TO_ENGINE_ADAPTER.get(dialect)
55    if engine_adapter is None:
56        return EngineAdapter(connection_factory, dialect, **kwargs)
57    if engine_adapter is EngineAdapterWithIndexSupport:
58        return EngineAdapterWithIndexSupport(
59            connection_factory,
60            dialect,
61            **kwargs,
62        )
63    return engine_adapter(connection_factory, **kwargs)

DIALECT_ALIASES = {'postgresql': 'postgres'}

def create_engine_adapter( connection_factory: Callable[[], Any], dialect: str, **kwargs: Any) -> sqlmesh.core.engine_adapter.base.EngineAdapter: View Source

50def create_engine_adapter(
51    connection_factory: t.Callable[[], t.Any], dialect: str, **kwargs: t.Any
52) -> EngineAdapter:
53    dialect = dialect.lower()
54    dialect = DIALECT_ALIASES.get(dialect, dialect)
55    engine_adapter = DIALECT_TO_ENGINE_ADAPTER.get(dialect)
56    if engine_adapter is None:
57        return EngineAdapter(connection_factory, dialect, **kwargs)
58    if engine_adapter is EngineAdapterWithIndexSupport:
59        return EngineAdapterWithIndexSupport(
60            connection_factory,
61            dialect,
62            **kwargs,
63        )
64    return engine_adapter(connection_factory, **kwargs)