Merge remote-tracking branch 'upstream/main'

williebsweet · williebsweet · commit 195bcdeb2000 · 2025-07-31T09:23:02.000-04:00
diff --git a/sqlmesh/core/linter/helpers.py b/sqlmesh/core/linter/helpers.py
@@ -158,3 +158,55 @@ def get_range_of_model_block(
     return Range(
         start=start_position.to_range(splitlines).start, end=end_position.to_range(splitlines).end
     )
+
+
+def get_range_of_a_key_in_model_block(
+    sql: str,
+    dialect: str,
+    key: str,
+) -> t.Optional[Range]:
+    """
+    Get the range of a specific key in the model block of an SQL file.
+    """
+    tokens = tokenize(sql, dialect=dialect)
+    if tokens is None:
+        return None
+
+    # Find the start of the model block
+    start_index = next(
+        (
+            i
+            for i, t in enumerate(tokens)
+            if t.token_type is TokenType.VAR and t.text.upper() == "MODEL"
+        ),
+        None,
+    )
+    end_index = next(
+        (i for i, t in enumerate(tokens) if t.token_type is TokenType.SEMICOLON),
+        None,
+    )
+    if start_index is None or end_index is None:
+        return None
+    if start_index >= end_index:
+        return None
+
+    tokens_of_interest = tokens[start_index + 1 : end_index]
+    # Find the key token
+    key_token = next(
+        (
+            t
+            for t in tokens_of_interest
+            if t.token_type is TokenType.VAR and t.text.upper() == key.upper()
+        ),
+        None,
+    )
+    if key_token is None:
+        return None
+
+    position = TokenPositionDetails(
+        line=key_token.line,
+        col=key_token.col,
+        start=key_token.start,
+        end=key_token.end,
+    )
+    return position.to_range(sql.splitlines())
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
@@ -523,8 +523,6 @@ def render(
             runtime_stage, start, end, execution_time, *kwargs.values()
         )
 
-        needs_optimization = needs_optimization and self._optimize_query_flag
-
         if should_cache and self._optimized_cache:
             query = self._optimized_cache
         else:
@@ -560,7 +558,7 @@ def render(
                 )
                 raise
 
-            if needs_optimization:
+            if needs_optimization and self._optimize_query_flag:
                 deps = d.find_tables(
                     query, default_catalog=self._default_catalog, dialect=self._dialect
                 )
diff --git a/tests/core/linter/test_helpers.py b/tests/core/linter/test_helpers.py
@@ -1,5 +1,9 @@
 from sqlmesh import Context
-from sqlmesh.core.linter.helpers import read_range_from_file, get_range_of_model_block
+from sqlmesh.core.linter.helpers import (
+    read_range_from_file,
+    get_range_of_model_block,
+    get_range_of_a_key_in_model_block,
+)
 from sqlmesh.core.model import SqlModel
 
 
@@ -34,3 +38,41 @@ def test_get_position_of_model_block():
         read_range = read_range_from_file(path, range)
         assert read_range.startswith("MODEL")
         assert read_range.endswith(";")
+
+
+def test_get_range_of_a_key_in_model_block_testing_on_sushi():
+    context = Context(paths=["examples/sushi"])
+
+    sql_models = [
+        model
+        for model in context.models.values()
+        if isinstance(model, SqlModel)
+        and model._path is not None
+        and str(model._path).endswith(".sql")
+    ]
+    assert len(sql_models) > 0
+
+    for model in sql_models:
+        possible_keys = ["name", "tags", "description", "columns", "owner", "cron", "dialect"]
+
+        dialect = model.dialect
+        assert dialect is not None
+
+        path = model._path
+        assert path is not None
+
+        with open(path, "r", encoding="utf-8") as file:
+            content = file.read()
+
+        count_properties_checked = 0
+
+        for key in possible_keys:
+            range = get_range_of_a_key_in_model_block(content, dialect, key)
+
+            # Check that the range starts with the key and ends with ;
+            if range:
+                read_range = read_range_from_file(path, range)
+                assert read_range.lower() == key.lower()
+                count_properties_checked += 1
+
+        assert count_properties_checked > 0
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
@@ -11061,3 +11061,38 @@ def entrypoint(context, **kwargs):
 
     assert model_daily is not None
     assert model_daily.cron == "@daily"
+
+
+def test_render_query_optimize_query_false(assert_exp_eq, sushi_context):
+    snapshots = sushi_context.snapshots
+
+    model = sushi_context.get_model("sushi.top_waiters")
+    model = model.copy(update={"optimize_query": False})
+
+    upstream_model_version = sushi_context.get_snapshot("sushi.waiter_revenue_by_day").version
+
+    assert_exp_eq(
+        model.render_query(snapshots=snapshots).sql(),
+        f"""
+        WITH "test_macros" AS (
+          SELECT
+            2 AS "lit_two",
+            "revenue" * 2.0 AS "sql_exp",
+            CAST("revenue" AS TEXT) AS "sql_lit"
+          FROM "memory"."sqlmesh__sushi"."sushi__waiter_revenue_by_day__{upstream_model_version}" AS "waiter_revenue_by_day" /* memory.sushi.waiter_revenue_by_day */
+        )
+        SELECT
+          CAST("waiter_id" AS INT) AS "waiter_id",
+          CAST("revenue" AS DOUBLE) AS "revenue"
+        FROM "memory"."sqlmesh__sushi"."sushi__waiter_revenue_by_day__{upstream_model_version}" AS "waiter_revenue_by_day" /* memory.sushi.waiter_revenue_by_day */
+        WHERE
+          "event_date" = (
+            SELECT
+              MAX("event_date")
+            FROM "memory"."sqlmesh__sushi"."sushi__waiter_revenue_by_day__{upstream_model_version}" AS "waiter_revenue_by_day" /* memory.sushi.waiter_revenue_by_day */
+          )
+        ORDER BY
+          "revenue" DESC
+        LIMIT 10
+        """,
+    )
diff --git a/vscode/extension/tests/lineage.spec.ts b/vscode/extension/tests/lineage.spec.ts
@@ -19,8 +19,6 @@ import {
  * Helper function to launch VS Code and test lineage with given project path config
  */
 async function testLineageWithProjectPath(page: Page): Promise<void> {
-  await page.waitForLoadState('networkidle')
-  await page.waitForLoadState('domcontentloaded')
   await openLineageView(page)
   await waitForLoadedSQLMesh(page)
 }
@@ -37,120 +35,87 @@ test('Lineage panel renders correctly - no project path config (default)', async
   await testLineageWithProjectPath(page)
 })
 
-test.skip('Lineage panel renders correctly - relative project path', async ({
+test('Lineage panel renders correctly - relative project path', async ({
   page,
   sharedCodeServer,
+  tempDir,
 }) => {
-  const workspaceDir = await fs.mkdtemp(
-    path.join(os.tmpdir(), 'vscode-test-workspace-'),
-  )
-  const projectDir = path.join(workspaceDir, 'projects', 'sushi')
+  const projectDir = path.join(tempDir, 'projects', 'sushi')
   await fs.copy(SUSHI_SOURCE_PATH, projectDir)
 
-  const context = await startCodeServer({
-    tempDir: workspaceDir,
-  })
-
   const settings = {
     'sqlmesh.projectPath': './projects/sushi',
-    'python.defaultInterpreterPath': context.defaultPythonInterpreter,
+    'python.defaultInterpreterPath': sharedCodeServer.defaultPythonInterpreter,
   }
-  await fs.ensureDir(path.join(workspaceDir, '.vscode'))
-  await fs.writeJson(
-    path.join(workspaceDir, '.vscode', 'settings.json'),
-    settings,
-    { spaces: 2 },
-  )
+  await fs.ensureDir(path.join(tempDir, '.vscode'))
+  await fs.writeJson(path.join(tempDir, '.vscode', 'settings.json'), settings, {
+    spaces: 2,
+  })
 
   try {
-    await openServerPage(page, workspaceDir, sharedCodeServer)
+    await openServerPage(page, tempDir, sharedCodeServer)
     await testLineageWithProjectPath(page)
   } finally {
-    await fs.remove(workspaceDir)
+    await fs.remove(tempDir)
   }
 })
 
-test.skip('Lineage panel renders correctly - absolute project path', async ({
+test('Lineage panel renders correctly - absolute project path', async ({
   page,
   tempDir,
   sharedCodeServer,
 }) => {
-  const workspaceDir = await fs.mkdtemp(
-    path.join(os.tmpdir(), 'vscode-test-workspace-'),
-  )
-  const projectDir = path.join(workspaceDir, 'projects', 'sushi')
-  await fs.ensureDir(path.join(workspaceDir, '.vscode'))
+  // Copy the sushi project to temporary directory
+  const projectDir = path.join(tempDir, 'projects', 'sushi')
+  await fs.ensureDir(path.join(tempDir, '.vscode'))
   await fs.copy(SUSHI_SOURCE_PATH, projectDir)
-  await fs.ensureDir(path.join(workspaceDir, '.vscode'))
-  const context = await startCodeServer({
-    tempDir: workspaceDir,
-  })
 
   const settings = {
     'sqlmesh.projectPath': projectDir,
-    'python.defaultInterpreterPath': context.defaultPythonInterpreter,
+    'python.defaultInterpreterPath': sharedCodeServer.defaultPythonInterpreter,
   }
-  await fs.writeJson(
-    path.join(workspaceDir, '.vscode', 'settings.json'),
-    settings,
-    { spaces: 2 },
-  )
+  await fs.writeJson(path.join(tempDir, '.vscode', 'settings.json'), settings, {
+    spaces: 2,
+  })
 
-  try {
-    await openServerPage(page, tempDir, sharedCodeServer)
-    await testLineageWithProjectPath(page)
-  } finally {
-    await stopCodeServer(context)
-  }
+  await openServerPage(page, tempDir, sharedCodeServer)
+  await testLineageWithProjectPath(page)
 })
 
-test.skip('Lineage panel renders correctly - relative project outside of workspace', async ({
+test('Lineage panel renders correctly - relative project outside of workspace', async ({
   page,
   sharedCodeServer,
   tempDir,
 }) => {
-  const tempFolder = await fs.mkdtemp(
-    path.join(os.tmpdir(), 'vscode-test-workspace-'),
-  )
-  const projectDir = path.join(tempFolder, 'projects', 'sushi')
+  const projectDir = path.join(tempDir, 'projects', 'sushi')
   await fs.copy(SUSHI_SOURCE_PATH, projectDir)
 
-  const workspaceDir = path.join(tempFolder, 'workspace')
+  const workspaceDir = path.join(tempDir, 'workspace')
   await fs.ensureDir(workspaceDir)
-  const context = await startCodeServer({
-    tempDir: workspaceDir,
-  })
 
   const settings = {
     'sqlmesh.projectPath': './../projects/sushi',
-    'python.defaultInterpreterPath': context.defaultPythonInterpreter,
+    'python.defaultInterpreterPath': sharedCodeServer.defaultPythonInterpreter,
   }
   await fs.ensureDir(path.join(workspaceDir, '.vscode'))
   await fs.writeJson(
     path.join(workspaceDir, '.vscode', 'settings.json'),
     settings,
     { spaces: 2 },
   )
-
-  try {
-    await openServerPage(page, tempDir, sharedCodeServer)
-    await testLineageWithProjectPath(page)
-  } finally {
-    await stopCodeServer(context)
-  }
+  await openServerPage(page, workspaceDir, sharedCodeServer)
+  await testLineageWithProjectPath(page)
 })
 
-test.skip('Lineage panel renders correctly - absolute path project outside of workspace', async ({
+test('Lineage panel renders correctly - absolute path project outside of workspace', async ({
   page,
   sharedCodeServer,
+  tempDir,
 }) => {
-  const tempFolder = await fs.mkdtemp(
-    path.join(os.tmpdir(), 'vscode-test-workspace-'),
-  )
-  const projectDir = path.join(tempFolder, 'projects', 'sushi')
+  const projectDir = path.join(tempDir, 'projects', 'sushi')
   await fs.copy(SUSHI_SOURCE_PATH, projectDir)
 
-  const workspaceDir = path.join(tempFolder, 'workspace')
+  const workspaceDir = path.join(tempDir, 'workspace')
   await fs.ensureDir(workspaceDir)
 
   const settings = {

Original file line number	Diff line number	Diff line change
`@@ -523,8 +523,6 @@ def render(`
`523`	`523`	`runtime_stage, start, end, execution_time, *kwargs.values()`
`524`	`524`	`)`
`525`	`525`
`526`		`- needs_optimization = needs_optimization and self._optimize_query_flag`
`527`		`-`
`528`	`526`	`if should_cache and self._optimized_cache:`
`529`	`527`	`query = self._optimized_cache`
`530`	`528`	`else:`
`@@ -560,7 +558,7 @@ def render(`
`560`	`558`	`)`
`561`	`559`	`raise`
`562`	`560`
`563`		`- if needs_optimization:`
	`561`	`+ if needs_optimization and self._optimize_query_flag:`
`564`	`562`	`deps = d.find_tables(`
`565`	`563`	`query, default_catalog=self._default_catalog, dialect=self._dialect`
`566`	`564`	`)`