test original sql

Author: izeigerman

sqlmesh/core/model/definition.py

@@ -2297,6 +2297,7 @@ def load_sql_based_model(
         variables=variables,
         inline_audits=inline_audits,
         blueprint_variables=blueprint_variables,
+        use_original_sql=True,
         **meta_fields,
     )
 
@@ -2519,6 +2520,7 @@ def _create_model(
     signal_definitions: t.Optional[SignalRegistry] = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
     blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
+    use_original_sql: bool = False,
     **kwargs: t.Any,
 ) -> Model:
     validate_extra_and_required_fields(
@@ -2555,7 +2557,7 @@ def _create_model(
     if "query" in kwargs:
         statements.append(kwargs["query"])
         kwargs["query"] = ParsableSql.from_parsed_expression(
-            kwargs["query"], dialect, use_meta_sql=True
+            kwargs["query"], dialect, use_meta_sql=use_original_sql
         )
 
     # Merge default statements with model-specific statements
@@ -2569,7 +2571,7 @@ def _create_model(
             is_metadata = statement_field == "on_virtual_update"
             statements.extend((stmt, is_metadata) for stmt in kwargs[statement_field])
             kwargs[statement_field] = [
-                ParsableSql.from_parsed_expression(stmt, dialect, use_meta_sql=True)
+                ParsableSql.from_parsed_expression(stmt, dialect, use_meta_sql=use_original_sql)
                 for stmt in kwargs[statement_field]
             ]
 

tests/core/test_model.py

@@ -11471,3 +11471,46 @@ def test_raw_jinja_raw_tag():
 
     model = load_sql_based_model(expressions)
     assert model.render_query().sql() == "SELECT '{{ foo }}' AS \"col\""
+
+
+def test_use_original_sql():
+    expressions = d.parse(
+        """
+        MODEL (name test);
+
+        CREATE TABLE pre (
+          a INT
+        );
+
+        SELECT
+          1,
+          2;
+
+        CREATE TABLE post (
+          b INT
+        );
+        """
+    )
+
+    model = load_sql_based_model(expressions)
+    assert model.query_.sql == "SELECT\n          1,\n          2"
+    assert model.pre_statements_[0].sql == "CREATE TABLE pre (\n          a INT\n        )"
+    assert model.post_statements_[0].sql == "CREATE TABLE post (\n          b INT\n        );"
+
+    # Now manually create the model and make sure that the original SQL is not used
+    model_query = d.parse_one("SELECT 1 AS one")
+    assert model_query.meta["sql"] == "SELECT 1 AS one"
+    model_query = model_query.select("2 AS two")
+
+    pre_statements = [d.parse_one("CREATE TABLE pre (\n          a INT\n        )")]
+    post_statements = [d.parse_one("CREATE TABLE post (\n          b INT\n        );")]
+
+    model = create_sql_model(
+        "test",
+        model_query,
+        pre_statements=pre_statements,
+        post_statements=post_statements,
+    )
+    assert model.query_.sql == "SELECT 1 AS one, 2 AS two"
+    assert model.pre_statements_[0].sql == "CREATE TABLE pre (a INT)"
+    assert model.post_statements_[0].sql == "CREATE TABLE post (b INT)"