sqlmesh 0.184.3.dev3__tar.gz → 0.184.3.dev5__tar.gz

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.184.3.dev3
+Version: 0.184.3.dev5
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.184.3.dev3'
-__version_tuple__ = version_tuple = (0, 184, 3, 'dev3')
+__version__ = version = '0.184.3.dev5'
+__version_tuple__ = version_tuple = (0, 184, 3, 'dev5')

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/sqlmesh/core/config/connection.py

@@ -4,6 +4,7 @@ import abc
 import base64
 import logging
 import os
+import importlib
 import pathlib
 import re
 import typing as t
@@ -24,6 +25,7 @@ from sqlmesh.core.config.common import (
 )
 from sqlmesh.core.engine_adapter.shared import CatalogSupport
 from sqlmesh.core.engine_adapter import EngineAdapter
+from sqlmesh.utils import str_to_bool
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import (
     ValidationInfo,
@@ -49,6 +51,30 @@ FORBIDDEN_STATE_SYNC_ENGINES = {
 MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)")
 
 
+def _get_engine_import_validator(
+    import_name: str, engine_type: str, extra_name: t.Optional[str] = None
+) -> t.Callable:
+    extra_name = extra_name or engine_type
+
+    @model_validator(mode="before")
+    def validate(cls: t.Any, data: t.Any) -> t.Any:
+        check_import = (
+            str_to_bool(str(data.pop("check_import", True))) if isinstance(data, dict) else True
+        )
+        if not check_import:
+            return data
+        try:
+            importlib.import_module(import_name)
+        except ImportError:
+            raise ConfigError(
+                f"Failed to import the '{engine_type}' engine library. Please run `pip install \"sqlmesh[{extra_name}]\"`."
+            )
+
+        return data
+
+    return validate
+
+
 class ConnectionConfig(abc.ABC, BaseConfig):
     type_: str
     concurrent_tasks: int
@@ -428,6 +454,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
     type_: t.Literal["snowflake"] = Field(alias="type", default="snowflake")
 
     _concurrent_tasks_validator = concurrent_tasks_validator
+    _engine_import_validator = _get_engine_import_validator("snowflake", "snowflake")
 
     @model_validator(mode="before")
     def _validate_authenticator(cls, data: t.Any) -> t.Any:
@@ -621,6 +648,7 @@ class DatabricksConnectionConfig(ConnectionConfig):
 
     _concurrent_tasks_validator = concurrent_tasks_validator
     _http_headers_validator = http_headers_validator
+    _engine_import_validator = _get_engine_import_validator("databricks", "databricks")
 
     @model_validator(mode="before")
     def _databricks_connect_validator(cls, data: t.Any) -> t.Any:
@@ -873,6 +901,8 @@ class BigQueryConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["bigquery"] = Field(alias="type", default="bigquery")
 
+    _engine_import_validator = _get_engine_import_validator("google.cloud.bigquery", "bigquery")
+
     @field_validator("execution_project")
     def validate_execution_project(
         cls,
@@ -1015,6 +1045,10 @@ class GCPPostgresConnectionConfig(ConnectionConfig):
     register_comments: bool = True
     pre_ping: bool = True
 
+    _engine_import_validator = _get_engine_import_validator(
+        "google.cloud.sql", "gcp_postgres", "gcppostgres"
+    )
+
     @model_validator(mode="before")
     def _validate_auth_method(cls, data: t.Any) -> t.Any:
         if not isinstance(data, dict):
@@ -1142,6 +1176,8 @@ class RedshiftConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["redshift"] = Field(alias="type", default="redshift")
 
+    _engine_import_validator = _get_engine_import_validator("redshift_connector", "redshift")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1201,6 +1237,8 @@ class PostgresConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["postgres"] = Field(alias="type", default="postgres")
 
+    _engine_import_validator = _get_engine_import_validator("psycopg2", "postgres")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1252,6 +1290,8 @@ class MySQLConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["mysql"] = Field(alias="type", default="mysql")
 
+    _engine_import_validator = _get_engine_import_validator("pymysql", "mysql")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         connection_keys = {
@@ -1302,6 +1342,8 @@ class MSSQLConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["mssql"] = Field(alias="type", default="mssql")
 
+    _engine_import_validator = _get_engine_import_validator("pymssql", "mssql")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1357,6 +1399,8 @@ class SparkConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["spark"] = Field(alias="type", default="spark")
 
+    _engine_import_validator = _get_engine_import_validator("pyspark", "spark")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1473,6 +1517,8 @@ class TrinoConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["trino"] = Field(alias="type", default="trino")
 
+    _engine_import_validator = _get_engine_import_validator("trino", "trino")
+
     @field_validator("schema_location_mapping", mode="before")
     @classmethod
     def _validate_regex_keys(
@@ -1623,6 +1669,8 @@ class ClickhouseConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["clickhouse"] = Field(alias="type", default="clickhouse")
 
+    _engine_import_validator = _get_engine_import_validator("clickhouse_connect", "clickhouse")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         kwargs = {
@@ -1727,6 +1775,8 @@ class AthenaConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["athena"] = Field(alias="type", default="athena")
 
+    _engine_import_validator = _get_engine_import_validator("pyathena", "athena")
+
     @model_validator(mode="after")
     def _root_validator(self) -> Self:
         work_group = self.work_group
@@ -1793,6 +1843,8 @@ class RisingwaveConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["risingwave"] = Field(alias="type", default="risingwave")
 
+    _engine_import_validator = _get_engine_import_validator("psycopg2", "risingwave")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.184.3.dev3
+Version: 0.184.3.dev5
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/tests/core/engine_adapter/integration/config.yaml

@@ -17,6 +17,7 @@ gateways:
       catalog: datalake
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_trino_iceberg:
@@ -28,6 +29,7 @@ gateways:
       catalog: datalake_iceberg
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_trino_delta:
@@ -39,6 +41,7 @@ gateways:
       catalog: datalake_delta
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_trino_nessie:
@@ -50,6 +53,7 @@ gateways:
       catalog: datalake_nessie
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_spark:
@@ -57,6 +61,7 @@ gateways:
       type: spark
       config:
         spark.remote: sc://{{ env_var('DOCKER_HOSTNAME', 'localhost') }}
+      check_import: false
     state_connection:
       type: duckdb
   inttest_mssql:
@@ -65,6 +70,7 @@ gateways:
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
       user: sa
       password: 1StrongPwd@@
+      check_import: false
   inttest_postgres:
     connection:
       type: postgres
@@ -73,6 +79,7 @@ gateways:
       database: postgres
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
       port: 5432
+      check_import: false
   inttest_mysql:
     connection:
       type: mysql
@@ -81,6 +88,7 @@ gateways:
       password: mysql
       port: 3306
       charset: utf8
+      check_import: false
   inttest_clickhouse_standalone:
     connection:
       type: clickhouse
@@ -88,6 +96,7 @@ gateways:
       port: 8123
       username: clickhouse
       password: clickhouse
+      check_import: false
     state_connection:
       type: duckdb
   inttest_clickhouse_cluster:
@@ -98,6 +107,7 @@ gateways:
       username: clickhouse
       password: clickhouse
       cluster: cluster1
+      check_import: false
     state_connection:
       type: duckdb
   inttest_risingwave:
@@ -107,6 +117,7 @@ gateways:
       database: dev
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
       port: 4566
+      check_import: false
 
 
   # Cloud databases
@@ -118,6 +129,7 @@ gateways:
       database: {{ env_var('SNOWFLAKE_DATABASE') }}
       user: {{ env_var('SNOWFLAKE_USER') }}
       password: {{ env_var('SNOWFLAKE_PASSWORD') }}
+      check_import: false
     state_connection:
       type: duckdb
 
@@ -128,6 +140,7 @@ gateways:
       server_hostname: {{ env_var('DATABRICKS_SERVER_HOSTNAME') }}
       http_path: {{ env_var('DATABRICKS_HTTP_PATH') }}
       access_token: {{ env_var('DATABRICKS_ACCESS_TOKEN') }}
+      check_import: false
 
   inttest_redshift:
     connection:
@@ -136,12 +149,14 @@ gateways:
       user: {{ env_var('REDSHIFT_USER') }}
       password: {{ env_var('REDSHIFT_PASSWORD') }}
       database: {{ env_var('REDSHIFT_DATABASE') }}
+      check_import: false
 
   inttest_bigquery:
     connection:
       type: bigquery
       method: service-account
       keyfile: {{ env_var('BIGQUERY_KEYFILE') }}
+      check_import: false
     state_connection:
       type: duckdb
 
@@ -155,6 +170,7 @@ gateways:
       connect_timeout: 30
       connection_pool_options:
         retries: 5
+      check_import: false
     state_connection:
       type: duckdb
 
@@ -166,6 +182,7 @@ gateways:
       region_name: {{ env_var("AWS_REGION") }}
       work_group: {{ env_var("ATHENA_WORK_GROUP", "primary") }}
       s3_warehouse_location: {{ env_var("ATHENA_S3_WAREHOUSE_LOCATION", "") }}
+      check_import: false
     state_connection:
       type: duckdb
 

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/tests/core/test_config.py

@@ -343,6 +343,7 @@ model_defaults:
         os.environ,
         {
             "SQLMESH__GATEWAYS__TESTING__STATE_CONNECTION__TYPE": "bigquery",
+            "SQLMESH__GATEWAYS__TESTING__STATE_CONNECTION__CHECK_IMPORT": "false",
             "SQLMESH__DEFAULT_GATEWAY": "testing",
         },
     ):
@@ -353,7 +354,7 @@ model_defaults:
             gateways={
                 "testing": GatewayConfig(
                     connection=MotherDuckConnectionConfig(database="blah"),
-                    state_connection=BigQueryConnectionConfig(),
+                    state_connection=BigQueryConnectionConfig(check_import=False),
                 ),
             },
             model_defaults=ModelDefaultsConfig(dialect="bigquery"),
@@ -373,6 +374,7 @@ config = Config(gateways={"duckdb_gateway": GatewayConfig(connection=DuckDBConne
         os.environ,
         {
             "SQLMESH__GATEWAYS__DUCKDB_GATEWAY__STATE_CONNECTION__TYPE": "bigquery",
+            "SQLMESH__GATEWAYS__DUCKDB_GATEWAY__STATE_CONNECTION__CHECK_IMPORT": "false",
             "SQLMESH__DEFAULT_GATEWAY": "duckdb_gateway",
         },
     ):
@@ -384,7 +386,7 @@ config = Config(gateways={"duckdb_gateway": GatewayConfig(connection=DuckDBConne
             gateways={  # type: ignore
                 "duckdb_gateway": GatewayConfig(
                     connection=DuckDBConnectionConfig(),
-                    state_connection=BigQueryConnectionConfig(),
+                    state_connection=BigQueryConnectionConfig(check_import=False),
                 ),
             },
             model_defaults=ModelDefaultsConfig(dialect=""),
@@ -823,6 +825,7 @@ def test_gcp_postgres_ip_and_scopes(tmp_path):
       gcp_postgres:
         connection:
           type: gcp_postgres
+          check_import: false
           instance_connection_string: something
           user: user
           password: password

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/tests/core/test_connection_config.py

@@ -20,8 +20,10 @@ from sqlmesh.core.config.connection import (
     TrinoAuthenticationMethod,
     AthenaConnectionConfig,
     _connection_config_validator,
+    _get_engine_import_validator,
 )
 from sqlmesh.utils.errors import ConfigError
+from sqlmesh.utils.pydantic import PydanticModel
 
 
 @pytest.fixture
@@ -750,6 +752,7 @@ def test_bigquery(make_config):
         project="project",
         execution_project="execution_project",
         quota_project="quota_project",
+        check_import=False,
     )
 
     assert isinstance(config, BigQueryConnectionConfig)
@@ -760,10 +763,10 @@ def test_bigquery(make_config):
     assert config.is_recommended_for_state_sync is False
 
     with pytest.raises(ConfigError, match="you must also specify the `project` field"):
-        make_config(type="bigquery", execution_project="execution_project")
+        make_config(type="bigquery", execution_project="execution_project", check_import=False)
 
     with pytest.raises(ConfigError, match="you must also specify the `project` field"):
-        make_config(type="bigquery", quota_project="quota_project")
+        make_config(type="bigquery", quota_project="quota_project", check_import=False)
 
 
 def test_postgres(make_config):
@@ -786,6 +789,7 @@ def test_gcp_postgres(make_config):
         user="user",
         password="password",
         db="database",
+        check_import=False,
     )
     assert isinstance(config, GCPPostgresConnectionConfig)
     assert config.is_recommended_for_state_sync is True
@@ -797,6 +801,7 @@ def test_gcp_postgres(make_config):
         password="password",
         db="database",
         ip_type="private",
+        check_import=False,
     )
     assert config.ip_type == "private"
 
@@ -807,6 +812,7 @@ def test_mysql(make_config):
         host="host",
         user="user",
         password="password",
+        check_import=False,
     )
     assert isinstance(config, MySQLConnectionConfig)
     assert config.is_recommended_for_state_sync is True
@@ -994,3 +1000,36 @@ def test_databricks(make_config):
             server_hostname="dbc-test.cloud.databricks.com",
             auth_type="databricks-oauth",
         )
+
+
+def test_engine_import_validator():
+    with pytest.raises(
+        ConfigError,
+        match=re.escape(
+            """Failed to import the 'bigquery' engine library. Please run `pip install "sqlmesh[bigquery]"`."""
+        ),
+    ):
+
+        class TestConfigA(PydanticModel):
+            _engine_import_validator = _get_engine_import_validator("missing", "bigquery")
+
+        TestConfigA()
+
+    with pytest.raises(
+        ConfigError,
+        match=re.escape(
+            """Failed to import the 'bigquery' engine library. Please run `pip install "sqlmesh[bigquery_extra]"`."""
+        ),
+    ):
+
+        class TestConfigB(PydanticModel):
+            _engine_import_validator = _get_engine_import_validator(
+                "missing", "bigquery", "bigquery_extra"
+            )
+
+        TestConfigB()
+
+    class TestConfigC(PydanticModel):
+        _engine_import_validator = _get_engine_import_validator("sqlmesh", "bigquery")
+
+    TestConfigC()

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/vscode/react/src/components/graph/ModelNode.tsx

@@ -8,6 +8,7 @@ import { Position, type NodeProps } from 'reactflow'
 import { type Column } from '@/api/client'
 import ModelNodeHeaderHandles from './ModelNodeHeaderHandles'
 import ModelColumns from './ModelColumns'
+import type { Model } from '@/domain/model'
 
 export const EnumLineageNodeModelType = {
   ...ModelType,

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/vscode/react/src/components/graph/context.tsx

@@ -9,6 +9,9 @@ import {
 import { getNodeMap, hasActiveEdge, hasActiveEdgeConnector } from './help'
 import { EnumSide } from './types'
 import { type Node } from 'reactflow'
+import type { Lineage } from '@/domain/lineage'
+import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
+import type { Model } from '@/domain/model'
 export interface Connections {
   left: string[]
   right: string[]
@@ -19,10 +22,6 @@ export type ActiveNodes = Set<string>
 export type SelectedNodes = Set<string>
 export type HighlightedNodes = Record<string, string[]>
 
-type Lineage = any
-type ModelSQLMeshModel = any
-type Model = any
-
 interface LineageFlow {
   lineage: Record<string, Lineage>
   lineageCache?: Record<string, Lineage>
@@ -33,7 +32,7 @@ interface LineageFlow {
   selectedNodes: SelectedNodes
   selectedEdges: any[]
   models: Record<string, Model>
-  unknownModels: Record<string, Model>
+  unknownModels: Set<string>
   connections: Map<string, Connections>
   withConnected: boolean
   withColumns: boolean
@@ -84,7 +83,7 @@ export const LineageFlowContext = createContext<LineageFlow>({
   mainNode: undefined,
   activeEdges: new Map(),
   activeNodes: new Set(),
-  models: new Map(),
+  models: {},
   unknownModels: new Set(),
   manuallySelectedColumn: undefined,
   connections: new Map(),

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/vscode/react/src/components/graph/help.ts

@@ -17,8 +17,7 @@ import {
 } from './ModelNode'
 import type { Lineage } from '@/domain/lineage'
 import type { ConnectedNode } from '@/workers/lineage'
-
-type ModelSQLMeshModel = any
+import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
 
 export interface GraphNodeData {
   label: string

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/vscode/react/src/domain/model.ts

@@ -1,11 +1,14 @@
-interface Model {
+import type { Column } from '@/api/client'
+
+export interface Model {
+  displayName: string
   name: string
   fqn: string
   path: string
   dialect: string
   full_path: string
   // type: ModelType
-  // columns: Column[]
+  columns: Column[]
   // default_catalog?: ModelDefaultCatalog
   // description?: ModelDescription
   // sql?: ModelSql

{sqlmesh-0.184.3.dev3 → sqlmesh-0.184.3.dev5}/vscode/react/src/pages/lineage.tsx

@@ -14,6 +14,7 @@ import { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
 import { useEventBus } from '@/hooks/eventBus'
 import type { VSCodeEvent } from '@bus/callbacks'
 import { URI } from 'vscode-uri'
+import type { Model } from '@/domain/model'
 
 export function LineagePage() {
   const { emit } = useEventBus()