SF Masked Tests (#430)

Co-authored-by: knassre-bodo <[email protected]>

Author: hadia206

.github/workflows/pr_testing.yml

@@ -41,6 +41,11 @@ on:
         type: boolean
         required: false
         default: false        
+      run-sf_masked:
+        description: "Run Snowflake Masked Tests"
+        type: boolean
+        required: false
+        default: false
 
 # Limit CI to cancel previous runs in the same PR
 concurrency:
@@ -143,7 +148,7 @@ jobs:
         run: uv run ruff check .
 
       - name: Run Tests
-        run: uv run pytest tests/ -m "not (snowflake or mysql or postgres)" -rs
+        run: uv run pytest tests/ -m "not (snowflake or mysql or postgres or sf_masked)" -rs
 
   run-sf-tests:
     name: Snowflake Tests
@@ -193,3 +198,24 @@ jobs:
       python-versions: ${{ github.event_name == 'workflow_dispatch'
                       && needs.get-py-ver-matrix.outputs.matrix
                       || '["3.10", "3.11", "3.12"]' }}
+
+  run-sf-masked-tests:
+    name: Masked Snowflake Tests
+    needs: [get-msg, get-py-ver-matrix]
+    if: |
+      (github.event_name == 'pull_request' && contains(needs.get-msg.outputs.commitMsg, '[run all]')) ||
+      (github.event_name == 'pull_request' && contains(needs.get-msg.outputs.commitMsg, '[run sf_masked]')) ||
+      (github.event_name == 'workflow_dispatch' && (inputs.run-all || inputs.run-sf_masked))
+    uses: ./.github/workflows/sf_masked_testing.yml
+    secrets:
+      SF_FULL_USERNAME: ${{ secrets.SF_FULL_USERNAME }}
+      SF_FULL_PASSWORD: ${{ secrets.SF_FULL_PASSWORD }}
+      SF_PARTIAL_USERNAME: ${{ secrets.SF_PARTIAL_USERNAME }}
+      SF_PARTIAL_PASSWORD: ${{ secrets.SF_PARTIAL_PASSWORD }}
+      SF_NONE_USERNAME: ${{ secrets.SF_NONE_USERNAME }}
+      SF_NONE_PASSWORD: ${{ secrets.SF_NONE_PASSWORD }}
+      SF_MASKED_ACCOUNT: ${{ secrets.SF_MASKED_ACCOUNT }}
+    with:
+      python-versions: ${{ github.event_name == 'workflow_dispatch'
+                      && needs.get-py-ver-matrix.outputs.matrix
+                      || '["3.10", "3.11", "3.12"]' }}

.github/workflows/sf_masked_testing.yml

@@ -0,0 +1,68 @@
+name: Run Snowflake Tests
+
+on:
+  workflow_call:
+    inputs:
+      python-versions:
+        description: "JSON string of Python versions"
+        type: string
+        required: true
+    secrets:
+      SF_FULL_USERNAME:
+        required: true
+      SF_FULL_PASSWORD:
+        required: true
+      SF_PARTIAL_USERNAME:
+        required: true
+      SF_PARTIAL_PASSWORD:
+        required: true
+      SF_NONE_USERNAME:
+        required: true
+      SF_NONE_PASSWORD:
+        required: true
+      SF_MASKED_ACCOUNT:
+        required: true
+
+jobs:
+  sf-tests:
+    name: Snowflake Tests (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ${{ fromJSON(inputs.python-versions) }}
+
+    env:
+      SF_FULL_USERNAME: ${{ secrets.SF_FULL_USERNAME }}
+      SF_FULL_PASSWORD: ${{ secrets.SF_FULL_PASSWORD }}
+      SF_PARTIAL_USERNAME: ${{ secrets.SF_PARTIAL_USERNAME }}
+      SF_PARTIAL_PASSWORD: ${{ secrets.SF_PARTIAL_PASSWORD }}
+      SF_NONE_USERNAME: ${{ secrets.SF_NONE_USERNAME }}
+      SF_NONE_PASSWORD: ${{ secrets.SF_NONE_PASSWORD }}
+      SF_MASKED_ACCOUNT: ${{ secrets.SF_MASKED_ACCOUNT }}
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup Python ${{ matrix.python-version }}
+        id: setup-python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "0.4.23"
+
+      - name: Create virtual environment
+        # See note in sf_testing.yml about why we use `uv venv` here.
+        run: uv venv
+
+      - name: Install dependencies
+        run: uv pip install -e ".[snowflake]"
+
+      - name: Confirm Snowflake connector is installed
+        run: uv run python -c "import snowflake.connector; print(snowflake.connector.__version__)"
+
+      - name: Run Snowflake Tests
+        run: uv run pytest -m sf_masked tests/ -rs

pytest.ini

@@ -4,3 +4,4 @@ markers =
     snowflake: marks tests that require Snowflake credentials
     mysql: marks tests that require MySQL credentials
     postgres: marks tests that require PostgresSQL credentials
+    sf_masked: marks tests that require Snowflake Masked credentials

tests/conftest.py

@@ -575,6 +575,24 @@ def sqlite_technograph_connection() -> DatabaseContext:
     return DatabaseContext(DatabaseConnection(connection), DatabaseDialect.SQLITE)
 
 
+@pytest.fixture(
+    params=[
+        pytest.param(("0", "raw"), id="raw"),
+        pytest.param(("1", "rewrite"), id="rewrite"),
+    ]
+)
+def enable_mask_rewrites(request):
+    """
+    Temporarily enable the mask rewrites by setting the environment variable to
+    the specified value, returning the name that should be used to identify
+    the situation ("raw" for disabled, "rewrite" for enabled).
+    """
+    old_value: str = os.environ.get("PYDOUGH_ENABLE_MASK_REWRITES", "0")
+    os.environ["PYDOUGH_ENABLE_MASK_REWRITES"] = request.param[0]
+    yield request.param[1]
+    os.environ["PYDOUGH_ENABLE_MASK_REWRITES"] = old_value
+
+
 @pytest.fixture(scope="session")
 def sqlite_cryptbank_connection() -> DatabaseContext:
     """