Snow 1313805 fix connection test

RELEASE-NOTES.md

@@ -51,6 +51,8 @@
 * More clear error messages in `snow snowpark build` command
 * Adding support for any source supported by `pip` in `snow snowpark`.
 * Fixed version parsing for packages lookup on Snowflake Anaconda Channel
+* Fix handling database/schema/role identifiers containing dashes
+* Fix schema override bug in `snow connection test`
 
 # v2.1.2
 

src/snowflake/cli/plugins/connection/commands.py

@@ -256,14 +256,19 @@ def test(
     # Test session attributes
     om = ObjectManager()
     try:
+        # "use database" operation changes schema to default "public",
+        # so to test schema set up by user we need to copy it here:
+        schema = conn.schema
+
         if conn.role:
-            om.use(object_type=ObjectType.ROLE, name=conn.role)
+            om.use(object_type=ObjectType.ROLE, name=f'"{conn.role}"')
         if conn.database:
-            om.use(object_type=ObjectType.DATABASE, name=conn.database)
-        if conn.schema:
-            om.use(object_type=ObjectType.SCHEMA, name=conn.schema)
+            om.use(object_type=ObjectType.DATABASE, name=f'"{conn.database}"')
+        if schema:
+            om.use(object_type=ObjectType.SCHEMA, name=f'"{schema}"')
         if conn.warehouse:
-            om.use(object_type=ObjectType.WAREHOUSE, name=conn.warehouse)
+            om.use(object_type=ObjectType.WAREHOUSE, name=f'"{conn.warehouse}"')
+
     except ProgrammingError as err:
         raise ClickException(str(err))
 
@@ -279,9 +284,9 @@ def test(
     }
 
     if enable_diag:
-        result[
-            "Diag Report Location"
-        ] = f"{diag_log_path}/SnowflakeConnectionTestReport.txt"
+        result["Diag Report Location"] = (
+            f"{diag_log_path}/SnowflakeConnectionTestReport.txt"
+        )
 
     return ObjectResult(result)
 

tests/test_connection.py

@@ -323,10 +323,10 @@ def test_connection_test(mock_connect, mock_om, runner):
 
     conn = mock_connect.return_value
     assert mock_om.return_value.use.mock_calls == [
-        mock.call(object_type=ObjectType.ROLE, name=conn.role),
-        mock.call(object_type=ObjectType.DATABASE, name=conn.database),
-        mock.call(object_type=ObjectType.SCHEMA, name=conn.schema),
-        mock.call(object_type=ObjectType.WAREHOUSE, name=conn.warehouse),
+        mock.call(object_type=ObjectType.ROLE, name=f'"{conn.role}"'),
+        mock.call(object_type=ObjectType.DATABASE, name=f'"{conn.database}"'),
+        mock.call(object_type=ObjectType.SCHEMA, name=f'"{conn.schema}"'),
+        mock.call(object_type=ObjectType.WAREHOUSE, name=f'"{conn.warehouse}"'),
     ]
 
 

tests_integration/config/connection_configs.toml

@@ -2,6 +2,3 @@
 [connections.integration]
 schema = "public"
 role = "INTEGRATION_TESTS"
-
-[cli.features]
-enable_snowgit = true

tests_integration/snowflake_connector.py

@@ -6,22 +6,51 @@
 import pytest
 from snowflake import connector
 from snowflake.cli.api.exceptions import EnvironmentVariableNotFoundError
+from contextlib import contextmanager
 
 _ENV_PARAMETER_PREFIX = "SNOWFLAKE_CONNECTIONS_INTEGRATION"
+SCHEMA_ENV_PARAMETER = f"{_ENV_PARAMETER_PREFIX}_SCHEMA"
+DATABASE_ENV_PARAMETER = f"{_ENV_PARAMETER_PREFIX}_DATABASE"
 
 
-@pytest.fixture(scope="function")
-def test_database(snowflake_session):
-    database_name = f"db_{uuid.uuid4().hex}"
+def add_uuid_to_name(name: str) -> str:
+    return f"{name}_{uuid.uuid4().hex}"
+
+
+@contextmanager
+def set_env(env_name: str, value: str):
+    os.environ[env_name] = value
+    yield
+    del os.environ[env_name]
+
+
+@contextmanager
+def setup_test_database(snowflake_session, database_name: str):
+    database_name = f'"{database_name}"'
     snowflake_session.execute_string(
         f"create database {database_name}; use database {database_name}; use schema public;"
     )
-    os.environ[f"{_ENV_PARAMETER_PREFIX}_DATABASE"] = database_name
+    with set_env(DATABASE_ENV_PARAMETER, value=database_name):
+        yield
+    snowflake_session.execute_string(f"drop database {database_name}")
+
+
+@contextmanager
+def setup_test_schema(snowflake_session, schema_name: str):
+    schema_name = f'"{schema_name}"'
+    snowflake_session.execute_string(
+        f"create schema {schema_name}; use schema {schema_name};"
+    )
+    with set_env(SCHEMA_ENV_PARAMETER, value=schema_name):
+        yield
+    snowflake_session.execute_string(f"drop schema {schema_name}")
 
-    yield database_name
 
-    snowflake_session.execute_string(f"drop database {database_name}")
-    del os.environ[f"{_ENV_PARAMETER_PREFIX}_DATABASE"]
+@pytest.fixture(scope="function")
+def test_database(snowflake_session):
+    database_name = add_uuid_to_name("db")
+    with setup_test_database(snowflake_session, database_name):
+        yield database_name
 
 
 @pytest.fixture(scope="function")

tests_integration/test_connection.py

@@ -0,0 +1,46 @@
+import pytest
+from tests_integration.snowflake_connector import (
+    setup_test_database,
+    setup_test_schema,
+    add_uuid_to_name,
+    set_env,
+    SCHEMA_ENV_PARAMETER,
+)
+
+
+@pytest.mark.integration
+def test_connection_test_simple(runner):
+    result = runner.invoke_with_connection_json(["connection", "test"])
+    assert result.exit_code == 0, result.output
+    assert result.json["Status"] == "OK"
+
+
+@pytest.mark.integration
+def test_connection_dashed_database(runner, snowflake_session):
+    database = add_uuid_to_name("dashed-database")
+    with setup_test_database(snowflake_session, database):
+        result = runner.invoke_with_connection_json(["connection", "test"])
+        assert result.exit_code == 0, result.output
+        assert result.json["Database"] == database
+
+
+@pytest.mark.integration
+def test_connection_dashed_schema(
+    runner, test_database, snowflake_session, snowflake_home
+):
+    schema = "dashed-schema-name"
+    with setup_test_schema(snowflake_session, schema):
+        result = runner.invoke_with_connection(["connection", "test", "--debug"])
+        assert result.exit_code == 0, result.output
+        assert f'use schema "{schema}"' in result.output
+
+
+@pytest.mark.integration
+def test_connection_not_existing_schema(
+    runner, test_database, snowflake_session, snowflake_home
+):
+    schema = "schema_which_does_not_exist"
+    with set_env(SCHEMA_ENV_PARAMETER, value=schema):
+        result = runner.invoke_with_connection(["connection", "test"])
+        assert result.exit_code == 1, result.output
+        assert "Object does not exist" in result.output