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

snowflake/cortex/_complete.py

@@ -1,7 +1,12 @@
-from typing import Optional, Union
+from typing import Iterator, Optional, Union
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    call_rest_function,
+    call_sql_function,
+    process_rest_response,
+)
 from snowflake.ml._internal import telemetry
 
 
@@ -10,19 +15,35 @@ from snowflake.ml._internal import telemetry
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
 def Complete(
-    model: Union[str, snowpark.Column], prompt: Union[str, snowpark.Column], session: Optional[snowpark.Session] = None
-) -> Union[str, snowpark.Column]:
+    model: Union[str, snowpark.Column],
+    prompt: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+    use_rest_api_experimental: bool = False,
+    stream: bool = False,
+) -> Union[str, Iterator[str], snowpark.Column]:
     """Complete calls into the LLM inference service to perform completion.
 
     Args:
         model: A Column of strings representing model types.
         prompt: A Column of prompts to send to the LLM.
         session: The snowpark session to use. Will be inferred by context if not specified.
+        use_rest_api_experimental (bool): Toggles between the use of SQL and REST implementation. This feature is
+            experimental and can be removed at any time.
+        stream (bool): Enables streaming. When enabled, a generator function is returned that provides the streaming
+            output as it is received. Each update is a string containing the new text content since the previous update.
+            The use of streaming requires the experimental use_rest_api_experimental flag to be enabled.
+
+    Raises:
+        ValueError: If `stream` is set to True and `use_rest_api_experimental` is set to False.
 
     Returns:
         A column of string responses.
     """
-
+    if stream is True and use_rest_api_experimental is False:
+        raise ValueError("If stream is set to True use_rest_api_experimental must also be set to True")
+    if use_rest_api_experimental:
+        response = call_rest_function("complete", model, prompt, session=session, stream=stream)
+        return process_rest_response(response)
     return _complete_impl("snowflake.cortex.complete", model, prompt, session=session)
 
 

snowflake/cortex/_sentiment.py

@@ -11,7 +11,7 @@ from snowflake.ml._internal import telemetry
 )
 def Sentiment(
     text: Union[str, snowpark.Column], session: Optional[snowpark.Session] = None
-) -> Union[str, snowpark.Column]:
+) -> Union[float, snowpark.Column]:
     """Sentiment calls into the LLM inference service to perform sentiment analysis on the input text.
 
     Args:
@@ -21,11 +21,14 @@ def Sentiment(
     Returns:
         A column of floats. 1 represents positive sentiment, -1 represents negative sentiment.
     """
-
     return _sentiment_impl("snowflake.cortex.sentiment", text, session=session)
 
 
 def _sentiment_impl(
     function: str, text: Union[str, snowpark.Column], session: Optional[snowpark.Session] = None
-) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text)
+) -> Union[float, snowpark.Column]:
+
+    output = call_sql_function(function, session, text)
+    if isinstance(output, snowpark.Column):
+        return output
+    return float(output)

snowflake/cortex/_sse_client.py

@@ -0,0 +1,81 @@
+from typing import Iterator, cast
+
+import requests
+
+
+class Event:
+    def __init__(self, event: str = "message", data: str = "") -> None:
+        self.event = event
+        self.data = data
+
+    def __str__(self) -> str:
+        s = f"{self.event} event"
+        if self.data:
+            s += f", {len(self.data)} bytes"
+        else:
+            s += ", no data"
+        return s
+
+
+class SSEClient:
+    def __init__(self, response: requests.Response) -> None:
+
+        self.response = response
+
+    def _read(self) -> Iterator[str]:
+
+        lines = b""
+        for chunk in self.response:
+            for line in chunk.splitlines(True):
+                lines += line
+                if lines.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")):
+                    yield cast(str, lines)
+                    lines = b""
+        if lines:
+            yield cast(str, lines)
+
+    def events(self) -> Iterator[Event]:
+        for raw_event in self._read():
+            event = Event()
+            # splitlines() only uses \r and \n
+            for line in raw_event.splitlines():
+
+                line = cast(bytes, line).decode("utf-8")
+
+                data = line.split(":", 1)
+                field = data[0]
+
+                if len(data) > 1:
+                    # "If value starts with a single U+0020 SPACE character,
+                    # remove it from value. .strip() would remove all white spaces"
+                    if data[1].startswith(" "):
+                        value = data[1][1:]
+                    else:
+                        value = data[1]
+                else:
+                    value = ""
+
+                # The data field may come over multiple lines and their values
+                # are concatenated with each other.
+                if field == "data":
+                    event.data += value + "\n"
+                elif field == "event":
+                    event.event = value
+
+            if not event.data:
+                continue
+
+            # If the data field ends with a newline, remove it.
+            if event.data.endswith("\n"):
+                event.data = event.data[0:-1]  # Replace trailing newline - rstrip would remove multiple.
+
+            # Empty event names default to 'message'
+            event.event = event.event or "message"
+
+            if event.event != "message":  # ignore anything but “message” or default event
+                continue
+
+            yield event
+
+    def close(self) -> None:
+        self.response.close()

snowflake/cortex/_util.py

@@ -1,15 +1,34 @@
-from typing import Optional, Union, cast
+import json
+from typing import Iterator, Optional, Union, cast
+from urllib.parse import urljoin, urlparse
+
+import requests
 
 from snowflake import snowpark
+from snowflake.cortex._sse_client import SSEClient
 from snowflake.snowpark import context, functions
 
 CORTEX_FUNCTIONS_TELEMETRY_PROJECT = "CortexFunctions"
 
 
+class SSEParseException(Exception):
+    """This exception is raised when an invalid server sent event is received from the server."""
+
+    pass
+
+
+class SnowflakeAuthenticationException(Exception):
+    """This exception is raised when the session object does not have session.connection.rest.token attribute."""
+
+    pass
+
+
 # Calls a sql function, handling both immediate (e.g. python types) and batch
 # (e.g. snowpark column and literal type modes).
 def call_sql_function(
-    function: str, session: Optional[snowpark.Session], *args: Union[str, snowpark.Column]
+    function: str,
+    session: Optional[snowpark.Session],
+    *args: Union[str, snowpark.Column],
 ) -> Union[str, snowpark.Column]:
     handle_as_column = False
     for arg in args:
@@ -17,21 +36,29 @@ def call_sql_function(
             handle_as_column = True
 
     if handle_as_column:
-        return cast(Union[str, snowpark.Column], call_sql_function_column(function, *args))
-    return cast(Union[str, snowpark.Column], call_sql_function_immediate(function, session, *args))
+        return cast(Union[str, snowpark.Column], _call_sql_function_column(function, *args))
+    return cast(
+        Union[str, snowpark.Column],
+        _call_sql_function_immediate(function, session, *args),
+    )
 
 
-def call_sql_function_column(function: str, *args: Union[str, snowpark.Column]) -> snowpark.Column:
+def _call_sql_function_column(function: str, *args: Union[str, snowpark.Column]) -> snowpark.Column:
     return cast(snowpark.Column, functions.builtin(function)(*args))
 
 
-def call_sql_function_immediate(
-    function: str, session: Optional[snowpark.Session], *args: Union[str, snowpark.Column]
+def _call_sql_function_immediate(
+    function: str,
+    session: Optional[snowpark.Session],
+    *args: Union[str, snowpark.Column],
 ) -> str:
     if session is None:
         session = context.get_active_session()
     if session is None:
-        raise Exception("No session available in the current context nor specified as an argument.")
+        raise SnowflakeAuthenticationException(
+            """Session required. Provide the session through a session=... argument or ensure an active session is
+            available in your environment."""
+        )
 
     lit_args = []
     for arg in args:
@@ -40,3 +67,73 @@ def call_sql_function_immediate(
     empty_df = session.create_dataframe([snowpark.Row()])
     df = empty_df.select(functions.builtin(function)(*lit_args))
     return cast(str, df.collect()[0][0])
+
+
+def call_rest_function(
+    function: str,
+    model: Union[str, snowpark.Column],
+    prompt: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+    stream: bool = False,
+) -> requests.Response:
+    if session is None:
+        session = context.get_active_session()
+    if session is None:
+        raise SnowflakeAuthenticationException(
+            """Session required. Provide the session through a session=... argument or ensure an active session is
+            available in your environment."""
+        )
+
+    if not hasattr(session.connection.rest, "token"):
+        raise SnowflakeAuthenticationException("Snowflake session error: REST token missing.")
+
+    if session.connection.rest.token is None or session.connection.rest.token == "":  # type: ignore[union-attr]
+        raise SnowflakeAuthenticationException("Snowflake session error: REST token is empty.")
+
+    url = urljoin(session.connection.host, f"api/v2/cortex/inference/{function}")
+    if urlparse(url).scheme == "":
+        url = "https://" + url
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f'Snowflake Token="{session.connection.rest.token}"',  # type: ignore[union-attr]
+        "Accept": "application/json, text/event-stream",
+    }
+
+    data = {
+        "model": model,
+        "messages": [{"content": prompt}],
+        "stream": stream,
+    }
+
+    response = requests.post(
+        url,
+        json=data,
+        headers=headers,
+        stream=stream,
+    )
+    response.raise_for_status()
+    return response
+
+
+def process_rest_response(response: requests.Response, stream: bool = False) -> Union[str, Iterator[str]]:
+    if not stream:
+        try:
+            message = response.json()["choices"][0]["message"]
+            output = str(message.get("content", ""))
+            return output
+        except (KeyError, IndexError) as e:
+            raise SSEParseException("Failed to parse streamed response.") from e
+    else:
+        return _return_gen(response)
+
+
+def _return_gen(response: requests.Response) -> Iterator[str]:
+    client = SSEClient(response)
+    for event in client.events():
+        response_loaded = json.loads(event.data)
+        try:
+            delta = response_loaded["choices"][0]["delta"]
+            output = str(delta.get("content", ""))
+            yield output
+        except (KeyError, IndexError) as e:
+            raise SSEParseException("Failed to parse streamed response.") from e

snowflake/ml/_internal/lineage/lineage_utils.py

@@ -1,21 +1,11 @@
 import copy
 import functools
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal.lineage import data_source
 
-DATA_SOURCES_ATTR = "_data_sources"
-
-
-def _get_datasources(*args: Any) -> List[data_source.DataSource]:
-    """Helper method for extracting data sources attribute from DataFrames in an argument list"""
-    result = []
-    for arg in args:
-        srcs = getattr(arg, DATA_SOURCES_ATTR, None)
-        if isinstance(srcs, list) and all(isinstance(s, data_source.DataSource) for s in srcs):
-            result += srcs
-    return result
+_DATA_SOURCES_ATTR = "_data_sources"
 
 
 def _wrap_func(
@@ -32,6 +22,37 @@ def _wrap_func(
     return wrapped
 
 
+def _wrap_class_func(fn: Callable[..., snowpark.DataFrame]) -> Callable[..., snowpark.DataFrame]:
+    @functools.wraps(fn)
+    def wrapped(*args: Any, **kwargs: Any) -> snowpark.DataFrame:
+        df = fn(*args, **kwargs)
+        data_sources = get_data_sources(*args, *kwargs.values())
+        if data_sources:
+            patch_dataframe(df, data_sources, inplace=True)
+        return df
+
+    return wrapped
+
+
+def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
+    """Helper method for extracting data sources attribute from DataFrames in an argument list"""
+    result: Optional[List[data_source.DataSource]] = None
+    for arg in args:
+        srcs = getattr(arg, _DATA_SOURCES_ATTR, None)
+        if isinstance(srcs, list) and all(isinstance(s, data_source.DataSource) for s in srcs):
+            if result is None:
+                result = []
+            result += srcs
+    return result
+
+
+def set_data_sources(obj: Any, data_sources: Optional[List[data_source.DataSource]]) -> None:
+    """Helper method for attaching data sources to an object"""
+    if data_sources:
+        assert all(isinstance(ds, data_source.DataSource) for ds in data_sources)
+    setattr(obj, _DATA_SOURCES_ATTR, data_sources)
+
+
 def patch_dataframe(
     df: snowpark.DataFrame, data_sources: List[data_source.DataSource], inplace: bool = False
 ) -> snowpark.DataFrame:
@@ -62,7 +83,7 @@ def patch_dataframe(
     ]
     if not inplace:
         df = copy.copy(df)
-    setattr(df, DATA_SOURCES_ATTR, data_sources)
+    set_data_sources(df, data_sources)
     for func in funcs:
         fn = getattr(df, func, None)
         if fn is not None:
@@ -70,18 +91,6 @@ def patch_dataframe(
     return df
 
 
-def _wrap_class_func(fn: Callable[..., snowpark.DataFrame]) -> Callable[..., snowpark.DataFrame]:
-    @functools.wraps(fn)
-    def wrapped(*args: Any, **kwargs: Any) -> snowpark.DataFrame:
-        df = fn(*args, **kwargs)
-        data_sources = _get_datasources(*args) + _get_datasources(*kwargs.values())
-        if data_sources:
-            patch_dataframe(df, data_sources, inplace=True)
-        return df
-
-    return wrapped
-
-
 # Class-level monkey-patches
 for klass, func_list in {
     snowpark.DataFrame: [

snowflake/ml/_internal/utils/temp_file_utils.py

@@ -8,14 +8,17 @@ from absl.logging import logging
 logger = logging.getLogger(__name__)
 
 
-def get_temp_file_path() -> str:
+def get_temp_file_path(prefix: str = "") -> str:
     """Returns a new random temp file path.
 
+    Args:
+        prefix: A prefix to the temp file path, this can help add stored file information. Defaults to None.
+
     Returns:
         A new temp file path.
     """
     # TODO(snandamuri): Use in-memory filesystem for temp files.
-    local_file = tempfile.NamedTemporaryFile(delete=True)
+    local_file = tempfile.NamedTemporaryFile(prefix=prefix, delete=True)
     local_file_name = local_file.name
     local_file.close()
     return local_file_name

snowflake/ml/dataset/dataset.py

@@ -65,6 +65,20 @@ class DatasetVersion:
         comment: Optional[str] = self._get_property("comment")
         return comment
 
+    @property
+    def label_cols(self) -> List[str]:
+        metadata = self._get_metadata()
+        if metadata is None or metadata.label_cols is None:
+            return []
+        return metadata.label_cols
+
+    @property
+    def exclude_cols(self) -> List[str]:
+        metadata = self._get_metadata()
+        if metadata is None or metadata.exclude_cols is None:
+            return []
+        return metadata.exclude_cols
+
     def _get_property(self, property_name: str, default: Any = None) -> Any:
         if self._properties is None:
             sql_result = (
@@ -91,17 +105,6 @@ class DatasetVersion:
                 warnings.warn(f"Metadata parsing failed with error: {e}", UserWarning, stacklevel=2)
         return self._metadata
 
-    def _get_exclude_cols(self) -> List[str]:
-        metadata = self._get_metadata()
-        if metadata is None:
-            return []
-        cols = []
-        if metadata.exclude_cols:
-            cols.extend(metadata.exclude_cols)
-        if metadata.label_cols:
-            cols.extend(metadata.label_cols)
-        return cols
-
     def url(self) -> str:
         """Returns the URL of the DatasetVersion contents in Snowflake.
 
@@ -168,7 +171,7 @@ class Dataset:
                         fully_qualified_name=self._fully_qualified_name,
                         version=v.name,
                         url=v.url(),
-                        exclude_cols=v._get_exclude_cols(),
+                        exclude_cols=(v.label_cols + v.exclude_cols),
                     )
                 ],
             )

snowflake/ml/dataset/dataset_factory.py

@@ -16,8 +16,7 @@ def create_from_dataframe(
     **version_kwargs: Any,
 ) -> dataset.Dataset:
     """
-    Create a new versioned Dataset from a DataFrame and returns
-    a DatasetReader for the newly created Dataset version.
+    Create a new versioned Dataset from a DataFrame.
 
     Args:
         session: The Snowpark Session instance to use.
@@ -39,7 +38,7 @@ def create_from_dataframe(
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def load_dataset(session: snowpark.Session, name: str, version: str) -> dataset.Dataset:
     """
-    Load a versioned Dataset into a DatasetReader.
+    Load a versioned Dataset.
 
     Args:
         session: The Snowpark Session instance to use.
@@ -47,7 +46,7 @@ def load_dataset(session: snowpark.Session, name: str, version: str) -> dataset.
         version: The dataset version name.
 
     Returns:
-        A DatasetReader object.
+        A Dataset object.
     """
     ds: dataset.Dataset = dataset.Dataset.load(session, name).select_version(version)
     return ds

snowflake/ml/feature_store/access_manager.py

@@ -42,6 +42,8 @@ class _SessionInfo:
 # Lists of permissions as tuples of (OBJECT_TYPE, [PRIVILEGES, ...])
 _PRE_INIT_PRIVILEGES: Dict[_FeatureStoreRole, List[_Privilege]] = {
     _FeatureStoreRole.PRODUCER: [
+        _Privilege("DATABASE", "{database}", ["USAGE"]),
+        _Privilege("SCHEMA", "{database}.{schema}", ["USAGE"]),
         _Privilege(
             "SCHEMA",
             "{database}.{schema}",
@@ -69,8 +71,7 @@ _PRE_INIT_PRIVILEGES: Dict[_FeatureStoreRole, List[_Privilege]] = {
         _Privilege("DYNAMIC TABLE", _ALL_OBJECTS, ["SELECT", "MONITOR"], "SCHEMA {database}.{schema}"),
         _Privilege("VIEW", _ALL_OBJECTS, ["SELECT", "REFERENCES"], "SCHEMA {database}.{schema}"),
         _Privilege("TABLE", _ALL_OBJECTS, ["SELECT", "REFERENCES"], "SCHEMA {database}.{schema}"),
-        # FIXME(dhung): FUTURE DATASETS not supported until 8.19
-        # _Privilege("DATASET", _ALL_OBJECTS, ["USAGE"], "SCHEMA {database}.{schema}"),
+        _Privilege("DATASET", _ALL_OBJECTS, ["USAGE"], "SCHEMA {database}.{schema}"),
         # User should decide whether they want to grant warehouse usage to CONSUMER
         # _Privilege("WAREHOUSE", "{warehouse}", ["USAGE"]),
     ],
@@ -128,8 +129,7 @@ def _grant_privileges(
 def _configure_pre_init_privileges(
     session: Session,
     session_info: _SessionInfo,
-    producer_role: str = "SNOWML_FEATURE_STORE_PRODUCER_RL",
-    consumer_role: str = "SNOWML_FEATURE_STORE_CONSUMER_RL",
+    roles_to_create: Dict[_FeatureStoreRole, str],
 ) -> None:
     """
     Configure Feature Store role privileges. Must be run with ACCOUNTADMIN
@@ -141,8 +141,7 @@ def _configure_pre_init_privileges(
     Args:
         session: Snowpark Session to interact with Snowflake backend.
         session_info: Session info like database and schema for the FeatureStore instance.
-        producer_role: Name of producer role to be configured.
-        consumer_role: Name of consumer role to be configured.
+        roles_to_create: Producer and optional consumer roles to create.
     """
 
     # Create schema if not already exists
@@ -159,29 +158,30 @@ def _configure_pre_init_privileges(
 
     # Pass schema ownership from admin to PRODUCER
     if schema_created:
+        # TODO: we are missing a test case for this code path
         session.sql(
-            f"GRANT OWNERSHIP ON SCHEMA {session_info.database}.{session_info.schema} TO ROLE {producer_role}"
+            f"GRANT OWNERSHIP ON SCHEMA {session_info.database}.{session_info.schema} "
+            f"TO ROLE {roles_to_create[_FeatureStoreRole.PRODUCER]}"
         ).collect()
 
     # Grant privileges to roles
-    _grant_privileges(session, producer_role, _PRE_INIT_PRIVILEGES[_FeatureStoreRole.PRODUCER], session_info)
-    _grant_privileges(session, consumer_role, _PRE_INIT_PRIVILEGES[_FeatureStoreRole.CONSUMER], session_info)
+    for role_type, role in roles_to_create.items():
+        _grant_privileges(session, role, _PRE_INIT_PRIVILEGES[role_type], session_info)
 
 
 def _configure_post_init_privileges(
     session: Session,
     session_info: _SessionInfo,
-    producer_role: str = "FS_PRODUCER",
-    consumer_role: str = "FS_CONSUMER",
+    roles_to_create: Dict[_FeatureStoreRole, str],
 ) -> None:
-    _grant_privileges(session, producer_role, _POST_INIT_PRIVILEGES[_FeatureStoreRole.PRODUCER], session_info)
-    _grant_privileges(session, consumer_role, _POST_INIT_PRIVILEGES[_FeatureStoreRole.CONSUMER], session_info)
+    for role_type, role in roles_to_create.items():
+        _grant_privileges(session, role, _POST_INIT_PRIVILEGES[role_type], session_info)
 
 
 def _configure_role_hierarchy(
     session: Session,
     producer_role: str,
-    consumer_role: str,
+    consumer_role: Optional[str],
 ) -> None:
     """
     Create Feature Store roles and configure role hierarchy hierarchy. Must be run with
@@ -195,18 +195,17 @@ def _configure_role_hierarchy(
         producer_role: Name of producer role to be configured.
         consumer_role: Name of consumer role to be configured.
     """
+    # Create the necessary roles and build role hierarchy
     producer_role = SqlIdentifier(producer_role)
-    consumer_role = SqlIdentifier(consumer_role)
-
-    # Create the necessary roles
     session.sql(f"CREATE ROLE IF NOT EXISTS {producer_role}").collect()
-    session.sql(f"CREATE ROLE IF NOT EXISTS {consumer_role}").collect()
-
-    # Build role hierarchy
-    session.sql(f"GRANT ROLE {consumer_role} TO ROLE {producer_role}").collect()
     session.sql(f"GRANT ROLE {producer_role} TO ROLE SYSADMIN").collect()
     session.sql(f"GRANT ROLE {producer_role} TO ROLE {session.get_current_role()}").collect()
 
+    if consumer_role is not None:
+        consumer_role = SqlIdentifier(consumer_role)
+        session.sql(f"CREATE ROLE IF NOT EXISTS {consumer_role}").collect()
+        session.sql(f"GRANT ROLE {consumer_role} TO ROLE {producer_role}").collect()
+
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def setup_feature_store(
@@ -215,7 +214,7 @@ def setup_feature_store(
     schema: str,
     warehouse: str,
     producer_role: str = "FS_PRODUCER",
-    consumer_role: str = "FS_CONSUMER",
+    consumer_role: Optional[str] = None,
 ) -> FeatureStore:
     """
     Sets up a new Feature Store including role/privilege setup. Must be run with ACCOUNTADMIN
@@ -230,7 +229,7 @@ def setup_feature_store(
         schema: Schema to create the FeatureStore instance.
         warehouse: Default warehouse for Feature Store compute.
         producer_role: Name of producer role to be configured.
-        consumer_role: Name of consumer role to be configured.
+        consumer_role: Name of consumer role to be configured. If not specified, consumer role won't be created.
 
     Returns:
         Feature Store instance.
@@ -249,20 +248,25 @@ def setup_feature_store(
     )
 
     try:
+        roles_to_create = {_FeatureStoreRole.PRODUCER: producer_role}
+        if consumer_role is not None:
+            roles_to_create.update({_FeatureStoreRole.CONSUMER: consumer_role})
         _configure_role_hierarchy(session, producer_role=producer_role, consumer_role=consumer_role)
     except exceptions.SnowparkSQLException:
         # Error can be safely ignored if roles already exist and hierarchy is already built
-        for role in (producer_role, consumer_role):
+        for _, role in roles_to_create.items():
             # Ensure roles already exist
             if session.sql(f"SHOW ROLES LIKE '{role}' STARTS WITH '{role}'").count() == 0:
                 raise
-        # Ensure hierarchy already configured
-        consumer_grants = session.sql(f"SHOW GRANTS ON ROLE {consumer_role}").collect()
-        if not any(r["granted_to"] == "ROLE" and r["grantee_name"] == producer_role for r in consumer_grants):
-            raise
+
+        if consumer_role is not None:
+            # Ensure hierarchy already configured
+            consumer_grants = session.sql(f"SHOW GRANTS ON ROLE {consumer_role}").collect()
+            if not any(r["granted_to"] == "ROLE" and r["grantee_name"] == producer_role for r in consumer_grants):
+                raise
 
     # Do any pre-FeatureStore.__init__() privilege setup
-    _configure_pre_init_privileges(session, session_info, producer_role, consumer_role)
+    _configure_pre_init_privileges(session, session_info, roles_to_create)
 
     # Use PRODUCER role to create and operate new Feature Store
     current_role = session.get_current_role()
@@ -274,6 +278,6 @@ def setup_feature_store(
         session.use_role(current_role)
 
     # Do any post-FeatureStore.__init__() privilege setup
-    _configure_post_init_privileges(session, session_info, producer_role, consumer_role)
+    _configure_post_init_privileges(session, session_info, roles_to_create)
 
     return fs

snowflake/ml/feature_store/feature_store.py

@@ -920,7 +920,7 @@ class FeatureStore:
         try:
             if output_type == "table":
                 table_name = f"{name}_{version}"
-                result_df.write.mode("errorifexists").save_as_table(table_name)  # type: ignore[call-overload]
+                result_df.write.mode("errorifexists").save_as_table(table_name)
                 ds_df = self._session.table(table_name)
                 return ds_df
             else:
@@ -1761,8 +1761,8 @@ class FeatureStore:
             self._session.sql(
                 f"""
                     SELECT * FROM TABLE(
-                        INFORMATION_SCHEMA.TAG_REFERENCES_INTERNAL(
-                            TAG_NAME => '{_FEATURE_STORE_OBJECT_TAG}'
+                        {self._config.database}.INFORMATION_SCHEMA.TAG_REFERENCES_INTERNAL(
+                            TAG_NAME => '{self._get_fully_qualified_name(_FEATURE_STORE_OBJECT_TAG)}'
                         )
                     ) LIMIT 1;
                 """