Repository: Snowflake-Labs/semantic-model-generator

Issue:
App needs to support MFA which required for all demo accounts and potential default for customer accounts
the app has to pass the Authenticator parameter of username_password_mfa along with a passcode.  Check out this doc:  https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#label-snowflake-connector-methods
(https://snowflake.slack.com/archives/D07DFRA3PPZ/p1721840365218259)

I did set the Authenticator and Passcode env variables, but they aren't called in the app.

So the error I get is this:  snowflake.connector.errors.OperationalError: 250001: Could not connect to Snowflake backend after 2 attempt(s).Aborting

and if I do that enough times it will lock me out

Discussion:

Solution patch:
diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml
index 217e23c..ddcb915 100644
--- a/.github/workflows/lint.yaml
+++ b/.github/workflows/lint.yaml
@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9"]
+        python-version: [ "3.10" ]
 
     steps:
       - name: Check out the code
@@ -50,16 +50,16 @@ jobs:
 
       - name: Run mypy
         run: |
-            make run_mypy
+          make run_mypy
 
       - name: Check with black
         run: |
-            make check_black
+          make check_black
 
       - name: Check with isort
         run: |
-            make check_isort
+          make check_isort
 
       - name: Run flake8
         run: |
-            make run_flake8
+          make run_flake8
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index e421ca5..9f1fc89 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9"]
+        python-version: [ "3.10" ]
     steps:
       - name: Check out the code
         uses: actions/checkout@v4
diff --git a/README.md b/README.md
index b04c683..71ab83c 100644
--- a/README.md
+++ b/README.md
@@ -38,7 +38,6 @@ account, [follow these instructions](https://docs.snowflake.com/en/user-guide/or
 export SNOWFLAKE_ROLE="<your-snowflake-role>"
 export SNOWFLAKE_WAREHOUSE="<your-snowflake-warehouse>"
 export SNOWFLAKE_USER="<your-snowflake-user>"
-export SNOWFLAKE_PASSWORD="<your-snowflake-password>"
 export SNOWFLAKE_ACCOUNT_LOCATOR="<your-snowflake-account-locator>"
 export SNOWFLAKE_HOST="<your-snowflake-host>"
 ```
@@ -49,7 +48,6 @@ export SNOWFLAKE_HOST="<your-snowflake-host>"
 set SNOWFLAKE_ROLE=<your-snowflake-role>
 set SNOWFLAKE_WAREHOUSE=<your-snowflake-warehouse>
 set SNOWFLAKE_USER=<your-snowflake-user>
-set SNOWFLAKE_PASSWORD=<your-snowflake-password>
 set SNOWFLAKE_ACCOUNT_LOCATOR=<your-snowflake-account-locator>
 set SNOWFLAKE_HOST=<your-snowflake-host>
 ```
@@ -63,15 +61,48 @@ import os
 os.environ['SNOWFLAKE_ROLE'] = '<your-snowflake-role>'
 os.environ['SNOWFLAKE_WAREHOUSE'] = '<your-snowflake-warehouse>'
 os.environ['SNOWFLAKE_USER'] = '<your-snowflake-user>'
-os.environ['SNOWFLAKE_PASSWORD'] = '<your-snowflake-password>'
 os.environ['SNOWFLAKE_ACCOUNT_LOCATOR'] = '<your-snowflake-account-locator>'
 os.environ['SNOWFLAKE_HOST'] = '<your-snowflake-host>'
 ```
 
-4. If you are using single sign-on authetication, you can leave password empty and set additional environ below:
+Our semantic model generators currently support three types of authentication. If no `SNOWFLAKE_AUTHENTICATOR`
+environment variable
+is set, the default is `snowflake`, which uses standard username/password support.
+
+1. Username and Password
+
+**Note**: If you have MFA enabled, using this default authenticator should send a push notification to your device.
 
- ```python
-os.environ['SNOWFLAKE_AUTHENTICATOR'] = 'EXTERNALBROWSER'
+```bash
+# no SNOWFLAKE_AUTHENTICATOR needed
+export SNOWFLAKE_PASSWORD="<your-snowflake-password>"
+```
+
+2. Username/Password with MFA passcode
+
+Using a passcode from your authenticator app:
+
+```bash
+export SNOWFLAKE_AUTHENTICATOR="username_password_mfa"
+export SNOWFLAKE_PASSWORD="<your-snowflake-password>"
+
+# if your authenticator app says "123 456", enter "123456" (no spaces)
+export SNOWFLAKE_MFA_PASSCODE="<your-snowflake-mfa-passcode>" 
+```
+
+Using a passcode embedded in the password:
+
+```bash
+export SNOWFLAKE_AUTHENTICATOR="username_password_mfa"
+export SNOWFLAKE_PASSWORD="<your-snowflake-password>"
+export SNOWFLAKE_MFA_PASSCODE_IN_PASSWORD="true"
+```
+
+3. Single Sign-On (SSO) with Okta
+
+```bash
+# no SNOWFLAKE_PASSWORD needed
+export SNOWFLAKE_AUTHENTICATOR="externalbrowser"
 ```
 
 ## Usage
diff --git a/admin_apps/app.py b/admin_apps/app.py
index aed5775..bcdce8e 100644
--- a/admin_apps/app.py
+++ b/admin_apps/app.py
@@ -49,7 +49,10 @@ def verify_environment_setup() -> None:
 
     # Instantiate the Snowflake connection that gets reused throughout the app.
     try:
-        get_snowflake_connection()
+        with st.spinner(
+            "Validating your connection to Snowflake. If you are using MFA, please check your authenticator app for a push notification."
+        ):
+            get_snowflake_connection()
     except DatabaseError:
         failed_connection_popup()
 
diff --git a/admin_apps/requirements.txt b/admin_apps/requirements.txt
index 1b2d7ac..e99bb6e 100644
--- a/admin_apps/requirements.txt
+++ b/admin_apps/requirements.txt
@@ -5,4 +5,4 @@ sqlglot==23.17.0
 streamlit==1.36.0
 streamlit-extras==0.4.3
 streamlit-monaco==0.1.3
-snowflake-connector-python==3.11.0
\ No newline at end of file
+snowflake-connector-python[secure-local-storage]==3.11.0
\ No newline at end of file
diff --git a/admin_apps/shared_utils.py b/admin_apps/shared_utils.py
index 8805464..1a37c25 100644
--- a/admin_apps/shared_utils.py
+++ b/admin_apps/shared_utils.py
@@ -48,7 +48,7 @@ def get_connector() -> SnowflakeConnector:
     )
 
 
-@st.cache_resource
+@st.cache_resource(show_spinner=False)
 def get_snowflake_connection() -> SnowflakeConnection:
     """
     Opens a general connection to Snowflake using the provided SnowflakeConnector
diff --git a/pyproject.toml b/pyproject.toml
index 8b3468e..44207b9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,14 +10,14 @@ readme = "README.md"
 python = ">=3.8.1,<3.9.7 || >3.9.7,<3.12"
 pandas = "^2.0.1"
 loguru = "^0.7.2"
-snowflake-connector-python = {extras = ["pandas"], version = "^3.8.1"}
+snowflake-connector-python = { extras = ["secure-local-storage", "pandas"], version = "^3.8.1" }
 protobuf = "^4.25.3"
 pydantic = "^2.6.4"
 PyYAML = "^6.0.1"
 "ruamel.yaml" = "^0.18.6"
 jsonargparse = "^4.27.7"
 tqdm = "^4.66.2"
-pyarrow = ">=10.0.1,<10.1.0" 
+pyarrow = ">=10.0.1,<10.1.0"
 sqlglot = "^23.10.0"
 strictyaml = "^1.7.3"
 streamlit-nightly = "^1.35.1.dev20240609"
diff --git a/semantic_model_generator/snowflake_utils/env_vars.py b/semantic_model_generator/snowflake_utils/env_vars.py
index cf1ed3d..d7985d2 100644
--- a/semantic_model_generator/snowflake_utils/env_vars.py
+++ b/semantic_model_generator/snowflake_utils/env_vars.py
@@ -9,6 +9,10 @@ SNOWFLAKE_HOST = os.getenv("SNOWFLAKE_HOST")
 SNOWFLAKE_AUTHENTICATOR = os.getenv("SNOWFLAKE_AUTHENTICATOR")
 SNOWFLAKE_ACCOUNT_LOCATOR = os.getenv("SNOWFLAKE_ACCOUNT_LOCATOR")
 
+# Optional MFA environment variables
+SNOWFLAKE_MFA_PASSCODE = os.getenv("SNOWFLAKE_MFA_PASSCODE")
+SNOWFLAKE_MFA_PASSCODE_IN_PASSWORD = os.getenv("SNOWFLAKE_MFA_PASSCODE_IN_PASSWORD")
+
 
 def assert_required_env_vars() -> list[str]:
     """
diff --git a/semantic_model_generator/snowflake_utils/snowflake_connector.py b/semantic_model_generator/snowflake_utils/snowflake_connector.py
index 6362f9d..749f529 100644
--- a/semantic_model_generator/snowflake_utils/snowflake_connector.py
+++ b/semantic_model_generator/snowflake_utils/snowflake_connector.py
@@ -347,6 +347,16 @@ class SnowflakeConnector:
         auth = env_vars.SNOWFLAKE_AUTHENTICATOR
         return auth
 
+    def _get_mfa_passcode(self) -> Optional[str]:
+        passcode = env_vars.SNOWFLAKE_MFA_PASSCODE
+        return passcode
+
+    def _is_mfa_passcode_in_password(self) -> bool:
+        passcode_in_password = env_vars.SNOWFLAKE_MFA_PASSCODE_IN_PASSWORD
+        if not passcode_in_password:
+            return False
+        return passcode_in_password.lower() == "true"
+
     @contextmanager
     def connect(
         self, db_name: str, schema_name: Optional[str] = None
@@ -382,6 +392,8 @@ class SnowflakeConnector:
             warehouse=self._get_warehouse(),
             host=self._get_host(),
             authenticator=self._get_authenticator(),
+            passcode=self._get_mfa_passcode(),
+            passcode_in_password=self._is_mfa_passcode_in_password(),
         )
         if db_name:
             set_database(connection, db_name=db_name)
diff --git a/semantic_model_generator/snowflake_utils/utils.py b/semantic_model_generator/snowflake_utils/utils.py
index 7e43316..6de58c6 100644
--- a/semantic_model_generator/snowflake_utils/utils.py
+++ b/semantic_model_generator/snowflake_utils/utils.py
@@ -28,8 +28,10 @@ def create_connection_parameters(
     database: Optional[str] = None,
     schema: Optional[str] = None,
     authenticator: Optional[str] = None,
-) -> Dict[str, str]:
-    connection_parameters: Dict[str, str] = dict(user=user, account=account)
+    passcode: Optional[str] = None,
+    passcode_in_password: Optional[bool] = None,
+) -> Dict[str, str | bool]:
+    connection_parameters: Dict[str, str | bool] = dict(user=user, account=account)
     if password:
         connection_parameters["password"] = password
     if role:
@@ -44,10 +46,14 @@ def create_connection_parameters(
         connection_parameters["authenticator"] = authenticator
     if host:
         connection_parameters["host"] = host
+    if passcode:
+        connection_parameters["passcode"] = passcode
+    if passcode_in_password:
+        connection_parameters["passcode_in_password"] = passcode_in_password
     return connection_parameters
 
 
-def _connection(connection_parameters: Dict[str, str]) -> SnowflakeConnection:
+def _connection(connection_parameters: Dict[str, str | bool]) -> SnowflakeConnection:
     # https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect
     return connect(**connection_parameters)
 
@@ -60,6 +66,8 @@ def snowflake_connection(
     password: Optional[str] = None,
     host: Optional[str] = None,
     authenticator: Optional[str] = None,
+    passcode: Optional[str] = None,
+    passcode_in_password: Optional[bool] = None,
 ) -> SnowflakeConnection:
     """
     Returns a Snowflake Connection to the specified account.
@@ -73,5 +81,7 @@ def snowflake_connection(
             role=role,
             warehouse=warehouse,
             authenticator=authenticator,
+            passcode=passcode,
+            passcode_in_password=passcode_in_password,
         )
     )


Test patch:
diff --git a/semantic_model_generator/tests/generate_model_test.py b/semantic_model_generator/tests/generate_model_test.py
index df974ee..5d0e54b 100644
--- a/semantic_model_generator/tests/generate_model_test.py
+++ b/semantic_model_generator/tests/generate_model_test.py
@@ -189,6 +189,10 @@ def mock_snowflake_connection_env(monkeypatch):
         SnowflakeConnector, "_get_host", return_value="test_host"
     ), patch.object(
         SnowflakeConnector, "_get_authenticator", return_value="test_authenticator"
+    ), patch.object(
+        SnowflakeConnector, "_get_mfa_passcode", return_value="123456"
+    ), patch.object(
+        SnowflakeConnector, "_is_mfa_passcode_in_password", return_value=False
     ):
         yield
 
@@ -360,6 +364,8 @@ def test_raw_schema_to_semantic_context(
         warehouse="test_warehouse",
         host="test_host",
         authenticator="test_authenticator",
+        passcode="123456",
+        passcode_in_password=False,
     )