SNOW-891470 adding implementation for default_connection_name (#1721)

Author: sfc-gh-mkeller

src/snowflake/connector/__init__.py

@@ -7,6 +7,8 @@
 #
 from __future__ import annotations
 
+from functools import wraps
+
 apilevel = "2.0"
 threadsafety = 2
 paramstyle = "pyformat"
@@ -47,6 +49,7 @@
 logging.getLogger(__name__).addHandler(NullHandler())
 
 
+@wraps(SnowflakeConnection.__init__)
 def Connect(**kwargs) -> SnowflakeConnection:
     return SnowflakeConnection(**kwargs)
 

src/snowflake/connector/config_manager.py

@@ -126,7 +126,7 @@ def value(self) -> Any:
                 value = self._get_config()
                 source = "configuration file"
             except MissingConfigOptionError:
-                if self.default:
+                if self.default is not None:
                     source = "default_value"
                     value = self.default
                 else:
@@ -166,8 +166,8 @@ def _get_env(self) -> tuple[bool, str | _T | None]:
         env_var = os.environ.get(env_name)
         if env_var is None:
             return False, None
-        loaded_var: str | _T | None = env_var
-        if env_var and self.parse_str is not None:
+        loaded_var: str | _T = env_var
+        if self.parse_str is not None:
             loaded_var = self.parse_str(env_var)
         if isinstance(loaded_var, (Table, tomlkit.TOMLDocument)):
             # If we got a TOML table we probably want it in dictionary form
@@ -449,6 +449,7 @@ def __getitem__(self, name: str) -> ConfigOption | ConfigManager:
 CONFIG_MANAGER.add_option(
     name="connections",
     parse_str=tomlkit.parse,
+    default=dict(),
 )
 CONFIG_MANAGER.add_option(
     name="default_connection_name",

src/snowflake/connector/connection.py

@@ -302,6 +302,21 @@ def __init__(
         connections_file_path: pathlib.Path | None = None,
         **kwargs,
     ) -> None:
+        """Create a new SnowflakeConnection.
+
+        Connections can be loaded from the TOML file located at
+        snowflake.connector.constants.CONNECTIONS_FILE.
+
+        When connection_name is supplied we will first load that connection
+        and then override any other values supplied.
+
+        When no arguments are given (other than connection_file_path) the
+        default connection will be loaded first. Note that no overwriting is
+        supported in this case.
+
+        If overwriting values from the default connection is desirable, supply
+        the name explicitly.
+        """
         self._lock_sequence_counter = Lock()
         self.sequence_counter = 0
         self._errorhandler = Error.default_errorhandler
@@ -324,6 +339,7 @@ def __init__(
             setattr(self, f"_{name}", value)
 
         self.heartbeat_thread = None
+        is_kwargs_empty = not kwargs
 
         if "application" not in kwargs:
             if ENV_VAR_PARTNER in os.environ.keys():
@@ -349,6 +365,17 @@ def __init__(
                     f" known ones are {list(connections.keys())}"
                 )
             kwargs = {**connections[connection_name], **kwargs}
+        elif is_kwargs_empty:
+            # connection_name is None and kwargs was empty when called
+            def_connection_name = CONFIG_MANAGER["default_connection_name"]
+            connections = CONFIG_MANAGER["connections"]
+            if def_connection_name not in connections:
+                raise Error(
+                    f"Default connection with name '{def_connection_name}' "
+                    "cannot be found, known ones are "
+                    f"{list(connections.keys())}"
+                )
+            kwargs = {**connections[def_connection_name]}
         self.__set_error_attributes()
         self.connect(**kwargs)
         self._telemetry = TelemetryClient(self._rest)

test/integ/test_connection.py

@@ -1226,7 +1226,7 @@ def test_connection_name_loading(monkeypatch, db_parameters, tmp_path, mode):
 
     doc = tomlkit.document()
     default_con = tomlkit.table()
-    tmp_config_file: None | pathlib.Path = None
+    tmp_connections_file: None | pathlib.Path = None
     try:
         # If anything unexpected fails here, don't want to expose password
         for k, v in db_parameters.items():
@@ -1236,11 +1236,11 @@ def test_connection_name_loading(monkeypatch, db_parameters, tmp_path, mode):
             if mode == "env":
                 m.setenv("SF_CONNECTIONS", tomlkit.dumps(doc))
             else:
-                tmp_config_file = tmp_path / "connections.toml"
-                tmp_config_file.write_text(tomlkit.dumps(doc))
+                tmp_connections_file = tmp_path / "connections.toml"
+                tmp_connections_file.write_text(tomlkit.dumps(doc))
             with snowflake.connector.connect(
                 connection_name="default",
-                connections_file_path=tmp_config_file,
+                connections_file_path=tmp_connections_file,
             ) as conn:
                 with conn.cursor() as cur:
                     assert cur.execute("select 1;").fetchall() == [
@@ -1253,6 +1253,32 @@ def test_connection_name_loading(monkeypatch, db_parameters, tmp_path, mode):
         pytest.fail("something failed", pytrace=False)
 
 
+@pytest.mark.skipolddriver
+def test_default_connection_name_loading(monkeypatch, db_parameters):
+    import tomlkit
+
+    doc = tomlkit.document()
+    default_con = tomlkit.table()
+    try:
+        # If anything unexpected fails here, don't want to expose password
+        for k, v in db_parameters.items():
+            default_con[k] = v
+        doc["default"] = default_con
+        with monkeypatch.context() as m:
+            m.setenv("SNOWFLAKE_CONNECTIONS", tomlkit.dumps(doc))
+            m.setenv("SNOWFLAKE_DEFAULT_CONNECTION_NAME", "default")
+            with snowflake.connector.connect() as conn:
+                with conn.cursor() as cur:
+                    assert cur.execute("select 1;").fetchall() == [
+                        (1,),
+                    ]
+    except Exception:
+        # This is my way of guaranteeing that we'll not expose the
+        # sensitive information that this test needs to handle.
+        # db_parameter contains passwords.
+        pytest.fail("something failed", pytrace=False)
+
+
 @pytest.mark.skipolddriver
 def test_not_found_connection_name():
     connection_name = random_string(5)

test/unit/test_connection.py

@@ -8,11 +8,15 @@
 import json
 import os
 import sys
+from textwrap import dedent
 from unittest.mock import patch
 
 import pytest
 
 import snowflake.connector
+from snowflake.connector.errors import Error
+
+from ..randomize import random_string
 
 try:
     from snowflake.connector.auth import (
@@ -28,10 +32,11 @@
 
 try:  # pragma: no cover
     from snowflake.connector.auth import AuthByUsrPwdMfa
+    from snowflake.connector.config_manager import CONFIG_MANAGER
     from snowflake.connector.constants import ENV_VAR_PARTNER, QueryStatus
 except ImportError:
     ENV_VAR_PARTNER = "SF_PARTNER"
-    QueryStatus = None
+    QueryStatus = CONFIG_MANAGER = None
 
     class AuthByUsrPwdMfa(AuthByDefault):
         def __init__(self, password: str, mfa_token: str) -> None:
@@ -211,3 +216,103 @@ def test_negative_custom_auth(auth_class):
             user="user",
             auth_class=auth_class,
         )
+
+
+def test_missing_default_connection(monkeypatch, tmp_path):
+    connections_file = tmp_path / "connections.toml"
+    config_file = tmp_path / "config.toml"
+    with monkeypatch.context() as m:
+        m.delenv("SNOWFLAKE_DEFAULT_CONNECTION_NAME", raising=False)
+        m.delenv("SNOWFLAKE_CONNECTIONS", raising=False)
+        m.setattr(CONFIG_MANAGER, "conf_file_cache", None)
+        m.setattr(CONFIG_MANAGER, "file_path", config_file)
+
+        with pytest.raises(
+            Error,
+            match="Default connection with name 'default' cannot be found, known ones are \\[\\]",
+        ):
+            snowflake.connector.connect(connections_file_path=connections_file)
+
+
+def test_missing_default_connection_conf_file(monkeypatch, tmp_path):
+    connection_name = random_string(5)
+    connections_file = tmp_path / "connections.toml"
+    config_file = tmp_path / "config.toml"
+    config_file.write_text(
+        dedent(
+            f"""\
+            default_connection_name = "{connection_name}"
+            """
+        )
+    )
+    with monkeypatch.context() as m:
+        m.delenv("SNOWFLAKE_DEFAULT_CONNECTION_NAME", raising=False)
+        m.delenv("SNOWFLAKE_CONNECTIONS", raising=False)
+        m.setattr(CONFIG_MANAGER, "conf_file_cache", None)
+        m.setattr(CONFIG_MANAGER, "file_path", config_file)
+
+        with pytest.raises(
+            Error,
+            match=f"Default connection with name '{connection_name}' cannot be found, known ones are \\[\\]",
+        ):
+            snowflake.connector.connect(connections_file_path=connections_file)
+
+
+def test_missing_default_connection_conn_file(monkeypatch, tmp_path):
+    connections_file = tmp_path / "connections.toml"
+    config_file = tmp_path / "config.toml"
+    connections_file.write_text(
+        dedent(
+            """\
+            [con_a]
+            user = "test user"
+            account = "test account"
+            password = "test password"
+            """
+        )
+    )
+    with monkeypatch.context() as m:
+        m.delenv("SNOWFLAKE_DEFAULT_CONNECTION_NAME", raising=False)
+        m.delenv("SNOWFLAKE_CONNECTIONS", raising=False)
+        m.setattr(CONFIG_MANAGER, "conf_file_cache", None)
+        m.setattr(CONFIG_MANAGER, "file_path", config_file)
+
+        with pytest.raises(
+            Error,
+            match="Default connection with name 'default' cannot be found, known ones are \\['con_a'\\]",
+        ):
+            snowflake.connector.connect(connections_file_path=connections_file)
+
+
+def test_missing_default_connection_conf_conn_file(monkeypatch, tmp_path):
+    connection_name = random_string(5)
+    connections_file = tmp_path / "connections.toml"
+    config_file = tmp_path / "config.toml"
+    config_file.write_text(
+        dedent(
+            f"""\
+            default_connection_name = "{connection_name}"
+            """
+        )
+    )
+    connections_file.write_text(
+        dedent(
+            """\
+            [con_a]
+            user = "test user"
+            account = "test account"
+            password = "test password"
+            """
+        )
+    )
+    with monkeypatch.context() as m:
+        m.delenv("SNOWFLAKE_DEFAULT_CONNECTION_NAME", raising=False)
+        m.delenv("SNOWFLAKE_CONNECTIONS", raising=False)
+        m.setattr(CONFIG_MANAGER, "conf_file_cache", None)
+        m.setattr(CONFIG_MANAGER, "file_path", config_file)
+
+        with pytest.raises(
+            Error,
+            match=f"Default connection with name '{connection_name}' cannot be found, known ones are \\['con_a'\\]",
+        ):
+            snowflake.connector.connect(connections_file_path=connections_file)