SNOW-945984: Fix sql simplifier when non select statement is used wit…

…h limit (#1216)

* fix sql simplifier when non select statement is used with limit

* changelog updates

CHANGELOG.md

@@ -34,12 +34,15 @@
     For instance, `df.select("a", seq1().alias("b")).select("a", "b").sort("a")` won't flatten the sort clause anymore.
   - a window or sequence-dependent data generator column is used after `DataFrame.limit()`. For instance, `df.limit(10).select(row_number().over())` won't flatten the limit and select in the generated SQL.
 - Fixed a bug that aliasing a DataFrame column raises an error when the DataFame is copied from another DataFrame with an aliased column. For instance,
+
   ```python
   df = df.select(col("a").alias("b"))
   df = copy(df)
   df.select(col("b").alias("c"))  # threw an error. Now it's fixed.
   ```
+
 - Fixed a bug in `Session.create_dataframe` that the non-nullable field in schema is not respected for boolean type. Note that this fix is only effective when the user to have the privilege to create a temp table.
+- Fixed a bug in sql simplifier where non-select statements in `session.sql` dropped sql query when used with `limit()`.
 - Fixed a bug that raised an exception when session parameter `ERROR_ON_NONDETERMINISTIC_UPDATE` is true.
 
 ### Behavior Changes (API Compatible)

src/snowflake/snowpark/_internal/analyzer/select_statement.py

@@ -811,6 +811,8 @@ def limit(self, n: int, *, offset: int = 0) -> "SelectStatement":
             new.limit_ = min(self.limit_, n) if self.limit_ else n
             new.offset = offset or self.offset
             new.column_states = self.column_states
+            new.pre_actions = new.from_.pre_actions
+            new.post_actions = new.from_.post_actions
         return new
 
 

tests/integ/test_simplifier_suite.py

@@ -735,6 +735,13 @@ def test_limit(session, simplifier_table):
         == f"SELECT  *  FROM (select * from {simplifier_table}) LIMIT 10"
     )
 
+    # test for non-select sql statement
+    temp_table_name = Utils.random_table_name()
+    query = f"create or replace temporary table {temp_table_name} (bar string)"
+    df = session.sql(query).limit(1)
+    assert df.queries["queries"][-2] == query
+    assert df.collect()[0][0] == f"Table {temp_table_name} successfully created."
+
 
 def test_filter_order_limit_together(session, simplifier_table):
     df = session.table(simplifier_table)