SNOW-1738538 Fix multi-ast eval validation in decoding expectation te…

…st ast
snowflakedb · Oct 23, 2024 · f2baf28 · f2baf28
1 parent e843c05
commit f2baf28
Show file tree

Hide file tree

Showing 18 changed files with 7,114 additions and 1 deletion.
diff --git a/src/snowflake/snowpark/_internal/ast_utils.py b/src/snowflake/snowpark/_internal/ast_utils.py
@@ -1342,7 +1342,19 @@ def ClearTempTables(message: proto.Request) -> None:
 
 def base64_str_to_request(base64_str: str) -> proto.Request:
     message = proto.Request()
-    message.ParseFromString(base64.b64decode(base64_str.strip()))
+
+    proto_strs = [base64.b64decode(s) for s in base64_str.split("\n")]
+
+    message = proto.Request()
+    message.ParseFromString(proto_strs[0])
+
+    for proto_str in proto_strs[1:]:
+        temp_msg = proto.Request()
+        temp_msg.ParseFromString(proto_str)
+        for temp_stmt in temp_msg.body:
+            stmt = message.body.add()
+            stmt.CopyFrom(temp_stmt)
+
     return message
 
 

diff --git a/tests/ast/data/DataFrame.collect.test b/tests/ast/data/DataFrame.collect.test
@@ -94,6 +94,206 @@ body {
     }
   }
 }
+body {
+  assign {
+    expr {
+      sp_dataframe_collect {
+        case_sensitive: true
+        id {
+          bitfield1: 1
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 29
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 4
+    var_id {
+      bitfield1: 4
+    }
+  }
+}
+body {
+  eval {
+    uid: 5
+    var_id {
+      bitfield1: 4
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_dataframe_collect {
+        block: true
+        case_sensitive: true
+        id {
+          bitfield1: 1
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 31
+        }
+        statement_params {
+          _1: "SF_PARTNER"
+          _2: "FAKE_PARTNER"
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 6
+    var_id {
+      bitfield1: 6
+    }
+  }
+}
+body {
+  eval {
+    uid: 7
+    var_id {
+      bitfield1: 6
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_dataframe_collect {
+        id {
+          bitfield1: 1
+        }
+        log_on_exception: true
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 33
+        }
+        statement_params {
+          _1: "SF_PARTNER"
+          _2: "FAKE_PARTNER"
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 8
+    var_id {
+      bitfield1: 8
+    }
+  }
+}
+body {
+  eval {
+    uid: 9
+    var_id {
+      bitfield1: 8
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_dataframe_collect {
+        case_sensitive: true
+        id {
+          bitfield1: 1
+        }
+        no_wait: true
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 35
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 10
+    var_id {
+      bitfield1: 10
+    }
+  }
+}
+body {
+  eval {
+    uid: 11
+    var_id {
+      bitfield1: 10
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_dataframe_collect {
+        case_sensitive: true
+        id {
+          bitfield1: 1
+        }
+        no_wait: true
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 37
+        }
+        statement_params {
+          _1: "SF_PARTNER"
+          _2: "FAKE_PARTNER"
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 12
+    var_id {
+      bitfield1: 12
+    }
+  }
+}
+body {
+  eval {
+    uid: 13
+    var_id {
+      bitfield1: 12
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_dataframe_collect {
+        id {
+          bitfield1: 1
+        }
+        log_on_exception: true
+        no_wait: true
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 39
+        }
+        statement_params {
+          _1: "SF_PARTNER"
+          _2: "FAKE_PARTNER"
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 14
+    var_id {
+      bitfield1: 14
+    }
+  }
+}
+body {
+  eval {
+    uid: 15
+    var_id {
+      bitfield1: 14
+    }
+  }
+}
 client_ast_version: 1
 client_language {
   python_language {

diff --git a/tests/ast/data/DataFrame.count.test b/tests/ast/data/DataFrame.count.test
@@ -106,6 +106,39 @@ body {
     }
   }
 }
+body {
+  assign {
+    expr {
+      sp_dataframe_count {
+        id {
+          bitfield1: 1
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 31
+        }
+        statement_params {
+          _1: "SF_PARTNER"
+          _2: "FAKE_PARTNER"
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 6
+    var_id {
+      bitfield1: 6
+    }
+  }
+}
+body {
+  eval {
+    uid: 7
+    var_id {
+      bitfield1: 6
+    }
+  }
+}
 client_ast_version: 1
 client_language {
   python_language {

diff --git a/tests/ast/data/DataFrame.count2.test b/tests/ast/data/DataFrame.count2.test
@@ -135,6 +135,39 @@ body {
     }
   }
 }
+body {
+  assign {
+    expr {
+      sp_dataframe_count {
+        id {
+          bitfield1: 1
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 34
+        }
+        statement_params {
+          _1: "SF_PARTNER"
+          _2: "FAKE_PARTNER"
+        }
+      }
+    }
+    symbol {
+    }
+    uid: 6
+    var_id {
+      bitfield1: 6
+    }
+  }
+}
+body {
+  eval {
+    uid: 7
+    var_id {
+      bitfield1: 6
+    }
+  }
+}
 client_ast_version: 1
 client_language {
   python_language {