From f99a296f9b2a5a702c6679e1f215a9ed43cb65fb Mon Sep 17 00:00:00 2001
From: Jan Sikorski <jan.sikorski@snowflake.com>
Date: Tue, 21 Jan 2025 14:26:26 +0100
Subject: [PATCH] Added root object
---
src/snowflake/cli/_plugins/cortex/commands.py | 4 ++--
src/snowflake/cli/api/cli_global_context.py | 7 +++++++
tests_integration/test_connection.py | 3 +++
tests_integration/test_cortex.py | 1 -
4 files changed, 12 insertions(+), 3 deletions(-)
diff --git a/src/snowflake/cli/_plugins/cortex/commands.py b/src/snowflake/cli/_plugins/cortex/commands.py
index 6d64c9ac2e..5cfd0a24d2 100644
--- a/src/snowflake/cli/_plugins/cortex/commands.py
+++ b/src/snowflake/cli/_plugins/cortex/commands.py
@@ -100,11 +100,11 @@ def search(
if not columns:
columns = []
-
conn = get_cli_context().connection
+ root = get_cli_context().root_object
search_service = (
- Root(conn)
+ root
.databases[conn.database]
.schemas[conn.schema]
.cortex_search_services[service]
diff --git a/src/snowflake/cli/api/cli_global_context.py b/src/snowflake/cli/api/cli_global_context.py
index abd4ce09a9..98b16ee655 100644
--- a/src/snowflake/cli/api/cli_global_context.py
+++ b/src/snowflake/cli/api/cli_global_context.py
@@ -21,12 +21,15 @@
from pathlib import Path
from typing import TYPE_CHECKING, Iterator
+from cryptography.utils import cached_property
+
from snowflake.cli.api.connections import ConnectionContext, OpenConnectionCache
from snowflake.cli.api.exceptions import MissingConfiguration
from snowflake.cli.api.metrics import CLIMetrics
from snowflake.cli.api.output.formats import OutputFormat
from snowflake.cli.api.rendering.jinja import CONTEXT_KEY
from snowflake.connector import SnowflakeConnection
+from snowflake.core import Root
if TYPE_CHECKING:
from snowflake.cli.api.project.definition_manager import DefinitionManager
@@ -196,6 +199,10 @@ def _should_force_mute_intermediate_output(self) -> bool:
"""Computes whether cli_console output should be muted."""
return self._manager.output_format == OutputFormat.JSON
+ @cached_property
+ def root_object(self) -> Root:
+ return Root(self.connection)
+
_CLI_CONTEXT_MANAGER: ContextVar[_CliGlobalContextManager | None] = ContextVar(
"cli_context", default=None
diff --git a/tests_integration/test_connection.py b/tests_integration/test_connection.py
index 4128c2e1fc..5d215a32d9 100644
--- a/tests_integration/test_connection.py
+++ b/tests_integration/test_connection.py
@@ -62,3 +62,6 @@ def test_connection_not_existing_schema(
f'Could not use schema "{schema.upper()}". Object does not exist'
in result.output
)
+
+def test_root_object(snowflake_session):
+ root = snowflake_session.root_object
\ No newline at end of file
diff --git a/tests_integration/test_cortex.py b/tests_integration/test_cortex.py
index 749630aff1..5fb2a6eba8 100644
--- a/tests_integration/test_cortex.py
+++ b/tests_integration/test_cortex.py
@@ -24,7 +24,6 @@
@pytest.mark.integration
-@pytest.mark.skip # TODO: when we have Cortex activated on test account, unskip this and setup part in sql script
@pytest.mark.skipif(
sys.version_info >= (3, 12),
reason="Snowflake Python API currently does not support Python 3.12 and greater",