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.trino import TrinoEngineAdapter
20from sqlmesh.core.engine_adapter.athena import AthenaEngineAdapter
21from sqlmesh.core.engine_adapter.risingwave import RisingwaveEngineAdapter
22from sqlmesh.core.engine_adapter.fabric import FabricEngineAdapter
23
24DIALECT_TO_ENGINE_ADAPTER = {
25    "hive": SparkEngineAdapter,
26    "spark": SparkEngineAdapter,
27    "bigquery": BigQueryEngineAdapter,
28    "clickhouse": ClickhouseEngineAdapter,
29    "duckdb": DuckDBEngineAdapter,
30    "snowflake": SnowflakeEngineAdapter,
31    "databricks": DatabricksEngineAdapter,
32    "redshift": RedshiftEngineAdapter,
33    "postgres": PostgresEngineAdapter,
34    "mysql": MySQLEngineAdapter,
35    "mssql": MSSQLEngineAdapter,
36    "trino": TrinoEngineAdapter,
37    "athena": AthenaEngineAdapter,
38    "risingwave": RisingwaveEngineAdapter,
39    "fabric": FabricEngineAdapter,
40}
41
42DIALECT_ALIASES = {
43    "postgresql": "postgres",
44}
45
46
47def create_engine_adapter(
48    connection_factory: t.Callable[[], t.Any], dialect: str, **kwargs: t.Any
49) -> EngineAdapter:
50    dialect = dialect.lower()
51    dialect = DIALECT_ALIASES.get(dialect, dialect)
52    engine_adapter = DIALECT_TO_ENGINE_ADAPTER.get(dialect)
53    if engine_adapter is None:
54        return EngineAdapter(connection_factory, dialect, **kwargs)
55    if engine_adapter is EngineAdapterWithIndexSupport:
56        return EngineAdapterWithIndexSupport(
57            connection_factory,
58            dialect,
59            **kwargs,
60        )
61    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

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