Merge branch 'main' into plan-dedup-snapshot-state-read

Author: fresioAS

pyproject.toml

@@ -78,7 +78,7 @@ dev = [
     "google-auth",
     "google-cloud-bigquery",
     "google-cloud-bigquery-storage",
-    "httpx",
+    "httpx2",
     "mypy~=1.13.0",
     "numpy",
     "pandas-stubs",

tests/core/engine_adapter/integration/test_integration_athena.py

@@ -48,8 +48,21 @@ def s3_list_objects(s3: t.Any, location: str, **list_objects_kwargs: t.Any) -> t
     return lst
 
 
+def s3_delete_objects(s3: t.Any, location: str) -> None:
+    # The S3 location for a given test is stable across pytest-rerunfailures retries within the same
+    # test session (it's derived from the session's testrun_uid + the test's name), so a prior failed
+    # attempt can leave objects behind that a subsequent retry would otherwise trip over. Proactively
+    # clearing the prefix makes these tests self-healing instead of just asserting it's already empty.
+    bucket, prefix = parse_s3_uri(location)
+    for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix):
+        objects = [{"Key": o["Key"]} for o in page.get("Contents", [])]
+        if objects:
+            s3.delete_objects(Bucket=bucket, Delete={"Objects": objects})
+
+
 def test_clear_partition_data(ctx: TestContext, engine_adapter: AthenaEngineAdapter, s3: t.Any):
     base_uri = engine_adapter.s3_warehouse_location_or_raise
+    s3_delete_objects(s3, base_uri)
     assert len(s3_list_objects(s3, base_uri)) == 0
 
     src_table = ctx.table("src_table")
@@ -239,6 +252,7 @@ def test_hive_truncate_table(ctx: TestContext, engine_adapter: AthenaEngineAdapt
         ]
     )
 
+    s3_delete_objects(s3, base_uri)
     assert len(s3_list_objects(s3, base_uri)) == 0
 
     engine_adapter.ctas(table_name=table_1, query_or_df=base_data)

tests/core/engine_adapter/integration/test_integration_bigquery.py

@@ -2,7 +2,6 @@
 from unittest import mock
 
 import pytest
-import time_machine
 from pathlib import Path
 from sqlglot import exp
 from sqlglot.optimizer.qualify_columns import quote_identifiers
@@ -474,7 +473,6 @@ def test_plan_correlation_id_in_job_labels(ctx: TestContext):
     assert labels == {correlation_id.job_type.value.lower(): correlation_id.job_id}
 
 
-@time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_run_correlation_id_in_job_labels(ctx: TestContext):
     run_id = "test_run_id"
     model_name = ctx.table("run_test")
@@ -497,7 +495,7 @@ def test_run_correlation_id_in_job_labels(ctx: TestContext):
             )
         )
     )
-    sqlmesh.plan(auto_apply=True, no_prompts=True)
+    sqlmesh.plan(auto_apply=True, no_prompts=True, execution_time="2023-01-08 15:00:00")
 
     captured_evaluators: t.List = []
     original_scheduler = sqlmesh.scheduler
@@ -510,12 +508,9 @@ def scheduler_wrapper(
             captured_evaluators.append(snapshot_evaluator)
         return original_scheduler(environment=environment, snapshot_evaluator=snapshot_evaluator)
 
-    with time_machine.travel("2023-01-09 00:00:00 UTC"):
-        with mock.patch(
-            "sqlmesh.core.context.analytics.collector.on_run_start", return_value=run_id
-        ):
-            with mock.patch.object(sqlmesh, "scheduler", scheduler_wrapper):
-                sqlmesh.run()
+    with mock.patch("sqlmesh.core.context.analytics.collector.on_run_start", return_value=run_id):
+        with mock.patch.object(sqlmesh, "scheduler", scheduler_wrapper):
+            sqlmesh.run(execution_time="2023-01-09 00:00:00")
 
     assert captured_evaluators
     adapter = t.cast(BigQueryEngineAdapter, captured_evaluators[-1].adapter)

tests/core/engine_adapter/integration/test_integration_postgres.py

@@ -1,4 +1,5 @@
 import typing as t
+import uuid
 from contextlib import contextmanager
 import pytest
 from pytest import FixtureRequest
@@ -53,7 +54,8 @@ def create_users(
             _cleanup_user(engine_adapter, user_name)
 
         for role_name in role_names:
-            user_name = f"test_{role_name}"
+            random_suffix = uuid.uuid4().hex[:6]
+            user_name = f"test_{role_name}_{random_suffix}"
             password = random_id()
             engine_adapter.execute(f"CREATE USER \"{user_name}\" WITH PASSWORD '{password}'")
             engine_adapter.execute(f'GRANT USAGE ON SCHEMA public TO "{user_name}"')
@@ -381,12 +383,14 @@ def test_grants_plan_target_layer_physical_only(
     with create_users(engine_adapter, "reader") as roles:
         (tmp_path / "models").mkdir(exist_ok=True)
 
-        model_def = """
+        reader_username = roles["reader"]["username"]
+
+        model_def = f"""
         MODEL (
             name test_schema.physical_grants_model,
             kind FULL,
             grants (
-                'select' = ['test_reader']
+                'select' = ['{reader_username}']
             ),
             grants_target_layer 'physical'
         );
@@ -421,12 +425,14 @@ def test_grants_plan_target_layer_virtual_only(
     with create_users(engine_adapter, "viewer") as roles:
         (tmp_path / "models").mkdir(exist_ok=True)
 
-        model_def = """
+        viewer_username = roles["viewer"]["username"]
+
+        model_def = f"""
         MODEL (
             name test_schema.virtual_grants_model,
             kind FULL,
             grants (
-                'select' = ['test_viewer']
+                'select' = ['{viewer_username}']
             ),
             grants_target_layer 'virtual'
         );

tests/web/test_main.py

@@ -7,7 +7,7 @@
 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
-from httpx import ASGITransport, AsyncClient
+from httpx2 import ASGITransport, AsyncClient
 from pytest_mock.plugin import MockerFixture
 
 from sqlmesh.core.context import Context

web/server/exceptions.py

@@ -3,7 +3,7 @@
 import typing as t
 
 from fastapi import HTTPException
-from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY
+from starlette.status import HTTP_422_UNPROCESSABLE_CONTENT
 
 from sqlmesh.utils.date import now_timestamp
 from web.server.models import ApiExceptionPayload
@@ -14,7 +14,7 @@ def __init__(
         self,
         origin: str,
         message: str,
-        status_code: int = HTTP_422_UNPROCESSABLE_ENTITY,
+        status_code: int = HTTP_422_UNPROCESSABLE_CONTENT,
         trigger: t.Optional[str] = None,
     ):
         super().__init__(status_code)

web/server/utils.py

@@ -9,7 +9,7 @@
 import pyarrow as pa  # type: ignore
 from fastapi import Depends, HTTPException
 from starlette.responses import StreamingResponse
-from starlette.status import HTTP_404_NOT_FOUND, HTTP_422_UNPROCESSABLE_ENTITY
+from starlette.status import HTTP_404_NOT_FOUND, HTTP_422_UNPROCESSABLE_CONTENT
 
 from sqlmesh.core import constants as c
 from web.server.console import api_console
@@ -45,7 +45,7 @@ def func_wrapper() -> R:
                     message="An unexpected error occurred",
                     origin="API -> utils -> run_in_executor",
                     trigger="An unexpected error occurred",
-                    status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                    status_code=HTTP_422_UNPROCESSABLE_CONTENT,
                 )
             )
             raise e