New scripting environment for integration with the monorepo (#2748)

THIS IS DEFINITELY WIP!

I've screwed up my Cloud Workspace trying to make this work, so I
haven't really tested scripts/copy-ast.sh.

I did test scripts/copy-remote-ast.sh whereby I run the script from my
local machine. AST tests pass after this.

---------

Co-authored-by: Varnika Budati <varnika.budati@snowflake.com>

Author: sfc-gh-azwiegincew

scripts/copy-ast.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+# Compute the source directory. We honor the MONOREPO_DIR environment variable if set,
+# and default to the standard Cloud Workspace location otherwise.
+DEFAULT_SRC_DIR=$HOME/Snowflake/trunk
+MONOREPO_DIR="${MONOREPO_DIR:-$DEFAULT_SRC_DIR}"
+if [ ! -d "$MONOREPO_DIR" ]; then
+  echo "Source directory '${MONOREPO_DIR}' doesn't exist"
+  exit 1
+fi
+
+# To allow this script to run from any subdirectory within snowpark-python, we use git rev-parse.
+SNOWPARK_ROOT=$(git rev-parse --show-toplevel)
+
+# Build the targets.
+pushd $MONOREPO_DIR
+bazel build //Snowpark/ast:ast_proto
+bazel build //Snowpark/unparser
+popd
+
+# Copy the AST.
+cp $MONOREPO_DIR/bazel-bin/Snowpark/ast/ast.proto $SNOWPARK_ROOT/src/snowflake/snowpark/_internal/proto/ast.proto
+
+# Generate Python from ast.proto.
+pushd $SNOWPARK_ROOT
+python -m tox -e protoc
+popd

scripts/copy-remote-ast.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+# This script assumes the target Cloud Workspace specified as the command-line argument has the build target.
+# To make sure this is the case, run bazel build //Snowpark/ast:ast_proto and bazel build //Snowpark/unparser.
+
+# N.B. The calling environment further requires:
+# export MONOREPO_DIR=$TMPDIR
+
+# To allow this script to run from any subdirectory within snowpark-python, we use git rev-parse.
+SNOWPARK_ROOT=$(git rev-parse --show-toplevel)
+
+if [ ! -d "$TMPDIR" ]; then
+  echo "TMPDIR not defined"
+  exit 1
+fi
+
+scp $1:~/Snowflake/trunk/bazel-bin/Snowpark/ast/ast.proto $SNOWPARK_ROOT/src/snowflake/snowpark/_internal/proto/ast.proto
+
+mkdir -p $TMPDIR/bazel-bin/Snowpark/unparser/unparser.runfiles
+scp -r $1:~/Snowflake/trunk/bazel-bin/Snowpark/unparser/unparser.runfiles/ $TMPDIR/bazel-bin/Snowpark/unparser/unparser.runfiles/
+
+scp $1:~/Snowflake/trunk/bazel-bin/Snowpark/unparser/unparser-lib.jar $TMPDIR/bazel-bin/Snowpark/unparser/
+scp $1:~/Snowflake/trunk/bazel-bin/Snowpark/unparser/unparser.jar $TMPDIR/bazel-bin/Snowpark/unparser/
+
+pushd $SNOWPARK_ROOT
+python -m tox -e protoc
+popd

src/snowflake/snowpark/_internal/proto/update-from-devvm.sh

@@ -35,8 +35,8 @@ pytest --update-expectations tests/ast
 For these tests to work, the Unparser must be built in the monorepo:
 ```bash
 cd my-monorepo-path
-cd Snowflake/unparser
-sbt assembly
+bazel build //Snowpark/unparser
 ```
 
-The location of the Unparser can be set either via the environment variable `SNOWPARK_UNPARSER_JAR` or via the _pytest_ commandline argument `--unparser-jar=<path>`.
+The location of the monorepo must be supplied in the environment variable `MONOREPO_DIR`.
+Alternatively, use `scripts/copy-remote-ast.sh` to copy the necessary artifacts from a cloud workspace.

src/snowflake/snowpark/_internal/proto/update-from-local.sh

@@ -33,7 +33,7 @@ def render(ast_base64: Union[str, List[str]], unparser_jar: Optional[str]) -> st
     """Uses the unparser to render the AST."""
     assert (
         unparser_jar
-    ), "A valid Unparser JAR path must be supplied either via --unparser-jar=<path> or the environment variable SNOWPARK_UNPARSER_JAR"
+    ), "A valid Unparser JAR path must be supplied either via --unparser-jar=<path> or the environment variable MONOREPO_DIR"
 
     if isinstance(ast_base64, str):
         ast_base64 = [ast_base64]
@@ -43,7 +43,7 @@ def render(ast_base64: Union[str, List[str]], unparser_jar: Optional[str]) -> st
             "java",
             "-cp",
             unparser_jar,
-            "com.snowflake.snowpark.experimental.unparser.UnparserCli",
+            "com.snowflake.snowpark.unparser.UnparserCli",
             ",".join(
                 ast_base64
             ),  # base64 strings will not contain , so pass multiple batches comma-separated.

tests/ast/README.md

@@ -11,9 +11,18 @@
 
 
 def default_unparser_path():
-    explicit = os.getenv("SNOWPARK_UNPARSER_JAR")
-    default_default = f"{os.getenv('HOME')}/Snowflake/trunk/Snowpark/unparser/target/scala-2.13/unparser-assembly-0.1.jar"
-    return explicit or default_default
+    explicit = os.getenv("MONOREPO_DIR")
+    default_default = os.path.join(os.getenv("HOME"), "Snowflake/trunk")
+    base_dir = explicit or default_default
+    unparser_dir = os.path.join(base_dir, "bazel-bin/Snowpark/unparser")
+
+    # Grab all *.jar files from the subtree.
+    jars = []
+    for path, _, files in os.walk(unparser_dir):
+        jars.extend(
+            [os.path.join(path, file) for file in files if file.endswith(".jar")]
+        )
+    return ":".join(jars)
 
 
 def pytest_addoption(parser):
@@ -22,7 +31,7 @@ def pytest_addoption(parser):
         action="store",
         default=default_unparser_path(),
         type=str,
-        help="Path to the Unparser JAR built in the monorepo. To build it, run `sbt assembly` from the unparser directory.",
+        help="Path to the Unparser JAR built in the monorepo.",
     )
     parser.addoption(
         "--update-expectations",
@@ -33,15 +42,18 @@ def pytest_addoption(parser):
 
 
 def pytest_configure(config):
-    pytest.unparser_jar = config.getoption("--unparser-jar")
-    if not os.path.exists(pytest.unparser_jar):
-        pytest.unparser_jar = None
+    unparser_jar = config.getoption("--unparser-jar")
+    pytest.unparser_jar = (
+        unparser_jar
+        if all(os.path.exists(file) for file in unparser_jar.split(":"))
+        else None
+    )
     pytest.update_expectations = config.getoption("--update-expectations")
 
     if pytest.unparser_jar is None and pytest.update_expectations:
         raise RuntimeError(
-            f"Unparser JAR not found at {pytest.unparser_jar}. "
-            f"Please set the correct path with --unparser-jar or SNOWPARK_UNPARSER_JAR."
+            f"Unparser JAR not found at {unparser_jar}. "
+            f"Please set the correct path with --unparser-jar or the MONOREPO_DIR environment variable."
         )
 
 

tests/ast/ast_test_utils.py

@@ -193,7 +193,7 @@ def test_ast(session, tables, test_case):
     if pytest.update_expectations:
         assert pytest.unparser_jar, (
             "Can only update expectations with unparser jar set. Either run the test with --unparser-jar=<path> or"
-            " update the environment variable SNOWPARK_UNPARSER_JAR."
+            " update the environment variable MONOREPO_DIR."
         )
         with open(DATA_DIR / test_case.filename, "w", encoding="utf-8") as f:
             f.writelines(

tests/ast/conftest.py

@@ -12,6 +12,7 @@
 import pytest
 
 from snowflake.snowpark._internal.utils import COMPATIBLE_WITH_MODIN, warning_dict
+from .ast.conftest import default_unparser_path
 
 logging.getLogger("snowflake.connector").setLevel(logging.ERROR)
 
@@ -45,12 +46,6 @@ def dict_exporter():
     opentelemetry_installed = False
 
 
-def default_unparser_path():
-    explicit = os.getenv("SNOWPARK_UNPARSER_JAR")
-    default_default = f"{os.getenv('HOME')}/Snowflake/trunk/Snowpark/unparser/target/scala-2.13/unparser-assembly-0.1.jar"
-    return explicit or default_default
-
-
 def is_excluded_frontend_file(path):
     for excluded in excluded_frontend_files:
         if str(path).endswith(excluded):
@@ -76,7 +71,7 @@ def pytest_addoption(parser, pluginmanager):
         action="store",
         default=default_unparser_path(),
         type=str,
-        help="Path to the Unparser JAR built in the monorepo. To build it, run `sbt assembly` from the unparser directory.",
+        help="Path to the Unparser JAR built in the monorepo.",
     )