snowflake-ml-python 1.5.3__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/ml/feature_store/feature_view.py

@@ -2,17 +2,25 @@ from __future__ import annotations
 
 import json
 import re
+import warnings
 from collections import OrderedDict
 from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
+from snowflake.ml._internal.exceptions import (
+    error_codes,
+    exceptions as snowml_exceptions,
+)
+from snowflake.ml._internal.utils import identifier
 from snowflake.ml._internal.utils.identifier import concat_names
 from snowflake.ml._internal.utils.sql_identifier import (
     SqlIdentifier,
     to_sql_identifiers,
 )
+from snowflake.ml.feature_store import feature_store
 from snowflake.ml.feature_store.entity import Entity
+from snowflake.ml.lineage import lineage_node
 from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark.types import (
     DateType,
@@ -67,6 +75,7 @@ class FeatureViewVersion(str):
 
 
 class FeatureViewStatus(Enum):
+    MASKED = "MASKED"  # for shared feature views where scheduling state is not available
     DRAFT = "DRAFT"
     STATIC = "STATIC"
     RUNNING = "RUNNING"  # This can be deprecated after BCR 2024_02 gets fully deployed
@@ -107,7 +116,7 @@ class FeatureViewSlice:
         return cls(**json_dict)
 
 
-class FeatureView:
+class FeatureView(lineage_node.LineageNode):
     """
     A FeatureView instance encapsulates a logical group of features.
     """
@@ -117,9 +126,11 @@ class FeatureView:
         name: str,
         entities: List[Entity],
         feature_df: DataFrame,
+        *,
         timestamp_col: Optional[str] = None,
         refresh_freq: Optional[str] = None,
         desc: str = "",
+        warehouse: Optional[str] = None,
         **_kwargs: Any,
     ) -> None:
         """
@@ -140,7 +151,33 @@ class FeatureView:
                 NOTE: If refresh_freq is not provided, then FeatureView will be registered as View on Snowflake backend
                     and there won't be extra storage cost.
             desc: description of the FeatureView.
+            warehouse: warehouse to refresh feature view. Not needed for static feature view (refresh_freq is None).
+                For managed feature view, this warehouse will overwrite the default warehouse of Feature Store if it is
+                specified, otherwise the default warehouse will be used.
             _kwargs: reserved kwargs for system generated args. NOTE: DO NOT USE.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> # draft_fv is a local object that hasn't materiaized to Snowflake backend yet.
+            >>> feature_df = session.sql("select f_1, f_2 from source_table")
+            >>> draft_fv = FeatureView(
+            ...     name="my_fv",
+            ...     entities=[e1, e2],
+            ...     feature_df=feature_df,
+            ...     timestamp_col='TS', # optional
+            ...     refresh_freq='1d',  # optional
+            ...     desc='A line about this feature view',  # optional
+            ...     warehouse='WH'      # optional, the warehouse used to refresh (managed) feature view
+            ... )
+            >>> print(draft_fv.status)
+            FeatureViewStatus.DRAFT
+            <BLANKLINE>
+            >>> # registered_fv is a local object that maps to a Snowflake backend object.
+            >>> registered_fv = fs.register_feature_view(draft_fv, "v1")
+            >>> print(registered_fv.status)
+            FeatureViewStatus.ACTIVE
+
         """
 
         self._name: SqlIdentifier = SqlIdentifier(name)
@@ -158,7 +195,7 @@ class FeatureView:
         self._refresh_freq: Optional[str] = refresh_freq
         self._database: Optional[SqlIdentifier] = None
         self._schema: Optional[SqlIdentifier] = None
-        self._warehouse: Optional[SqlIdentifier] = None
+        self._warehouse: Optional[SqlIdentifier] = SqlIdentifier(warehouse) if warehouse is not None else None
         self._refresh_mode: Optional[str] = None
         self._refresh_mode_reason: Optional[str] = None
         self._owner: Optional[str] = None
@@ -176,6 +213,33 @@ class FeatureView:
 
         Raises:
             ValueError: if selected feature names is not found in the FeatureView.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> # feature_df contains 3 features and 1 entity
+            >>> feature_df = session.table(source_table).select(
+            ...     'TRIPDURATION',
+            ...     'START_STATION_LATITUDE',
+            ...     'END_STATION_LONGITUDE',
+            ...     'TRIP_ID'
+            ... )
+            >>> darft_fv = FeatureView(name='F_TRIP', entities=[e], feature_df=feature_df)
+            >>> fv = fs.register_feature_view(darft_fv, version='1.0')
+            >>> # shows all 3 features
+            >>> fv.feature_names
+            ['TRIPDURATION', 'START_STATION_LATITUDE', 'END_STATION_LONGITUDE']
+            <BLANKLINE>
+            >>> # slice a subset of features
+            >>> fv_slice = fv.slice(['TRIPDURATION', 'START_STATION_LATITUDE'])
+            >>> fv_slice.names
+            ['TRIPDURATION', 'START_STATION_LATITUDE']
+            <BLANKLINE>
+            >>> # query the full set of features in original feature view
+            >>> fv_slice.feature_view_ref.feature_names
+            ['TRIPDURATION', 'START_STATION_LATITUDE', 'END_STATION_LONGITUDE']
+
         """
 
         res = []
@@ -187,14 +251,30 @@ class FeatureView:
         return FeatureViewSlice(self, res)
 
     def fully_qualified_name(self) -> str:
-        """Returns the fully qualified name (<database_name>.<schema_name>.<feature_view_name>) for the
-            FeatureView in Snowflake.
+        """
+        Returns the fully qualified name (<database_name>.<schema_name>.<feature_view_name>) for the
+        FeatureView in Snowflake.
 
         Returns:
             fully qualified name string.
 
         Raises:
             RuntimeError: if the FeatureView is not registered.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> feature_df = session.table(source_table).select(
+            ...     'TRIPDURATION',
+            ...     'START_STATION_LATITUDE',
+            ...     'TRIP_ID'
+            ... )
+            >>> darft_fv = FeatureView(name='F_TRIP', entities=[e], feature_df=feature_df)
+            >>> registered_fv = fs.register_feature_view(darft_fv, version='1.0')
+            >>> registered_fv.fully_qualified_name()
+            'MY_DB.MY_SCHEMA."F_TRIP$1.0"'
+
         """
         if self.status == FeatureViewStatus.DRAFT or self.version is None:
             raise RuntimeError(f"FeatureView {self.name} has not been registered.")
@@ -212,6 +292,22 @@ class FeatureView:
 
         Raises:
             ValueError: if feature name is not found in the FeatureView.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> feature_df = session.table(source_table).select('TRIPDURATION', 'START_STATION_LATITUDE', 'TRIP_ID')
+            >>> draft_fv = FeatureView(name='F_TRIP', entities=[e], feature_df=feature_df)
+            >>> draft_fv = draft_fv.attach_feature_desc({
+            ...     "TRIPDURATION": "Duration of a trip.",
+            ...     "START_STATION_LATITUDE": "Latitude of the start station."
+            ... })
+            >>> registered_fv = fs.register_feature_view(draft_fv, version='1.0')
+            >>> registered_fv.feature_descs
+            OrderedDict([('TRIPDURATION', 'Duration of a trip.'),
+                ('START_STATION_LATITUDE', 'Latitude of the start station.')])
+
         """
         for f, d in descs.items():
             f = SqlIdentifier(f)
@@ -243,6 +339,41 @@ class FeatureView:
     def desc(self) -> str:
         return self._desc
 
+    @desc.setter
+    def desc(self, new_value: str) -> None:
+        """Set the description of feature view.
+
+        Args:
+            new_value: new value of description.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> darft_fv = FeatureView(
+            ...     name='F_TRIP',
+            ...     entities=[e],
+            ...     feature_df=feature_df,
+            ...     desc='old desc'
+            ... )
+            >>> fv_1 = fs.register_feature_view(darft_fv, version='1.0')
+            >>> print(fv_1.desc)
+            old desc
+            <BLANKLINE>
+            >>> darft_fv.desc = 'NEW DESC'
+            >>> fv_2 = fs.register_feature_view(darft_fv, version='2.0')
+            >>> print(fv_2.desc)
+            NEW DESC
+
+        """
+        warnings.warn(
+            "You must call register_feature_view() to make it effective. "
+            "Or use update_feature_view(desc=<new_value>).",
+            stacklevel=2,
+            category=UserWarning,
+        )
+        self._desc = new_value
+
     @property
     def query(self) -> str:
         return self._query
@@ -269,10 +400,37 @@ class FeatureView:
 
     @refresh_freq.setter
     def refresh_freq(self, new_value: str) -> None:
-        if self.status == FeatureViewStatus.DRAFT or self.status == FeatureViewStatus.STATIC:
-            raise RuntimeError(
-                f"Feature view {self.name}/{self.version} must be registered and non-static to update refresh_freq."
-            )
+        """Set refresh frequency of feature view.
+
+        Args:
+            new_value: The new value of refresh frequency.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> darft_fv = FeatureView(
+            ...     name='F_TRIP',
+            ...     entities=[e],
+            ...     feature_df=feature_df,
+            ...     refresh_freq='1d'
+            ... )
+            >>> fv_1 = fs.register_feature_view(darft_fv, version='1.0')
+            >>> print(fv_1.refresh_freq)
+            1 day
+            <BLANKLINE>
+            >>> darft_fv.refresh_freq = '12h'
+            >>> fv_2 = fs.register_feature_view(darft_fv, version='2.0')
+            >>> print(fv_2.refresh_freq)
+            12 hours
+
+        """
+        warnings.warn(
+            "You must call register_feature_view() to make it effective. "
+            "Or use update_feature_view(refresh_freq=<new_value>).",
+            stacklevel=2,
+            category=UserWarning,
+        )
         self._refresh_freq = new_value
 
     @property
@@ -289,10 +447,38 @@ class FeatureView:
 
     @warehouse.setter
     def warehouse(self, new_value: str) -> None:
-        if self.status == FeatureViewStatus.DRAFT or self.status == FeatureViewStatus.STATIC:
-            raise RuntimeError(
-                f"Feature view {self.name}/{self.version} must be registered and non-static to update warehouse."
-            )
+        """Set warehouse of feature view.
+
+        Args:
+            new_value: The new value of warehouse.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = fs.get_entity('TRIP_ID')
+            >>> darft_fv = FeatureView(
+            ...     name='F_TRIP',
+            ...     entities=[e],
+            ...     feature_df=feature_df,
+            ...     refresh_freq='1d',
+            ...     warehouse='WH1',
+            ... )
+            >>> fv_1 = fs.register_feature_view(darft_fv, version='1.0')
+            >>> print(fv_1.warehouse)
+            WH1
+            <BLANKLINE>
+            >>> darft_fv.warehouse = 'WH2'
+            >>> fv_2 = fs.register_feature_view(darft_fv, version='2.0')
+            >>> print(fv_2.warehouse)
+            WH2
+
+        """
+        warnings.warn(
+            "You must call register_feature_view() to make it effective. "
+            "Or use update_feature_view(warehouse=<new_value>).",
+            stacklevel=2,
+            category=UserWarning,
+        )
         self._warehouse = SqlIdentifier(new_value)
 
     @property
@@ -406,6 +592,11 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             feature_desc_dict[k.identifier()] = v
         fv_dict["_feature_desc"] = feature_desc_dict
 
+        lineage_node_keys = [key for key in fv_dict if key.startswith("_node") or key == "_session"]
+
+        for key in lineage_node_keys:
+            fv_dict.pop(key)
+
         return fv_dict
 
     def to_df(self, session: Session) -> DataFrame:
@@ -428,7 +619,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
 
         entities = []
         for e_json in json_dict["_entities"]:
-            e = Entity(e_json["name"], e_json["join_keys"], e_json["desc"])
+            e = Entity(e_json["name"], e_json["join_keys"], desc=e_json["desc"])
             e.owner = e_json["owner"]
             entities.append(e)
 
@@ -449,6 +640,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             refresh_mode_reason=json_dict["_refresh_mode_reason"],
             owner=json_dict["_owner"],
             infer_schema_df=session.sql(json_dict.get("_infer_schema_query", None)),
+            session=session,
         )
 
     @staticmethod
@@ -463,6 +655,21 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             )
         )
 
+    @staticmethod
+    def _load_from_lineage_node(session: Session, name: str, version: str) -> FeatureView:
+        db_name, feature_store_name, feature_view_name, _ = identifier.parse_schema_level_object_identifier(name)
+
+        session_warehouse = session.get_current_warehouse()
+
+        if not session_warehouse:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.NOT_FOUND,
+                original_exception=ValueError("No active warehouse selected in the current session"),
+            )
+
+        fs = feature_store.FeatureStore(session, db_name, feature_store_name, default_warehouse=session_warehouse)
+        return fs.get_feature_view(feature_view_name, version)  # type: ignore[no-any-return]
+
     @staticmethod
     def _construct_feature_view(
         name: str,
@@ -481,6 +688,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         refresh_mode_reason: Optional[str],
         owner: Optional[str],
         infer_schema_df: Optional[DataFrame],
+        session: Session,
     ) -> FeatureView:
         fv = FeatureView(
             name=name,
@@ -500,4 +708,11 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         fv._refresh_mode_reason = refresh_mode_reason
         fv._owner = owner
         fv.attach_feature_desc(feature_descs)
+
+        lineage_node.LineageNode.__init__(
+            fv, session=session, name=f"{fv.database}.{fv._schema}.{name}", domain="feature_view", version=version
+        )
         return fv
+
+
+lineage_node.DOMAIN_LINEAGE_REGISTRY["feature_view"] = FeatureView

snowflake/ml/fileset/embedded_stage_fs.py

@@ -11,11 +11,17 @@ from snowflake.ml._internal.exceptions import (
     fileset_errors,
 )
 from snowflake.ml._internal.utils import identifier
+from snowflake.ml.fileset import stage_fs
 from snowflake.snowpark import exceptions as snowpark_exceptions
 
-from . import stage_fs
-
-_SNOWURL_PATH_RE = re.compile(r"versions/(?P<version>[^/]+)(?:/+(?P<filepath>.*))?")
+PROTOCOL_NAME = "snow"
+_SNOWURL_ENTITY_PATTERN = (
+    f"(?:{PROTOCOL_NAME}://)?"
+    r"(?<!@)(?P<domain>\w+)/"
+    rf"(?P<name>(?:{identifier._SF_IDENTIFIER}\.){{,2}}{identifier._SF_IDENTIFIER})/"
+)
+_SNOWURL_VERSION_PATTERN = r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/+(?P<relpath>.*))?)?)"
+_SNOWURL_PATH_RE = re.compile(f"(?:{_SNOWURL_ENTITY_PATTERN})?" + _SNOWURL_VERSION_PATTERN)
 
 
 class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
@@ -76,8 +82,8 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
         versions_dict = defaultdict(list)
         for file in files:
             match = _SNOWURL_PATH_RE.fullmatch(file)
-            assert match is not None and match.group("filepath") is not None
-            versions_dict[match.group("version")].append(match.group("filepath"))
+            assert match is not None and match.group("relpath") is not None
+            versions_dict[match.group("version")].append(match.group("relpath"))
         try:
             async_jobs: List[snowpark.AsyncJob] = []
             for version, version_files in versions_dict.items():
@@ -98,10 +104,8 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
                 (r["NAME"], r["URL"]) for job in async_jobs for r in stage_fs._resolve_async_job(job)
             ]
             return presigned_urls
-        except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST) or e.message.startswith(
-                fileset_errors.ERRNO_STAGE_NOT_EXIST
-            ):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code in {fileset_errors.ERRNO_DOMAIN_NOT_EXIST, fileset_errors.ERRNO_STAGE_NOT_EXIST}:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_NOT_FOUND,
                     original_exception=fileset_errors.StageNotFoundError(
@@ -118,7 +122,7 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
     def _parent(cls, path: str) -> str:
         """Get parent of specified path up to minimally valid root path.
 
-        For SnowURL, the minimum valid path is snow://<domain>/<entity>/versions/<version>
+        For SnowURL, the minimum valid relative path is versions/<version>
 
         Args:
             path: File or directory path
@@ -128,22 +132,22 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
 
         Examples:
         ----
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version/file.ext")
-        "snow://dataset/my_ds/versions/my_version/"
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version/subdir/file.ext")
-        "snow://dataset/my_ds/versions/my_version/subdir/"
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version/")
-        "snow://dataset/my_ds/versions/my_version/"
-        >>> fs._parent("snow://dataset/my_ds/versions/my_version")
-        "snow://dataset/my_ds/versions/my_version"
+        >>> fs._parent("versions/my_version/file.ext")
+        "versions/my_version"
+        >>> fs._parent("versions/my_version/subdir/file.ext")
+        "versions/my_version/subdir"
+        >>> fs._parent("versions/my_version/")
+        "versions/my_version"
+        >>> fs._parent("versions/my_version")
+        "versions/my_version"
         """
         path_match = _SNOWURL_PATH_RE.fullmatch(path)
         if not path_match:
             return super()._parent(path)  # type: ignore[no-any-return]
-        filepath: str = path_match.group("filepath") or ""
-        root: str = path[: path_match.start("filepath")] if filepath else path
+        filepath: str = path_match.group("relpath") or ""
+        root: str = path[: path_match.start("relpath")] if filepath else path
         if "/" in filepath:
             parent = filepath.rsplit("/", 1)[0]
             return root + parent
         else:
-            return root
+            return root.rstrip("/")

snowflake/ml/fileset/fileset.py

@@ -256,9 +256,9 @@ class FileSet:
                     api_calls=[snowpark.DataFrameWriter.copy_into_location],
                 ),
             )
-        except snowpark_exceptions.SnowparkClientException as e:
+        except snowpark_exceptions.SnowparkSQLException as e:
             # Snowpark wraps the Python Connector error code in the head of the error message.
-            if e.message.startswith(fileset_errors.ERRNO_FILE_EXIST_IN_STAGE):
+            if e.sql_error_code == fileset_errors.ERRNO_FILE_EXIST_IN_STAGE:
                 raise fileset_errors.FileSetExistError(fileset_error_messages.FILESET_ALREADY_EXISTS.format(name))
             else:
                 raise fileset_errors.FileSetError(str(e))

snowflake/ml/fileset/snowfs.py

@@ -14,18 +14,10 @@ from snowflake.ml._internal.exceptions import (
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.fileset import embedded_stage_fs, sfcfs
 
-PROTOCOL_NAME = "snow"
-
 _SFFileEntityPath = collections.namedtuple(
     "_SFFileEntityPath", ["domain", "name", "filepath", "version", "relative_path"]
 )
-_PROJECT = "FileSet"
-_SNOWURL_PATTERN = re.compile(
-    f"({PROTOCOL_NAME}://)?"
-    r"(?<!@)(?P<domain>\w+)/"
-    rf"(?P<name>(?:{identifier._SF_IDENTIFIER}\.){{,2}}{identifier._SF_IDENTIFIER})/"
-    r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)"
-)
+_SNOWURL_PATTERN = re.compile(embedded_stage_fs._SNOWURL_ENTITY_PATTERN + embedded_stage_fs._SNOWURL_VERSION_PATTERN)
 
 
 class SnowFileSystem(sfcfs.SFFileSystem):
@@ -38,7 +30,7 @@ class SnowFileSystem(sfcfs.SFFileSystem):
     See `sfcfs.SFFileSystem` documentation for example usage patterns.
     """
 
-    protocol = PROTOCOL_NAME
+    protocol = embedded_stage_fs.PROTOCOL_NAME
     _IS_BUGGED_VERSION = None
 
     def __init__(
@@ -75,10 +67,7 @@ class SnowFileSystem(sfcfs.SFFileSystem):
         """Convert the relative path in a stage to an absolute path starts with the location of the stage."""
         # Strip protocol from absolute path, since backend needs snow:// prefix to resolve correctly
         # but fsspec logic strips protocol when doing any searching and globbing
-        stage_name = stage_fs.stage_name
-        protocol = f"{PROTOCOL_NAME}://"
-        if stage_name.startswith(protocol):
-            stage_name = stage_name[len(protocol) :]
+        stage_name: str = self._strip_protocol(stage_fs.stage_name)
         abs_path = stage_name + "/" + path
         return abs_path
 
@@ -128,4 +117,4 @@ class SnowFileSystem(sfcfs.SFFileSystem):
             )
 
 
-fsspec.register_implementation(PROTOCOL_NAME, SnowFileSystem)
+fsspec.register_implementation(SnowFileSystem.protocol, SnowFileSystem)

snowflake/ml/fileset/stage_fs.py

@@ -170,8 +170,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
             path = path.lstrip("/")
             async_job: snowpark.AsyncJob = self._session.sql(f"LIST '{loc}/{path}'").collect(block=False)
             objects: List[snowpark.Row] = _resolve_async_job(async_job)
-        except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code == fileset_errors.ERRNO_DOMAIN_NOT_EXIST:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_NOT_FOUND,
                     original_exception=fileset_errors.StageNotFoundError(
@@ -234,21 +234,29 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
 
         Raises:
             SnowflakeMLException: An error occurred when the given path points to a file that cannot be found.
+            snowpark_exceptions.SnowparkClientException: File access failed with a Snowpark exception
         """
         path = path.lstrip("/")
         if self._USE_FALLBACK_FILE_ACCESS:
             return self._open_with_snowpark(path)
         cached_presigned_url = self._url_cache.get(path, None)
-        if not cached_presigned_url:
-            res = self._fetch_presigned_urls([path])
-            url = res[0][1]
-            expire_at = time.time() + _PRESIGNED_URL_LIFETIME_SEC
-            cached_presigned_url = _PresignedUrl(url, expire_at)
-            self._url_cache[path] = cached_presigned_url
-            logging.debug(f"Retrieved presigned url for {path}.")
-        elif cached_presigned_url.is_expiring():
-            self.optimize_read()
-            cached_presigned_url = self._url_cache[path]
+        try:
+            if not cached_presigned_url:
+                res = self._fetch_presigned_urls([path])
+                url = res[0][1]
+                expire_at = time.time() + _PRESIGNED_URL_LIFETIME_SEC
+                cached_presigned_url = _PresignedUrl(url, expire_at)
+                self._url_cache[path] = cached_presigned_url
+                logging.debug(f"Retrieved presigned url for {path}.")
+            elif cached_presigned_url.is_expiring():
+                self.optimize_read()
+                cached_presigned_url = self._url_cache[path]
+        except snowpark_exceptions.SnowparkClientException as e:
+            if self._USE_FALLBACK_FILE_ACCESS == False:  # noqa: E712 # Fallback disabled
+                raise
+            # This may be an intermittent failure, so don't set _USE_FALLBACK_FILE_ACCESS = True
+            logging.warning(f"Pre-signed URL generation failed with {e.message}, trying fallback file access")
+            return self._open_with_snowpark(path)
         url = cached_presigned_url.url
         try:
             return self._fs._open(url, mode=mode, **kwargs)
@@ -387,10 +395,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                     api_calls=[snowpark.DataFrame.collect],
                 ),
             )
-        except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST) or e.message.startswith(
-                fileset_errors.ERRNO_STAGE_NOT_EXIST
-            ):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code in {fileset_errors.ERRNO_DOMAIN_NOT_EXIST, fileset_errors.ERRNO_STAGE_NOT_EXIST}:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_NOT_FOUND,
                     original_exception=fileset_errors.StageNotFoundError(
@@ -406,9 +412,9 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
 
 
 def _match_error_code(ex: snowpark_exceptions.SnowparkSQLException, error_code: int) -> bool:
-    # Snowpark writes error code to message instead of populating e.error_code
+    # Snowpark writes error code to message instead of populating e.sql_error_code
     error_code_str = str(error_code)
-    return ex.error_code == error_code_str or error_code_str in ex.message
+    return ex.sql_error_code == error_code_str or error_code_str in ex.message
 
 
 @snowflake_plan.SnowflakePlan.Decorator.wrap_exception  # type: ignore[misc]

snowflake/ml/lineage/__init__.py

@@ -0,0 +1,3 @@
+from .lineage_node import LineageNode
+
+__all__ = ["LineageNode"]