[SNOW-1465372] Strip version numbers from the Snowpark lib during cod…

…egen (#1173)
snowflakedb · Jun 7, 2024 · 641347d · 641347d
1 parent be72cb9
commit 641347d
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 19 deletions.
diff --git a/src/snowflake/cli/plugins/nativeapp/codegen/snowpark/python_processor.py b/src/snowflake/cli/plugins/nativeapp/codegen/snowpark/python_processor.py
@@ -43,7 +43,10 @@
 DEFAULT_TIMEOUT = 30
 TEMPLATE_PATH = Path(__file__).parent / "callback_source.py.jinja"
 SNOWPARK_LIB_NAME = "snowflake-snowpark-python"
-SNOWPARK_LIB_REGEX = rf"'{SNOWPARK_LIB_NAME}\s*((<|<=|!=|==|>=|>|~=|===)\s*[a-zA-Z0-9_.*+!-]+)?'"  # support PEP 508, even though not all of it is supported in Snowflake yet
+SNOWPARK_LIB_REGEX = re.compile(
+    # support PEP 508, even though not all of it is supported in Snowflake yet
+    rf"'{SNOWPARK_LIB_NAME}\s*((<|<=|!=|==|>=|>|~=|===)\s*[a-zA-Z0-9_.*+!-]+)?'"
+)
 STAGE_PREFIX = "@"
 
 
@@ -266,14 +269,12 @@ def _normalize(
         # If user defined their udf as @udf(lambda: x, ...) then extension_fn.handler is <lambda>.
         extension_fn.handler = f"{py_file.stem}.{extension_fn.handler}"
 
-        snowpark_lib_found = False
-        snowpark_lib_regex = re.compile(SNOWPARK_LIB_REGEX)
-        for pkg in extension_fn.packages:
-            if snowpark_lib_regex.fullmatch(ensure_string_literal(pkg.strip())):
-                snowpark_lib_found = True
-                break
-        if not snowpark_lib_found:
-            extension_fn.packages.append(SNOWPARK_LIB_NAME)
+        extension_fn.packages = [
+            self._normalize_package_name(pkg) for pkg in extension_fn.packages
+        ]
+        snowpark_lib_name = ensure_string_literal(SNOWPARK_LIB_NAME)
+        if snowpark_lib_name not in extension_fn.packages:
+            extension_fn.packages.append(snowpark_lib_name)
 
         if extension_fn.imports is None:
             extension_fn.imports = []
@@ -283,6 +284,23 @@ def _normalize(
             deploy_root=deploy_root,
         )
 
+    def _normalize_package_name(self, pkg: str) -> str:
+        """
+        Returns a normalized version of the provided package name, as a Snowflake SQL string literal. Since the
+        Snowpark library can sometimes add a spurious version to its own package name, we strip this here too so
+        that the native application does not accidentally rely on stale packages once the snowpark library is updated
+        in the cloud.
+
+        Args:
+            pkg (str): The package name to normalize.
+        Returns:
+            A normalized version of the package name, as a Snowflake SQL string literal.
+        """
+        normalized_package_name = ensure_string_literal(pkg.strip())
+        if SNOWPARK_LIB_REGEX.fullmatch(normalized_package_name):
+            return ensure_string_literal(SNOWPARK_LIB_NAME)
+        return normalized_package_name
+
     def collect_extension_functions(
         self, bundle_map: BundleMap, processor_mapping: Optional[ProcessorMapping]
     ) -> Dict[Path, List[NativeAppExtensionFunction]]:

diff --git a/tests/nativeapp/codegen/snowpark/__snapshots__/test_python_processor.ambr b/tests/nativeapp/codegen/snowpark/__snapshots__/test_python_processor.ambr
@@ -181,7 +181,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python===0.15.0-rc1')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -203,7 +203,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python===0.15.0-rc1')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -269,7 +269,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python  ==  0.15.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -335,7 +335,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python==0.15.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -357,7 +357,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python==0.15.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -379,7 +379,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python<0.16.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -401,7 +401,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python<=0.17.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -423,7 +423,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python>=0.14.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -445,7 +445,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python>0.13.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'
@@ -467,7 +467,7 @@
   LANGUAGE PYTHON
   RUNTIME_VERSION=3.11
   IMPORTS=('/path/to/import1.py', '/path/to/import2.zip', '/stagepath/main.py')
-  PACKAGES=('snowflake-snowpark-python===0.15.0')
+  PACKAGES=('snowflake-snowpark-python')
   EXTERNAL_ACCESS_INTEGRATIONS=('integration_one', 'integration_two')
   SECRETS=('key1'=secret_one, 'key2'=integration_two)
   HANDLER='main.my_function_handler'