From c68e154e9fbfbb6b7b2f42b3d859d276d72c60b4 Mon Sep 17 00:00:00 2001
From: Patryk Czajka <patryk.czajka@snowflake.com>
Date: Mon, 8 Apr 2024 15:07:26 +0200
Subject: [PATCH] add quotes for 'use' sql queries

---
 src/snowflake/cli/api/sql_execution.py        |  2 +-
 .../config/connection_configs.toml            |  3 --
 tests_integration/snowflake_connector.py      | 34 +++++++++++++++---
 tests_integration/test_connection.py          | 36 +++++++++++++++++++
 4 files changed, 66 insertions(+), 9 deletions(-)
 create mode 100644 tests_integration/test_connection.py

diff --git a/src/snowflake/cli/api/sql_execution.py b/src/snowflake/cli/api/sql_execution.py
index c9c2b9db34..08e93075e0 100644
--- a/src/snowflake/cli/api/sql_execution.py
+++ b/src/snowflake/cli/api/sql_execution.py
@@ -64,7 +64,7 @@ def _execute_queries(self, queries: str, **kwargs):
         return list(self._execute_string(dedent(queries), **kwargs))
 
     def use(self, object_type: ObjectType, name: str):
-        return self._execute_query(f"use {object_type.value.sf_name} {name}")
+        return self._execute_query(f'use {object_type.value.sf_name} "{name}"')
 
     @contextmanager
     def use_role(self, new_role: str):
diff --git a/tests_integration/config/connection_configs.toml b/tests_integration/config/connection_configs.toml
index 0768b7de5c..2570480b06 100644
--- a/tests_integration/config/connection_configs.toml
+++ b/tests_integration/config/connection_configs.toml
@@ -2,6 +2,3 @@
 [connections.integration]
 schema = "public"
 role = "INTEGRATION_TESTS"
-
-[cli.features]
-enable_snowgit = true
diff --git a/tests_integration/snowflake_connector.py b/tests_integration/snowflake_connector.py
index 8edb61444a..101beda6cf 100644
--- a/tests_integration/snowflake_connector.py
+++ b/tests_integration/snowflake_connector.py
@@ -6,24 +6,48 @@
 import pytest
 from snowflake import connector
 from snowflake.cli.api.exceptions import EnvironmentVariableNotFoundError
+from contextlib import contextmanager
 
 _ENV_PARAMETER_PREFIX = "SNOWFLAKE_CONNECTIONS_INTEGRATION"
 
 
-@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 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
+    os.environ[f"{_ENV_PARAMETER_PREFIX}_DATABASE"] = f"{database_name}"
 
-    yield database_name
+    yield
 
     snowflake_session.execute_string(f"drop database {database_name}")
     del os.environ[f"{_ENV_PARAMETER_PREFIX}_DATABASE"]
 
 
+@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};"
+    )
+    os.environ[f"{_ENV_PARAMETER_PREFIX}_SCHEMA"] = f"{schema_name}"
+    yield
+    snowflake_session.execute_string(f"drop schema {schema_name}")
+    del os.environ[f"{_ENV_PARAMETER_PREFIX}_SCHEMA"]
+
+
+@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")
 def test_role(snowflake_session):
     role_name = f"role_{uuid.uuid4().hex}"
diff --git a/tests_integration/test_connection.py b/tests_integration/test_connection.py
new file mode 100644
index 0000000000..928bb4b804
--- /dev/null
+++ b/tests_integration/test_connection.py
@@ -0,0 +1,36 @@
+import pytest
+from tests_integration.snowflake_connector import (
+    setup_test_database,
+    setup_test_schema,
+    add_uuid_to_name,
+)
+
+
+@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
+@pytest.mark.skip(
+    reason="BUG: connections test command seem to override every setting with schema PUBLIC"
+)
+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