[DO NOT MERGE] [SNOW-1491199] merge phase0 server side

.github/workflows/precommit.yml

@@ -243,6 +243,152 @@ jobs:
             .tox/.coverage
             .tox/coverage.xml
 
+  test-ast-encoding:
+    name: Test AST expectations (phase0) py-${{ matrix.os }}-${{ matrix.python-version }}
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [macos-latest, windows-latest, ubuntu-latest]
+        # Protobuf encoding differs per python version, use 3.8 by default.
+        python-version: [ "3.8"]
+        cloud-provider: [aws]
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Decrypt parameters.py
+        shell: bash
+        run: .github/scripts/decrypt_parameters.sh
+        env:
+          PARAMETER_PASSWORD: ${{ secrets.PARAMETER_PASSWORD }}
+          CLOUD_PROVIDER: ${{ matrix.cloud-provider }}
+      - name: Download wheel(s)
+        uses: actions/download-artifact@v4
+        with:
+          name: wheel
+          path: dist
+      - name: Show wheels downloaded
+        run: ls -lh dist
+        shell: bash
+      - name: Upgrade setuptools, pip and wheel
+        run: python -m pip install -U setuptools pip wheel
+      - name: Install tox
+        run: python -m pip install tox
+      - name: Set up windows environment
+        if: runner.os == 'Windows'
+        run: |
+            "TZ=EST" >> $env:GITHUB_ENV
+            echo "TZ=EST" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
+      - name: Run tests
+        run: TZ=EST python -m tox -e "py${PYTHON_VERSION/\./}-ast"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short
+          TOX_PARALLEL_NO_SPINNER: 1
+          SNOWFLAKE_IS_PYTHON_RUNTIME_TEST: 1
+          TZ: EST
+        shell: bash
+
+  doctest-ast-phase0:
+    name: Doctest AST (phase0) py-${{ matrix.os }}-${{ matrix.python-version }}
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ ubuntu-latest ]
+        python-version: [ "3.9"]
+        cloud-provider: [ aws ]
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Decrypt parameters.py
+        shell: bash
+        run: .github/scripts/decrypt_parameters.sh
+        env:
+          PARAMETER_PASSWORD: ${{ secrets.PARAMETER_PASSWORD }}
+          CLOUD_PROVIDER: ${{ matrix.cloud-provider }}
+      - name: Download wheel(s)
+        uses: actions/download-artifact@v4
+        with:
+          name: wheel
+          path: dist
+      - name: Show wheels downloaded
+        run: ls -lh dist
+        shell: bash
+      - name: Upgrade setuptools, pip and wheel
+        run: python -m pip install -U setuptools pip wheel
+      - name: Install tox
+        run: python -m pip install tox
+      - name: Run doctests for AST (I/II)
+        run: python -m tox -e "py${PYTHON_VERSION}-doctest-notudf-ci"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short --enable-ast
+          TOX_PARALLEL_NO_SPINNER: 1
+        shell: bash
+
+  test-ast-phase0:
+    name: Test AST enabled (phase0) py-${{ matrix.os }}-${{ matrix.python-version }}
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ ubuntu-latest ]
+        python-version: [ "3.9", "3.11" ]
+        cloud-provider: [ aws ]
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Decrypt parameters.py
+        shell: bash
+        run: .github/scripts/decrypt_parameters.sh
+        env:
+          PARAMETER_PASSWORD: ${{ secrets.PARAMETER_PASSWORD }}
+          CLOUD_PROVIDER: ${{ matrix.cloud-provider }}
+      - name: Download wheel(s)
+        uses: actions/download-artifact@v4
+        with:
+          name: wheel
+          path: dist
+      - name: Show wheels downloaded
+        run: ls -lh dist
+        shell: bash
+      - name: Upgrade setuptools, pip and wheel
+        run: python -m pip install -U setuptools pip wheel
+      - name: Install tox
+        run: python -m pip install tox
+      - name: Run tests for AST (II/II)
+        run: python -m tox -e "py${PYTHON_VERSION/\./}-notdoctest-ci"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short --enable-ast
+          TOX_PARALLEL_NO_SPINNER: 1
+        shell: bash
+
   test-enable-cte-optimization:
     name: Test Enable CTE Optimization py-${{ matrix.os }}-${{ matrix.python-version }}-${{ matrix.cloud-provider }}
     needs: build

.pre-commit-config.yaml

@@ -1,4 +1,4 @@
-exclude: '^(.*egg.info.*|.*/parameters.py|docs/).*$'
+exclude: '^(.*egg.info.*|.*/parameters.py|.*/_vendored|docs/).*$'
 default_language_version:
     python: python3.9
 repos:
@@ -80,7 +80,10 @@ repos:
           - flake8-init-return == 1.0.0
           - flake8-print
         args:
-          - --per-file-ignores=tests/*.py:T201 # prints are allowed in test files
+          # prints are allowed in test files
+          - --per-file-ignores=tests/*.py:T201
+          # Ignore errors for protobuf generated file
+          - --exclude=src/_vendored,src/snowflake/snowpark/_internal/proto/ast_pb2.py,src/snowflake/snowpark/_internal/proto/scalapb_pb2.py
     # Use mypy for static type checking.
 -   repo: https://github.com/pre-commit/mirrors-mypy
     rev: 'v0.991'

scripts/tox_install_cmd.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 
-set -e
+set -euxo pipefail
 
 # The options are passed to us with spaces and we need to get rid
 # of those spaces when we pump it into pip install

setup.py

@@ -13,13 +13,18 @@
 MODIN_DEPENDENCY_VERSION = (
     "==0.28.1"  # Snowpark pandas requires modin 0.28.1, which depends on pandas 2.2.1
 )
-CONNECTOR_DEPENDENCY_VERSION = ">=3.10.0, <4.0.0"
+# Use HEAD of main branch in connector. This doesn't work with [pandas] extra.
+# CONNECTOR_DEPENDENCY = "snowflake-connector-python @ git+https://github.com/snowflakedb/snowflake-connector-python@main#egg=snowflake-connector-python"
+CONNECTOR_DEPENDENCY_VERSION = ">=3.12.0, <4.0.0"
+CONNECTOR_DEPENDENCY = f"snowflake-connector-python{CONNECTOR_DEPENDENCY_VERSION}"
 INSTALL_REQ_LIST = [
     "setuptools>=40.6.0",
     "wheel",
-    f"snowflake-connector-python{CONNECTOR_DEPENDENCY_VERSION}",
+    CONNECTOR_DEPENDENCY,
     # snowpark directly depends on typing-extension, so we should not remove it even if connector also depends on it.
     "typing-extensions>=4.1.0, <5.0.0",
+    "protobuf",
+    "tzlocal",
     "pyyaml",
     "cloudpickle>=1.6.0,<=2.2.1,!=2.1.0,!=2.2.0;python_version<'3.11'",
     "cloudpickle==2.2.1;python_version~='3.11'",  # backend only supports cloudpickle 2.2.1 + python 3.11 at the moment
@@ -31,6 +36,8 @@
 
 PANDAS_REQUIREMENTS = [
     f"snowflake-connector-python[pandas]{CONNECTOR_DEPENDENCY_VERSION}",
+    # When using HEAD of connector.
+    # f"{CONNECTOR_DEPENDENCY}[pandas]",
 ]
 MODIN_REQUIREMENTS = [
     *PANDAS_REQUIREMENTS,
@@ -39,6 +46,7 @@
 DEVELOPMENT_REQUIREMENTS = [
     "pytest<8.0.0",  # check SNOW-1022240 for more details on the pin here
     "pytest-cov",
+    "wrapt",
     "coverage",
     "sphinx==5.0.2",
     "cachetools",  # used in UDF doctest
@@ -47,6 +55,12 @@
     "openpyxl",  # used in read_excel test, not a requirement for distribution
     "matplotlib",  # used in plot tests
     "pre-commit",
+    "protoc-wheel-0",
+    "aiohttp",  # vcrpy requirements.
+    "boto",  # vcrpy requirements.
+    "httplib2",  # vcrpy requirements.
+    "httpx",  # vcrpy requirements.
+    "tornado",  # vcrpy requirements.
 ]
 
 # read the version
@@ -88,6 +102,7 @@
         "snowflake.snowpark._internal",
         "snowflake.snowpark._internal.analyzer",
         "snowflake.snowpark._internal.compiler",
+        "snowflake.snowpark._internal.proto",
         "snowflake.snowpark.mock",
         "snowflake.snowpark.modin",
         "snowflake.snowpark.modin.config",
@@ -115,6 +130,8 @@
         "modin": MODIN_REQUIREMENTS,
         "secure-local-storage": [
             f"snowflake-connector-python[secure-local-storage]{CONNECTOR_DEPENDENCY_VERSION}",
+            # When using HEAD, use this.
+            # f"{CONNECTOR_DEPENDENCY}[secure-local-storage]",
         ],
         "development": DEVELOPMENT_REQUIREMENTS,
         "modin-development": [

src/snowflake/snowpark/__init__.py

@@ -35,6 +35,7 @@
     "WhenNotMatchedClause",
     "QueryRecord",
     "QueryHistory",
+    "QueryListener",
     "AsyncJob",
 ]
 
@@ -53,7 +54,7 @@
 from snowflake.snowpark.dataframe_stat_functions import DataFrameStatFunctions
 from snowflake.snowpark.dataframe_writer import DataFrameWriter
 from snowflake.snowpark.file_operation import FileOperation, GetResult, PutResult
-from snowflake.snowpark.query_history import QueryHistory, QueryRecord
+from snowflake.snowpark.query_history import QueryHistory, QueryListener, QueryRecord
 from snowflake.snowpark.relational_grouped_dataframe import (
     GroupingSets,
     RelationalGroupedDataFrame,

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

@@ -63,6 +63,7 @@ def __init__(self, child: Optional["Expression"] = None) -> None:
         self.children = [child] if child else None
         self.datatype: Optional[DataType] = None
         self._cumulative_node_complexity: Optional[Dict[PlanNodeCategory, int]] = None
+        self._ast = None
 
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         # TODO: consider adding it to __init__ or use cached_property.
@@ -521,6 +522,9 @@ def __init__(self, expr: Expression, order_by_cols: List[Expression]) -> None:
         self.expr = expr
         self.order_by_cols = order_by_cols
         self.datatype = expr.datatype
+        assert all(
+            isinstance(order_by_col, Expression) for order_by_col in order_by_cols
+        )
 
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr, *self.order_by_cols)
@@ -591,13 +595,15 @@ def __init__(
         datatype: DataType,
         nullable: bool = True,
         api_call_source: Optional[str] = None,
+        is_aggregate_function: bool = False,
     ) -> None:
         super().__init__()
         self.udf_name = udf_name
         self.children = children
         self.datatype = datatype
         self.nullable = nullable
         self.api_call_source = api_call_source
+        self.is_aggregate_function = is_aggregate_function
 
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(*self.children)

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

@@ -88,7 +88,7 @@ def __init__(
         self.path = path
         self.outer = outer
         self.recursive = recursive
-        self.mode = mode
+        self.mode = mode.upper()
 
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]: