Merge branch 'main' into external-access-db-cleanup

.github/workflows/test_cli_action.yaml

@@ -18,23 +18,11 @@ jobs:
         with:
           cli-version: "latest"
           default-config-file-path: "tests_integration/config/connection_configs.toml"
-      - name: Set up key (Ubuntu / MacOS)
-        env:
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY: ${{ secrets.SNOWFLAKE_PRIVATE_KEY }}
-        run: |
-          PARENT_DIR=$(dirname "${{ github.workspace }}")
-          PRIVATE_KEY_FILE=$PARENT_DIR/.ssh/key.p8
-          echo "PARENT_DIR=$PARENT_DIR" >> $GITHUB_ENV
-          echo "PRIVATE_KEY_FILE=$PRIVATE_KEY_FILE" >> $GITHUB_ENV
-
-          mkdir $PARENT_DIR/.ssh
-          echo "${SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY}" > $PRIVATE_KEY_FILE
-          sudo chmod 600 $PRIVATE_KEY_FILE
       - name: Test connection
         env:
           TERM: unknown
           SNOWFLAKE_CONNECTIONS_INTEGRATION_AUTHENTICATOR: SNOWFLAKE_JWT
           SNOWFLAKE_CONNECTIONS_INTEGRATION_USER: ${{ secrets.SNOWFLAKE_USER }}
           SNOWFLAKE_CONNECTIONS_INTEGRATION_ACCOUNT: ${{ secrets.SNOWFLAKE_ACCOUNT }}
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_FILE: ${{ env.PRIVATE_KEY_FILE }}
+          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_RAW: ${{ secrets.SNOWFLAKE_PRIVATE_KEY_RAW }}
         run: snow connection test -c integration | grep Status

.github/workflows/test_fork.yaml

@@ -31,43 +31,14 @@ jobs:
         run: |
           python -m pip install --upgrade pip hatch
           python -m hatch env create ${{ inputs.python-env }}
-
-      - name: Set up key (Ubuntu / MacOS)
-        if: inputs.runs-on == 'ubuntu-latest' || inputs.runs-on == 'macos-latest'
-        env:
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY: ${{ secrets.SNOWFLAKE_PRIVATE_KEY }}
-        run: |
-          PARENT_DIR=$(dirname "${{ github.workspace }}")
-          PRIVATE_KEY_FILE=$PARENT_DIR/.ssh/key.p8
-          echo "PARENT_DIR=$PARENT_DIR" >> $GITHUB_ENV
-          echo "PRIVATE_KEY_FILE=$PRIVATE_KEY_FILE" >> $GITHUB_ENV
-
-          mkdir $PARENT_DIR/.ssh
-          echo "${SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY}" > $PRIVATE_KEY_FILE
-          sudo chmod 600 $PRIVATE_KEY_FILE
-
-      - name: Set up key (Windows)
-        if: inputs.runs-on == 'windows-latest'
-        env:
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY: ${{ secrets.SNOWFLAKE_PRIVATE_KEY }}
-        run: |
-          $parentDir = Split-Path -Parent "${{ github.workspace }}"
-          $privateKeyPath = $parentDir + "\\ssh\\key.p8"
-          echo "PARENT_DIR=$parentDir" >> $env:GITHUB_ENV
-          echo "PRIVATE_KEY_FILE=$privateKeyPath" >> $env:GITHUB_ENV
-
-          mkdir $parentDir\\ssh
-          Set-Content -Path $privateKeyPath -Value $env:SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY
-        shell: pwsh
-
       - name: Run integration tests
         env:
           TERM: unknown
           SNOWFLAKE_CONNECTIONS_INTEGRATION_AUTHENTICATOR: SNOWFLAKE_JWT
           SNOWFLAKE_CONNECTIONS_INTEGRATION_USER: ${{ secrets.SNOWFLAKE_USER }}
           SNOWFLAKE_CONNECTIONS_INTEGRATION_ACCOUNT: ${{ secrets.SNOWFLAKE_ACCOUNT }}
           SNOWFLAKE_CONNECTIONS_INTEGRATION_DATABASE: ${{ secrets.SNOWFLAKE_DATABASE }}
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_FILE: ${{ env.PRIVATE_KEY_FILE }}
+          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_RAW: ${{ secrets.SNOWFLAKE_PRIVATE_KEY_RAW }}
         run: python -m hatch run ${{ inputs.hatch-run }}
 
       # Update check run called "integration-fork"

.github/workflows/test_trusted.yaml

@@ -31,35 +31,6 @@ jobs:
         run: |
           python -m pip install --upgrade pip hatch
           python -m hatch env create ${{ inputs.python-env }}
-
-      - name: Set up key (Ubuntu / MacOS)
-        if: inputs.runs-on == 'ubuntu-latest' || inputs.runs-on == 'macos-latest'
-        env:
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY: ${{ secrets.SNOWFLAKE_PRIVATE_KEY }}
-        run: |
-          PARENT_DIR=$(dirname "${{ github.workspace }}")
-          PRIVATE_KEY_FILE=$PARENT_DIR/.ssh/key.p8
-          echo "PARENT_DIR=$PARENT_DIR" >> $GITHUB_ENV
-          echo "PRIVATE_KEY_FILE=$PRIVATE_KEY_FILE" >> $GITHUB_ENV
-
-          mkdir $PARENT_DIR/.ssh
-          echo "${SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY}" > $PRIVATE_KEY_FILE
-          sudo chmod 600 $PRIVATE_KEY_FILE
-
-      - name: Set up key (Windows)
-        if: inputs.runs-on == 'windows-latest'
-        env:
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY: ${{ secrets.SNOWFLAKE_PRIVATE_KEY }}
-        run: |
-          $parentDir = Split-Path -Parent "${{ github.workspace }}"
-          $privateKeyPath = $parentDir + "\\ssh\\key.p8"
-          echo "PARENT_DIR=$parentDir" >> $env:GITHUB_ENV
-          echo "PRIVATE_KEY_FILE=$privateKeyPath" >> $env:GITHUB_ENV
-
-          mkdir $parentDir\\ssh
-          Set-Content -Path $privateKeyPath -Value $env:SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY
-        shell: pwsh
-
       - name: Run integration tests
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -68,5 +39,5 @@ jobs:
           SNOWFLAKE_CONNECTIONS_INTEGRATION_USER: ${{ secrets.SNOWFLAKE_USER }}
           SNOWFLAKE_CONNECTIONS_INTEGRATION_ACCOUNT: ${{ secrets.SNOWFLAKE_ACCOUNT }}
           SNOWFLAKE_CONNECTIONS_INTEGRATION_DATABASE: ${{ secrets.SNOWFLAKE_DATABASE }}
-          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_FILE: ${{ env.PRIVATE_KEY_FILE }}
+          SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_RAW: ${{ secrets.SNOWFLAKE_PRIVATE_KEY_RAW }}
         run: python -m hatch run ${{ inputs.hatch-run }}

CONTRIBUTING.md

@@ -82,7 +82,7 @@ where ``<key>`` is the name of the key. The following environment variables are
 - `SNOWFLAKE_CONNECTIONS_INTEGRATION_HOST`
 - `SNOWFLAKE_CONNECTIONS_INTEGRATION_ACCOUNT`
 - `SNOWFLAKE_CONNECTIONS_INTEGRATION_USER`
-- `SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_FILE` (`SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_PATH` will work too, but preferable is with `_FILE` ending)
+- `SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_RAW` (Loads the private key directly from the environment variable)
 - `SNOWFLAKE_CONNECTIONS_INTEGRATION_ROLE`
 - `SNOWFLAKE_CONNECTIONS_INTEGRATION_DATABASE`
 - `SNOWFLAKE_CONNECTIONS_INTEGRATION_WAREHOUSE`

RELEASE-NOTES.md

@@ -42,6 +42,7 @@
 * Added `snow ws migrate` command to migrate `snowflake.yml` file from V1 to V2.
 * Added `--package-entity-id` and `--app-entity-id` options to `snow app` commands to allow targeting specific entities when the `definition_version` in `snowflake.yml` is `2` or higher and it contains multiple `application package` or `application` entities.
 * Added templates expansion of arbitrary files for Native Apps through `templates` processor.
+* Added `SNOWFLAKE_..._PRIVATE_KEY_RAW` environment variable to pass private key as a raw string.
 
 ## Fixes and improvements
 * Fixed problem with whitespaces in `snow connection add` command.

src/snowflake/cli/_app/snow_connector.py

@@ -125,10 +125,7 @@ def connect_to_snowflake(
         k: v for k, v in connection_parameters.items() if v is not None
     }
 
-    if "private_key_file" in connection_parameters:
-        update_connection_details_with_private_key(
-            connection_parameters, "private_key_file"
-        )
+    update_connection_details_with_private_key(connection_parameters)
 
     if mfa_passcode:
         connection_parameters["passcode"] = mfa_passcode
@@ -192,11 +189,37 @@ def _raise_errors_related_to_session_token(
         )
 
 
-def update_connection_details_with_private_key(
-    connection_parameters: Dict, private_key_var_name: str = "private_key_file"
+def update_connection_details_with_private_key(connection_parameters: Dict):
+    if "private_key_file" in connection_parameters:
+        _load_private_key(connection_parameters, "private_key_file")
+    elif "private_key_path" in connection_parameters:
+        _load_private_key(connection_parameters, "private_key_path")
+    elif "private_key_raw" in connection_parameters:
+        _load_private_key_from_parameters(connection_parameters, "private_key_raw")
+    return connection_parameters
+
+
+def _load_private_key(connection_parameters: Dict, private_key_var_name: str) -> None:
+    if connection_parameters.get("authenticator") == "SNOWFLAKE_JWT":
+        private_key_pem = _load_pem_from_file(
+            connection_parameters[private_key_var_name]
+        )
+        private_key = _load_pem_to_der(private_key_pem)
+        connection_parameters["private_key"] = private_key
+        del connection_parameters[private_key_var_name]
+    else:
+        raise ClickException(
+            "Private Key authentication requires authenticator set to SNOWFLAKE_JWT"
+        )
+
+
+def _load_private_key_from_parameters(
+    connection_parameters: Dict, private_key_var_name: str
 ) -> None:
     if connection_parameters.get("authenticator") == "SNOWFLAKE_JWT":
-        private_key = _load_pem_to_der(connection_parameters[private_key_var_name])
+        private_key_pem = connection_parameters[private_key_var_name]
+        private_key_pem = private_key_pem.encode("utf-8")
+        private_key = _load_pem_to_der(private_key_pem)
         connection_parameters["private_key"] = private_key
         del connection_parameters[private_key_var_name]
     else:
@@ -213,17 +236,19 @@ def _update_connection_application_name(connection_parameters: Dict):
     connection_parameters.update(connection_application_params)
 
 
-def _load_pem_to_der(private_key_file: str) -> bytes:
-    """
-    Given a private key file path (in PEM format), decode key data into DER
-    format
-    """
-
+def _load_pem_from_file(private_key_file: str) -> bytes:
     with SecurePath(private_key_file).open(
         "rb", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB
     ) as f:
         private_key_pem = f.read()
+    return private_key_pem
+
 
+def _load_pem_to_der(private_key_pem: bytes) -> bytes:
+    """
+    Given a private key file path (in PEM format), decode key data into DER
+    format
+    """
     private_key_passphrase = os.getenv("PRIVATE_KEY_PASSPHRASE", None)
     if (
         private_key_pem.startswith(ENCRYPTED_PKCS8_PK_HEADER)

tests/test_connection.py

@@ -698,13 +698,14 @@ def test_token_file_path_tokens(mock_connector, mock_ctx, runner, temp_dir):
     clear=True,
 )
 @mock.patch("snowflake.connector.connect")
+@mock.patch("snowflake.cli._app.snow_connector._load_pem_from_file")
 @mock.patch("snowflake.cli._app.snow_connector._load_pem_to_der")
 def test_key_pair_authentication_from_config(
-    mock_load, mock_connector, mock_ctx, temp_dir, runner
+    mock_convert, mock_load_file, mock_connector, mock_ctx, temp_dir, runner
 ):
     ctx = mock_ctx()
     mock_connector.return_value = ctx
-    mock_load.return_value = "secret value"
+    mock_convert.return_value = "secret value"
 
     with NamedTemporaryFile("w+", suffix="toml") as tmp_file:
         tmp_file.write(
@@ -726,7 +727,7 @@ def test_key_pair_authentication_from_config(
         )
 
     assert result.exit_code == 0, result.output
-    mock_load.assert_called_once_with("~/sf_private_key.p8")
+    mock_load_file.assert_called_once_with("~/sf_private_key.p8")
     mock_connector.assert_called_once_with(
         application="SNOWCLI.OBJECT.LIST",
         account="my_account",

tests/test_snow_connector.py

@@ -89,7 +89,9 @@ def test_command_context_is_passed_to_snowflake_connection(
 @mock.patch("snowflake.connector.connect")
 @mock.patch("snowflake.cli._app.snow_connector.command_info")
 @mock.patch("snowflake.cli._app.snow_connector._load_pem_to_der")
+@mock.patch("snowflake.cli._app.snow_connector._load_pem_from_file")
 def test_private_key_loading_and_aliases(
+    mock_load_pem_from_file,
     mock_load_pem_to_der,
     mock_command_info,
     mock_connect,
@@ -117,6 +119,7 @@ def test_private_key_loading_and_aliases(
             overrides[user_input] = override_value
 
     mock_command_info.return_value = "SNOWCLI.SQL"
+    mock_load_pem_from_file.return_value = b"bytes"
     mock_load_pem_to_der.return_value = b"bytes"
 
     conn_dict = get_connection_dict(connection_name)
@@ -141,7 +144,8 @@ def test_private_key_loading_and_aliases(
             **expected_private_key_args,
         )
         if expected_private_key_file_value is not None:
-            mock_load_pem_to_der.assert_called_with(expected_private_key_file_value)
+            mock_load_pem_from_file.assert_called_with(expected_private_key_file_value)
+            mock_load_pem_to_der.assert_called_with(b"bytes")
 
 
 @mock.patch.dict(os.environ, {}, clear=True)

tests_integration/snowflake_connector.py

@@ -95,6 +95,7 @@ def snowflake_session():
         "account": _get_from_env("ACCOUNT"),
         "user": _get_from_env("USER"),
         "private_key_file": _get_private_key_file(),
+        "private_key_raw": _get_from_env("PRIVATE_KEY_RAW", allow_none=True),
         "host": _get_from_env("HOST", allow_none=True),
         "warehouse": _get_from_env("WAREHOUSE", allow_none=True),
         "role": _get_from_env("ROLE", allow_none=True),
@@ -121,4 +122,4 @@ def _get_private_key_file() -> Optional[str]:
     private_key_file = _get_from_env("PRIVATE_KEY_PATH", allow_none=True)
     if private_key_file is not None:
         return private_key_file
-    return _get_from_env("PRIVATE_KEY_FILE")
+    return _get_from_env("PRIVATE_KEY_FILE", allow_none=True)

tests_integration/test_temporary_connection.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import tempfile
 
 import pytest
 import os
@@ -28,30 +29,34 @@
         "SNOWFLAKE_CONNECTIONS_INTEGRATION_USER": os.environ.get(
             "SNOWFLAKE_CONNECTIONS_INTEGRATION_USER", None
         ),
-        "SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_FILE": os.environ.get(
-            "SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_PATH",
-            os.environ.get("SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_FILE"),
+        "SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_RAW": os.environ.get(
+            "SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_RAW",
         ),
     },
     clear=True,
 )
 def test_temporary_connection(runner, snapshot):
 
-    result = runner.invoke(
-        [
-            "sql",
-            "-q",
-            "select 1",
-            "--temporary-connection",
-            "--authenticator",
-            "SNOWFLAKE_JWT",
-            "--account",
-            os.environ["SNOWFLAKE_CONNECTIONS_INTEGRATION_ACCOUNT"],
-            "--user",
-            os.environ["SNOWFLAKE_CONNECTIONS_INTEGRATION_USER"],
-            "--private-key-file",
-            os.environ["SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_FILE"],
-        ]
-    )
-    assert result.exit_code == 0
-    assert result.output == snapshot
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        private_key_path = os.path.join(tmp_dir, "private_key.p8")
+        with open(private_key_path, "w") as f:
+            f.write(os.environ["SNOWFLAKE_CONNECTIONS_INTEGRATION_PRIVATE_KEY_RAW"])
+
+        result = runner.invoke(
+            [
+                "sql",
+                "-q",
+                "select 1",
+                "--temporary-connection",
+                "--authenticator",
+                "SNOWFLAKE_JWT",
+                "--account",
+                os.environ["SNOWFLAKE_CONNECTIONS_INTEGRATION_ACCOUNT"],
+                "--user",
+                os.environ["SNOWFLAKE_CONNECTIONS_INTEGRATION_USER"],
+                "--private-key-file",
+                str(private_key_path),
+            ]
+        )
+        assert result.exit_code == 0
+        assert result.output == snapshot