add mock executor; fix loader; adapt unit tests

Author: themisvaltinos

sqlmesh/core/loader.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
 
 import abc
+import concurrent.futures
 import glob
 import itertools
 import linecache
 import multiprocessing as mp
 import os
 import re
 import typing as t
+import concurrent
 from collections import Counter, defaultdict
 from dataclasses import dataclass
 from pathlib import Path
-from concurrent.futures import ProcessPoolExecutor, as_completed
 
 from sqlglot.errors import SqlglotError
 from sqlglot import exp
@@ -310,7 +311,10 @@ def _load_external_models(
             # external models with no explicit gateway defined form the base set
             for model in external_models:
                 if model.gateway is None:
-                    models[model.fqn] = model
+                    try:
+                        models[model.fqn] = model
+                    except Exception as ex:
+                        raise ConfigError(f"Failed to add model: {model.fqn}\n\n{ex}")
 
             # however, if there is a gateway defined, gateway-specific models take precedence
             if gateway:
@@ -463,20 +467,15 @@ def _load_models(
         audits into a Dict and creates the dag
         """
         cache = SqlMeshLoader._Cache(self, self.config_path)
-        import time
 
-        now = time.time()
         sql_models = self._load_sql_models(macros, jinja_macros, audits, signals, cache, gateway)
-        print("sql models", time.time() - now)
-        now = time.time()
         external_models = self._load_external_models(audits, cache, gateway)
-        print("external models", time.time() - now)
         python_models = self._load_python_models(macros, jinja_macros, audits, signals)
 
         all_model_names = list(sql_models) + list(external_models) + list(python_models)
         duplicates = [name for name, count in Counter(all_model_names).items() if count > 1]
         if duplicates:
-            raise ValueError(f"Duplicate model name(s) found: {', '.join(duplicates)}.")
+            raise ConfigError(f"Duplicate model name(s) found: {', '.join(duplicates)}.")
 
         return UniqueKeyDict("models", **sql_models, **external_models, **python_models)
 
@@ -491,8 +490,7 @@ def _load_sql_models(
     ) -> UniqueKeyDict[str, Model]:
         """Loads the sql models into a Dict"""
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
-
-        paths = set()
+        paths: t.Set[Path] = set()
 
         for path in self._glob_paths(
             self.config_path / c.MODELS,
@@ -507,14 +505,11 @@ def _load_sql_models(
 
         for path in paths.copy():
             cached_models = cache.get(path)
-
             if cached_models:
                 paths.remove(path)
-
                 for model in cached_models:
-                    models[model.fqn] = model
-
-        error = False
+                    if model.enabled:
+                        models[model.fqn] = model
 
         if paths:
             defaults = dict(
@@ -535,31 +530,31 @@ def _load_sql_models(
                 default_catalog_per_gateway=self.context.default_catalog_per_gateway,
             )
 
-            with ProcessPoolExecutor(
+            errors: t.List[str] = []
+            with concurrent.futures.ProcessPoolExecutor(
                 mp_context=mp.get_context("fork"),
                 initializer=_init_model_defaults,
                 initargs=(self.config, gateway, defaults, cache),
                 max_workers=c.MAX_FORK_WORKERS,
             ) as pool:
-                for fut in as_completed(pool.submit(load_sql_models, path) for path in paths):
+                futures_to_paths = {pool.submit(load_sql_models, path): path for path in paths}
+                for fut, path in futures_to_paths.items():
                     try:
-                        path, loaded = fut.result()
-
+                        _, loaded = fut.result()
                         if loaded:
                             for model in loaded:
-                                model._path = path
-                                models[model.fqn] = model
+                                if model.enabled:
+                                    model._path = path
+                                    models[model.fqn] = model
                         else:
                             for model in cache.get(path):
-                                models[model.fqn] = model
+                                if model.enabled:
+                                    models[model.fqn] = model
                     except Exception as ex:
-                        self._console.log_error(
-                            f"Failed to load model definition at '{path}'.\n{ex}"
-                        )
-                        error = True
+                        errors.append(f"Failed to load model definition at '{path}'.\n\n{ex}")
 
-        if error:
-            raise ConfigError("Failed to load models")
+            if errors:
+                raise ConfigError(f"Failed to load models\n\n{'\n'.join(errors)}")
 
         return models
 

tests/conftest.py

@@ -506,3 +506,15 @@ def _make_function(table_name: str, random_id: str) -> exp.Table:
         return temp_table
 
     return _make_function
+
+
+@pytest.fixture(autouse=True)
+def patch_process_pool_executor(mocker: MockerFixture, request):
+    """Patch ProcessPoolExecutor with MockProcessPoolExecutor in all tests except test_forking.py."""
+    # Skip mocking for test_forking.py
+    if request.node.fspath.basename == "test_forking.py":
+        return
+
+    from tests.mock_executor import MockProcessPoolExecutor
+
+    mocker.patch("concurrent.futures.ProcessPoolExecutor", MockProcessPoolExecutor)

tests/core/test_loader.py

@@ -85,14 +85,14 @@ def test_duplicate_model_names_different_kind(tmp_path: Path, sample_models):
         path_3.write_text(model_3["contents"])
 
     with pytest.raises(
-        ValueError, match=r'Duplicate model name\(s\) found: "memory"."test_schema"."test_model".'
+        ConfigError, match=r'Duplicate model name\(s\) found: "memory"."test_schema"."test_model".'
     ):
         Context(paths=tmp_path, config=config)
 
 
 @pytest.mark.parametrize("sample_models", ["sql", "external"], indirect=True)
 def test_duplicate_model_names_same_kind(tmp_path: Path, sample_models):
-    """Test same (SQL and external) models with duplicate model names raises ValueError."""
+    """Test same (SQL and external) models with duplicate model names raises ConfigError."""
 
     def duplicate_model_path(fpath):
         return Path(fpath).parent / ("duplicate" + Path(fpath).suffix)
@@ -109,7 +109,7 @@ def duplicate_model_path(fpath):
     duplicate_fpath.write_text(model["contents"])
 
     with pytest.raises(
-        ValueError,
+        ConfigError,
         match=r'Duplicate key \'"memory"."test_schema"."test_model"\' found in UniqueKeyDict<models>. Call dict.update\(\.\.\.\) if this is intentional.',
     ):
         Context(paths=tmp_path, config=config)

tests/core/test_model.py

@@ -2941,26 +2941,33 @@ def test_model_cache(tmp_path: Path, mocker: MockerFixture):
     expressions = d.parse(
         """
         MODEL (
-            name db.seed,
+            name db.model_sql,
         );
         SELECT 1, ds;
     """
     )
 
     model = load_sql_based_model([e for e in expressions if e])
 
-    loader = mocker.Mock(return_value=[model])
-
-    assert cache.get_or_load("test_model", "test_entry_a", loader=loader)[0].dict() == model.dict()
-    assert cache.get_or_load("test_model", "test_entry_a", loader=loader)[0].dict() == model.dict()
+    assert cache.put([model], "test_model", "test_entry_a")
+    assert cache.get("test_model", "test_entry_a")[0].dict() == model.dict()
 
-    assert cache.get_or_load("test_model", "test_entry_b", loader=loader)[0].dict() == model.dict()
-    assert cache.get_or_load("test_model", "test_entry_b", loader=loader)[0].dict() == model.dict()
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.model_seed,
+            kind SEED (
+              path '../seeds/waiter_names.csv',
+            ),
+        );
+    """
+    )
 
-    assert cache.get_or_load("test_model", "test_entry_a", loader=loader)[0].dict() == model.dict()
-    assert cache.get_or_load("test_model", "test_entry_a", loader=loader)[0].dict() == model.dict()
+    seed_model = load_sql_based_model(
+        expressions, path=Path("./examples/sushi/models/test_model.sql")
+    )
 
-    assert loader.call_count == 2
+    assert not cache.put([seed_model], "test_model", "test_entry_b")
 
 
 @pytest.mark.slow
@@ -2983,7 +2990,7 @@ def test_model_cache_gateway(tmp_path: Path, mocker: MockerFixture):
     assert patched_cache_put.call_count == 0
 
     Context(paths=tmp_path, config=config, gateway="secondary")
-    assert patched_cache_put.call_count == 4
+    assert patched_cache_put.call_count == 2
 
 
 @pytest.mark.slow
@@ -3001,7 +3008,7 @@ def test_model_cache_default_catalog(tmp_path: Path, mocker: MockerFixture):
         PropertyMock(return_value=None),
     ):
         Context(paths=tmp_path)
-        assert patched_cache_put.call_count == 4
+        assert patched_cache_put.call_count == 2
 
 
 def test_model_ctas_query():

tests/mock_executor.py

@@ -0,0 +1,32 @@
+from concurrent.futures import Future
+
+
+class MockProcessPoolExecutor:
+    """A mock implementation of ProcessPoolExecutor for use in tests.
+
+    This executor runs functions synchronously in the same process, avoiding the issues
+    with forking in test environments.
+    """
+
+    def __init__(self, max_workers=None, mp_context=None, initializer=None, initargs=()):
+        if initializer is not None:
+            try:
+                initializer(*initargs)
+            except BaseException as ex:
+                raise RuntimeError(f"Exception in initializer: {ex}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        return True
+
+    def submit(self, fn, *args, **kwargs):
+        """Execute the function synchronously and return a Future with the result."""
+        future = Future()
+        try:
+            result = fn(*args, **kwargs)
+            future.set_result(result)
+        except Exception as e:
+            future.set_exception(e)
+        return future