sqlmesh 0.186.2.dev18__py3-none-any.whl → 0.186.2.dev20__py3-none-any.whl

sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.186.2.dev18'
-__version_tuple__ = version_tuple = (0, 186, 2, 'dev18')
+__version__ = version = '0.186.2.dev20'
+__version_tuple__ = version_tuple = (0, 186, 2, 'dev20')

sqlmesh/core/context.py

@@ -1129,33 +1129,15 @@ class GenericContext(BaseContext, t.Generic[C]):
 
             with open(target._path, "r+", encoding="utf-8") as file:
                 before = file.read()
-                expressions = parse(before, default_dialect=self.config_for_node(target).dialect)
-                if transpile and is_meta_expression(expressions[0]):
-                    for prop in expressions[0].expressions:
-                        if prop.name.lower() == "dialect":
-                            prop.replace(
-                                exp.Property(
-                                    this="dialect",
-                                    value=exp.Literal.string(transpile or target.dialect),
-                                )
-                            )
-
-                format_config = self.config_for_node(target).format
-                after = format_model_expressions(
-                    expressions,
-                    transpile or target.dialect,
-                    rewrite_casts=(
-                        rewrite_casts
-                        if rewrite_casts is not None
-                        else not format_config.no_rewrite_casts
-                    ),
-                    **{**format_config.generator_options, **kwargs},
-                )
 
-                if append_newline is None:
-                    append_newline = format_config.append_newline
-                if append_newline:
-                    after += "\n"
+                after = self._format(
+                    target,
+                    before,
+                    transpile=transpile,
+                    rewrite_casts=rewrite_casts,
+                    append_newline=append_newline,
+                    **kwargs,
+                )
 
                 if not check:
                     file.seek(0)
@@ -1174,6 +1156,44 @@ class GenericContext(BaseContext, t.Generic[C]):
 
         return True
 
+    def _format(
+        self,
+        target: Model | Audit,
+        before: str,
+        *,
+        transpile: t.Optional[str] = None,
+        rewrite_casts: t.Optional[bool] = None,
+        append_newline: t.Optional[bool] = None,
+        **kwargs: t.Any,
+    ) -> str:
+        expressions = parse(before, default_dialect=self.config_for_node(target).dialect)
+        if transpile and is_meta_expression(expressions[0]):
+            for prop in expressions[0].expressions:
+                if prop.name.lower() == "dialect":
+                    prop.replace(
+                        exp.Property(
+                            this="dialect",
+                            value=exp.Literal.string(transpile or target.dialect),
+                        )
+                    )
+
+        format_config = self.config_for_node(target).format
+        after = format_model_expressions(
+            expressions,
+            transpile or target.dialect,
+            rewrite_casts=(
+                rewrite_casts if rewrite_casts is not None else not format_config.no_rewrite_casts
+            ),
+            **{**format_config.generator_options, **kwargs},
+        )
+
+        if append_newline is None:
+            append_newline = format_config.append_newline
+        if append_newline:
+            after += "\n"
+
+        return after
+
     @python_api_analytics
     def plan(
         self,

sqlmesh/lsp/main.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 """A Language Server Protocol (LSP) server for SQL with SQLMesh integration, refactored without globals."""
 
+from itertools import chain
 import logging
 import typing as t
 from pathlib import Path
@@ -20,7 +21,11 @@ from sqlmesh.lsp.api import (
     ApiResponseGetModels,
 )
 from sqlmesh.lsp.completions import get_sql_completions
-from sqlmesh.lsp.context import LSPContext, ModelTarget, render_model as render_model_context
+from sqlmesh.lsp.context import (
+    LSPContext,
+    ModelTarget,
+    render_model as render_model_context,
+)
 from sqlmesh.lsp.custom import (
     ALL_MODELS_FEATURE,
     RENDER_MODEL_FEATURE,
@@ -213,15 +218,29 @@ class SQLMeshLanguageServer:
                 uri = URI(params.text_document.uri)
                 self._ensure_context_for_document(uri)
                 document = ls.workspace.get_text_document(params.text_document.uri)
+                before = document.source
                 if self.lsp_context is None:
                     raise RuntimeError(f"No context found for document: {document.path}")
 
-                # Perform formatting using the loaded context
-                self.lsp_context.context.format(paths=(str(uri.to_path()),))
-                with open(uri.to_path(), "r+", encoding="utf-8") as file:
-                    new_text = file.read()
-
-                # Return a single edit that replaces the entire file.
+                target = next(
+                    (
+                        target
+                        for target in chain(
+                            self.lsp_context.context._models.values(),
+                            self.lsp_context.context._audits.values(),
+                        )
+                        if target._path is not None
+                        and target._path.suffix == ".sql"
+                        and (target._path.samefile(uri.to_path()))
+                    ),
+                    None,
+                )
+                if target is None:
+                    return []
+                after = self.lsp_context.context._format(
+                    target=target,
+                    before=before,
+                )
                 return [
                     types.TextEdit(
                         range=types.Range(
@@ -231,7 +250,7 @@ class SQLMeshLanguageServer:
                                 character=len(document.lines[-1]) if document.lines else 0,
                             ),
                         ),
-                        new_text=new_text,
+                        new_text=after,
                     )
                 ]
             except Exception as e:
@@ -279,21 +298,31 @@ class SQLMeshLanguageServer:
                     raise RuntimeError(f"No context found for document: {document.path}")
 
                 references = get_references(self.lsp_context, uri, params.position)
-                return [
-                    types.LocationLink(
-                        target_uri=reference.uri,
-                        target_selection_range=types.Range(
+                location_links = []
+                for reference in references:
+                    # Use target_range if available (for CTEs), otherwise default to start of file
+                    if reference.target_range:
+                        target_range = reference.target_range
+                        target_selection_range = reference.target_range
+                    else:
+                        target_range = types.Range(
                             start=types.Position(line=0, character=0),
                             end=types.Position(line=0, character=0),
-                        ),
-                        target_range=types.Range(
+                        )
+                        target_selection_range = types.Range(
                             start=types.Position(line=0, character=0),
                             end=types.Position(line=0, character=0),
-                        ),
-                        origin_selection_range=reference.range,
+                        )
+
+                    location_links.append(
+                        types.LocationLink(
+                            target_uri=reference.uri,
+                            target_selection_range=target_selection_range,
+                            target_range=target_range,
+                            origin_selection_range=reference.range,
+                        )
                     )
-                    for reference in references
-                ]
+                return location_links
             except Exception as e:
                 ls.show_message(f"Error getting references: {e}", types.MessageType.Error)
                 return []

sqlmesh/lsp/reference.py

@@ -6,23 +6,27 @@ from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
 from sqlglot import exp
 from sqlmesh.lsp.description import generate_markdown_description
+from sqlglot.optimizer.scope import build_scope
 from sqlmesh.lsp.uri import URI
 from sqlmesh.utils.pydantic import PydanticModel
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 
 class Reference(PydanticModel):
     """
-    A reference to a model.
+    A reference to a model or CTE.
 
     Attributes:
         range: The range of the reference in the source file
         uri: The uri of the referenced model
         markdown_description: The markdown description of the referenced model
+        target_range: The range of the definition for go-to-definition (optional, used for CTEs)
     """
 
     range: Range
     uri: str
     markdown_description: t.Optional[str] = None
+    target_range: t.Optional[Range] = None
 
 
 def by_position(position: Position) -> t.Callable[[Reference], bool]:
@@ -88,6 +92,7 @@ def get_model_definitions_for_a_path(
     - Need to normalize it before matching
     - Try get_model before normalization
     - Match to models that the model refers to
+    - Also find CTE references within the query
     """
     path = document_uri.to_path()
     if path.suffix != ".sql":
@@ -126,66 +131,95 @@ def get_model_definitions_for_a_path(
     # Find all possible references
     references = []
 
-    # Get SQL query and find all table references
-    tables = list(query.find_all(exp.Table))
-    if len(tables) == 0:
-        return []
-
     with open(file_path, "r", encoding="utf-8") as file:
         read_file = file.readlines()
 
-    for table in tables:
-        # Normalize the table reference
-        unaliased = table.copy()
-        if unaliased.args.get("alias") is not None:
-            unaliased.set("alias", None)
-        reference_name = unaliased.sql(dialect=dialect)
-        try:
-            normalized_reference_name = normalize_model_name(
-                reference_name,
-                default_catalog=lint_context.context.default_catalog,
-                dialect=dialect,
-            )
-            if normalized_reference_name not in depends_on:
-                continue
-        except Exception:
-            # Skip references that cannot be normalized
-            continue
-
-        # Get the referenced model uri
-        referenced_model = lint_context.context.get_model(
-            model_or_snapshot=normalized_reference_name, raise_if_missing=False
-        )
-        if referenced_model is None:
-            continue
-        referenced_model_path = referenced_model._path
-        # Check whether the path exists
-        if not referenced_model_path.is_file():
-            continue
-        referenced_model_uri = URI.from_path(referenced_model_path)
-
-        # Extract metadata for positioning
-        table_meta = TokenPositionDetails.from_meta(table.this.meta)
-        table_range = _range_from_token_position_details(table_meta, read_file)
-        start_pos = table_range.start
-        end_pos = table_range.end
-
-        # If there's a catalog or database qualifier, adjust the start position
-        catalog_or_db = table.args.get("catalog") or table.args.get("db")
-        if catalog_or_db is not None:
-            catalog_or_db_meta = TokenPositionDetails.from_meta(catalog_or_db.meta)
-            catalog_or_db_range = _range_from_token_position_details(catalog_or_db_meta, read_file)
-            start_pos = catalog_or_db_range.start
-
-        description = generate_markdown_description(referenced_model)
-
-        references.append(
-            Reference(
-                uri=referenced_model_uri.value,
-                range=Range(start=start_pos, end=end_pos),
-                markdown_description=description,
-            )
-        )
+    # Build scope tree to properly handle nested CTEs
+    query = normalize_identifiers(query.copy(), dialect=dialect)
+    root_scope = build_scope(query)
+
+    if root_scope:
+        # Traverse all scopes to find CTE definitions and table references
+        for scope in root_scope.traverse():
+            for table in scope.tables:
+                table_name = table.name
+
+                # Check if this table reference is a CTE in the current scope
+                if cte_scope := scope.cte_sources.get(table_name):
+                    cte = cte_scope.expression.parent
+                    alias = cte.args["alias"]
+                    if isinstance(alias, exp.TableAlias):
+                        identifier = alias.this
+                        if isinstance(identifier, exp.Identifier):
+                            target_range = _range_from_token_position_details(
+                                TokenPositionDetails.from_meta(identifier.meta), read_file
+                            )
+                            table_range = _range_from_token_position_details(
+                                TokenPositionDetails.from_meta(table.this.meta), read_file
+                            )
+                            references.append(
+                                Reference(
+                                    uri=document_uri.value,  # Same file
+                                    range=table_range,
+                                    target_range=target_range,
+                                )
+                            )
+                    continue
+
+                # For non-CTE tables, process as before (external model references)
+                # Normalize the table reference
+                unaliased = table.copy()
+                if unaliased.args.get("alias") is not None:
+                    unaliased.set("alias", None)
+                reference_name = unaliased.sql(dialect=dialect)
+                try:
+                    normalized_reference_name = normalize_model_name(
+                        reference_name,
+                        default_catalog=lint_context.context.default_catalog,
+                        dialect=dialect,
+                    )
+                    if normalized_reference_name not in depends_on:
+                        continue
+                except Exception:
+                    # Skip references that cannot be normalized
+                    continue
+
+                # Get the referenced model uri
+                referenced_model = lint_context.context.get_model(
+                    model_or_snapshot=normalized_reference_name, raise_if_missing=False
+                )
+                if referenced_model is None:
+                    continue
+                referenced_model_path = referenced_model._path
+                # Check whether the path exists
+                if not referenced_model_path.is_file():
+                    continue
+                referenced_model_uri = URI.from_path(referenced_model_path)
+
+                # Extract metadata for positioning
+                table_meta = TokenPositionDetails.from_meta(table.this.meta)
+                table_range = _range_from_token_position_details(table_meta, read_file)
+                start_pos = table_range.start
+                end_pos = table_range.end
+
+                # If there's a catalog or database qualifier, adjust the start position
+                catalog_or_db = table.args.get("catalog") or table.args.get("db")
+                if catalog_or_db is not None:
+                    catalog_or_db_meta = TokenPositionDetails.from_meta(catalog_or_db.meta)
+                    catalog_or_db_range = _range_from_token_position_details(
+                        catalog_or_db_meta, read_file
+                    )
+                    start_pos = catalog_or_db_range.start
+
+                description = generate_markdown_description(referenced_model)
+
+                references.append(
+                    Reference(
+                        uri=referenced_model_uri.value,
+                        range=Range(start=start_pos, end=end_pos),
+                        markdown_description=description,
+                    )
+                )
 
     return references
 

{sqlmesh-0.186.2.dev18.dist-info → sqlmesh-0.186.2.dev20.dist-info}/METADATA

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

{sqlmesh-0.186.2.dev18.dist-info → sqlmesh-0.186.2.dev20.dist-info}/RECORD

@@ -1,5 +1,5 @@
 sqlmesh/__init__.py,sha256=69xw6bTmi4cW90M-sl3g6wLDgxVcLHcKg1q8GS3lqyA,7466
-sqlmesh/_version.py,sha256=23uEi0q4YZ9-xEEuuYzX2qYiIPhDRAIz-zuV9UYt5Ng,530
+sqlmesh/_version.py,sha256=md758L-6fCOTdjqRZ-XWWVl4pSKbkoIiqBeQD-HssQk,530
 sqlmesh/magics.py,sha256=JHorAZNtezf123jmTyNSoUC9yHmv-08hMP205PJmmIo,39750
 sqlmesh/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlmesh/cicd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -13,7 +13,7 @@ sqlmesh/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlmesh/core/_typing.py,sha256=e1f7xO2aZCy38M_3VoDy-XI3shyd99-IFctgIFSgXWM,497
 sqlmesh/core/console.py,sha256=ZAaGLOWBJGBxPzLV7KZZ9akPExgD5m5AboCM8O0Q0nY,150944
 sqlmesh/core/constants.py,sha256=oqt6zjepUhQoWlgapKn_HpG7Pu9mrFwepdX5caNfJUI,2499
-sqlmesh/core/context.py,sha256=paXJMdaydEQ9Tp80fVUnxveBlbCN7kGGEvesqAiezmw,115276
+sqlmesh/core/context.py,sha256=gOwToFkXSQjb71Ki3cEqfxYjS64p2T2XKfZPVBtjGQA,115624
 sqlmesh/core/context_diff.py,sha256=s9toP17d4tJfOpFH91LcjGK2CM64PVuB0MPkeFgF4hQ,20415
 sqlmesh/core/dialect.py,sha256=fPTV517IOS0nWMxYznNNQnXkQkFFrPOSKe8ghsQ0o8A,49759
 sqlmesh/core/environment.py,sha256=0TTQf4C4LXwf37fm-dfYv1qyG7vtZfQI5rizgm5g0j0,11442
@@ -160,8 +160,8 @@ sqlmesh/lsp/completions.py,sha256=Uo8lUzPBeoB58Eu2642YoAtNI0pgIFWFSvwmuJALTv4,35
 sqlmesh/lsp/context.py,sha256=nCWg2uVH28-IJC2uDYuCbiXTh1lUFElSpNu4Q-fDWi4,2918
 sqlmesh/lsp/custom.py,sha256=1YVcQz0y-4cdsvTdOQvlhUQXwOqPBWi3v_u1X2EqQxE,901
 sqlmesh/lsp/description.py,sha256=4QIGC9ix8IzmWF6DVUNhY4m1Crk3sNQA6hZ1RbyqU78,827
-sqlmesh/lsp/main.py,sha256=5EF1xA1rESYq4sTKDJPPG9Zl6mx3xfy6dbzWBbDee0M,18553
-sqlmesh/lsp/reference.py,sha256=bO2edkI2avmEYayQI11_OsDX1gjp36kq9ANwHczWSC0,8285
+sqlmesh/lsp/main.py,sha256=s8e9mD-J6pbxCrEZiMlRL0DvUu1rD3vxVdebgx5kjhA,19609
+sqlmesh/lsp/reference.py,sha256=v4VXMYQY9rPC8b5w-jzob3i77VdZ1FZ8sA_-sWuPx7M,10502
 sqlmesh/lsp/uri.py,sha256=mvSLbZGmMTHPpliThOtVWiFcHnfJzp1w39q9l-xD2SI,802
 sqlmesh/migrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlmesh/migrations/v0001_init.py,sha256=h8FwMYJj9_O8yGpk12guaSaGwX5A3we2GiKeHrXi15o,2048
@@ -270,7 +270,7 @@ sqlmesh/utils/pydantic.py,sha256=I898jAFqa3qU6lG2uPpRebVFncOUo8ZrLWM_R3K3DHU,116
 sqlmesh/utils/rich.py,sha256=xNZyPCqEc4TAlJlsEthg-5eYUWpaizz0QjBzO7a1Yv4,1509
 sqlmesh/utils/windows.py,sha256=r7Yin4WEwbvH7sKA3hOJtRsclmscNuBeCklCfVPvtQs,611
 sqlmesh/utils/yaml.py,sha256=np-WVcwmwsKghCgAj1OBBMsK1SXSRn4wR-H37cGiFr8,2282
-sqlmesh-0.186.2.dev18.dist-info/licenses/LICENSE,sha256=OlMefUjgWJdULtf84BLW0AZZcY8DwdgQqb_1j2862j8,11346
+sqlmesh-0.186.2.dev20.dist-info/licenses/LICENSE,sha256=OlMefUjgWJdULtf84BLW0AZZcY8DwdgQqb_1j2862j8,11346
 web/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 web/client/dist/index.html,sha256=M50OuwjW63dRCS0Dz40Yecgvw9ecLFqizbTa6gv5Nd4,5472
 web/client/dist/assets/Audits-BPlxzqeG.js,sha256=TCx1qCjCGgmUSvH6HoPr5t8ztbRXMlOxrcW4o7hRjlA,1041
@@ -388,8 +388,8 @@ web/server/api/endpoints/models.py,sha256=ud-lRwriyI_OYg9tZchrBgeiXELiRTLNEodraf
 web/server/api/endpoints/modules.py,sha256=8hqqgonGay_mJmpCw0IdbjsPhWlQH2VLdKAqha-myac,468
 web/server/api/endpoints/plan.py,sha256=OYuQ_4dgbhG9ibTcX1a8mvYgH3ANxw9Fhbb4wg6QWqc,9222
 web/server/api/endpoints/table_diff.py,sha256=Wd5YxeKbZB7AkrDvqFgvV9ZQlnKYOF1n5W_w2EHE_Zw,2085
-sqlmesh-0.186.2.dev18.dist-info/METADATA,sha256=zjC2hE6KzopYHThETS-RzPimIwcnBi-IBr4Ogt_7rB4,25524
-sqlmesh-0.186.2.dev18.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-sqlmesh-0.186.2.dev18.dist-info/entry_points.txt,sha256=pIYKtg-u-uO7BK_Y0qrdmfyj0bYrcMc3Xiqr5WE3VFk,121
-sqlmesh-0.186.2.dev18.dist-info/top_level.txt,sha256=bspjo5vjELQkuN0qKuTFIbsrUt3Zeo8D2mlKeHG1LdE,12
-sqlmesh-0.186.2.dev18.dist-info/RECORD,,
+sqlmesh-0.186.2.dev20.dist-info/METADATA,sha256=7njLBJTTA1FSif3Pvs0NSw7SoMzwz3DzupTHsU1ZS40,25524
+sqlmesh-0.186.2.dev20.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+sqlmesh-0.186.2.dev20.dist-info/entry_points.txt,sha256=pIYKtg-u-uO7BK_Y0qrdmfyj0bYrcMc3Xiqr5WE3VFk,121
+sqlmesh-0.186.2.dev20.dist-info/top_level.txt,sha256=bspjo5vjELQkuN0qKuTFIbsrUt3Zeo8D2mlKeHG1LdE,12
+sqlmesh-0.186.2.dev20.dist-info/RECORD,,