Merge branch 'main' into feature/DRM/resolve-starlette-deprecation-warnings

Author: cmgoffena13

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'
         );