[
{
"path": ".cursor/rules/overall-guidelines.mdc",
"content": "---\nalwaysApply: true\n---\n\n# Cursor Rules for Go Snowflake Driver\n\n## General Development Standards\n\n### Code Quality\n- Follow Go formatting standards (use `gofmt`)\n- Use meaningful variable and function names\n- Include error handling for all operations that can fail\n- Write comprehensive documentation for public APIs\n\n### Project Structure\n- Place test files in the same package as the code being tested\n- Use `test_data/` directory for test fixtures and sample data\n- Group related functionality in logical packages\n\n### Testing\n- Test files should be named `*_test.go`\n- **For test-specific rules, see `testing.mdc`**\n- Write both positive and negative test cases\n- Use table-driven tests for testing multiple scenarios\n\n### Code Review Guidelines\n- Ensure code follows Go best practices\n- Verify comprehensive test coverage\n- Check that error messages are descriptive and helpful for debugging\n- Validate that public APIs are properly documented\n"
},
{
"path": ".cursor/rules/testing.mdc",
"content": "---\nalwaysApply: true\n---\n\n# Cursor Rules for Go Test Files\n\nThis file automatically applies when working on `*_test.go` files.\n\n## Testing Standards\n\n### Assertion Helper Usage\n- **ALWAYS** Attempt to use assertion helpers from `assert_test.go` instead of direct `t.Fatal`, `t.Fatalf`, `t.Error`, or `t.Errorf` calls. Where it makes sense, add new assertion helpers.\n- **NEVER** write manual if-then-fatal patterns in test functions when a suitable assertion helper exists.\n\n#### Common Assertion Patterns:\n\n**Error Checking:**\n```go\n// ❌ WRONG\nif err != nil {\n t.Fatalf(\"Unexpected error: %v\", err)\n}\n\n// ✅ CORRECT \nassertNilF(t, err, \"Unexpected error\")\n```\n\n**Nil Checking:**\n```go\n// ❌ WRONG\nif obj == nil {\n t.Fatal(\"Expected non-nil object\")\n}\n\n// ✅ CORRECT\nassertNotNilF(t, obj, \"Expected non-nil object\")\n```\n\n**Equality Checking:**\n```go\n// ❌ WRONG\nif actual != expected {\n t.Fatalf(\"Expected %v, got %v\", expected, actual)\n}\n\n// ✅ CORRECT\nassertEqualF(t, actual, expected, \"Values should match\")\n```\n\n**Error Message Validation:**\n```go\n// ❌ WRONG\nif err.Error() != expectedMsg {\n t.Fatalf(\"Expected error: %s, got: %s\", expectedMsg, err.Error())\n}\n\n// ✅ CORRECT\nassertEqualF(t, err.Error(), expectedMsg, \"Error message should match\")\n```\n\n**Boolean Assertions:**\n```go\n// ❌ WRONG\nif !condition {\n t.Fatal(\"Condition should be true\")\n}\n\n// ✅ CORRECT\nassertTrueF(t, condition, \"Condition should be true\")\n```\n\n#### Helper Function Reference:\nAlways examine `assertion_helpers.go` for the latest set of helpers. Consider these existing examples below.\n- `assertNilF/E(t, value, description)` - Assert value is nil\n- `assertNotNilF/E(t, value, description)` - Assert value is not nil \n- `assertEqualF/E(t, actual, expected, description)` - Assert equality\n- `assertNotEqualF/E(t, actual, expected, description)` - Assert inequality\n- `assertTrueF/E(t, value, description)` - Assert boolean is true\n- `assertFalseF/E(t, value, description)` - Assert boolean is false\n- `assertStringContainsF/E(t, str, substring, description)` - Assert string contains substring\n- `assertErrIsF/E(t, actual, expected, description)` - Assert error matches expected error\n\n#### When to Use F vs E:\n- Use `F` suffix (Fatal) for critical failures that should stop the test immediately as well as for preconditions\n- Use `E` suffix (Error) for non-critical failures that allow the test to continue\n\n## Code Review Guidelines:\n- Flag any direct use of `t.Fatal*` or `t.Error*` in new code\n- Ensure all test functions use appropriate assertion helpers\n- Verify that error messages are descriptive and helpful for debugging\n- Check that tests are comprehensive and cover edge cases# Cursor Rules for Go Test Files"
},
{
"path": ".github/CODEOWNERS",
"content": "* @snowflakedb/Client\n\n/transport.go @snowflakedb/pki-oversight @snowflakedb/Client\n/crl.go @snowflakedb/pki-oversight @snowflakedb/Client\n/ocsp.go @snowflakedb/pki-oversight @snowflakedb/Client\n\n# GitHub Advanced Security Secret Scanning config\n/.github/secret_scanning.yml @snowflakedb/prodsec-security-manager-write"
},
{
"path": ".github/ISSUE_TEMPLATE/BUG_REPORT.md",
"content": "---\nname: Bug Report 🐞\nabout: Something isn't working as expected? Here is the right place to report.\nlabels: bug\n---\n\n\n:exclamation: If you need **urgent assistance** then [file a case with Snowflake Support](https://community.snowflake.com/s/article/How-To-Submit-a-Support-Case-in-Snowflake-Lodge).\nOtherwise continue here.\n\n\nPlease answer these questions before submitting your issue. \nIn order to accurately debug the issue this information is required. Thanks!\n\n1. What version of GO driver are you using?\n\n \n2. What operating system and processor architecture are you using?\n\n \n3. What version of GO are you using?\nrun `go version` in your console\n\n4.Server version:* E.g. 1.90.1\nYou may get the server version by running a query:\n```\nSELECT CURRENT_VERSION();\n```\n5. What did you do?\n\n If possible, provide a recipe for reproducing the error.\n A complete runnable program is good.\n\n6. What did you expect to see?\n\n What should have happened and what happened instead?\n\n7. Can you set logging to DEBUG and collect the logs?\n\n https://community.snowflake.com/s/article/How-to-generate-log-file-on-Snowflake-connectors\n \n Before sharing any information, please be sure to review the log and remove any sensitive\n information.\n"
},
{
"path": ".github/ISSUE_TEMPLATE/FEATURE_REQUEST.md",
"content": "---\nname: Feature Request 💡\nabout: Suggest a new idea for the project.\nlabels: feature\n---\n\n\n## What is the current behavior?\n\n## What is the desired behavior?\n\n## How would this improve `gosnowflake`?\n\n## References, Other Background\n\n"
},
{
"path": ".github/ISSUE_TEMPLATE.md",
"content": "### Issue description\nTell us what should happen and what happens instead\n\n### Example code\n```go\nIf possible, please enter some example code here to reproduce the issue.\n```\n\n### Error log\n```\nIf you have an error log, please paste it here.\n```\nAdd ``glog` option to your application to collect log files.\n\n### Configuration\n*Driver version (or git SHA):*\n\n*Go version:* run `go version` in your console\n\n*Server version:* E.g. 1.90.1\nYou may get the server version by running a query:\n```\nSELECT CURRENT_VERSION();\n```\n\n*Client OS:* E.g. Debian 8.1 (Jessie), Windows 10\n"
},
{
"path": ".github/PULL_REQUEST_TEMPLATE.md",
"content": "### Description\n\nSNOW-XXX Please explain the changes you made here.\n\n### Checklist\n- [ ] Added proper logging (if possible)\n- [ ] Created tests which fail without the change (if possible)\n- [ ] Extended the README / documentation, if necessary\n"
},
{
"path": ".github/repo_meta.yaml",
"content": "point_of_contact: @snowflakedb/client\nproduction: true\ncode_owners_file_present: false\njira_area: Developer Platform\n"
},
{
"path": ".github/secret_scanning.yml",
"content": "paths-ignore:\n - \"**/test_data/**\"\n"
},
{
"path": ".github/workflows/build-test.yml",
"content": "name: Build and Test\n\npermissions:\n contents: read\n\non:\n push:\n branches:\n - master\n tags:\n - v*\n pull_request:\n schedule:\n - cron: '7 3 * * *'\n workflow_dispatch:\n inputs:\n goTestParams:\n default:\n description: 'Parameters passed to go test'\n sequentialTests:\n type: boolean\n default: false\n description: 'Run tests sequentially (no buffering, slower)'\n\nconcurrency:\n # older builds for the same pull request numer or branch should be cancelled\n group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}\n cancel-in-progress: true\n\njobs:\n lint:\n runs-on: ubuntu-latest\n name: Check linter\n steps:\n - uses: actions/checkout@v4\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: '1.26'\n - name: golangci-lint\n uses: golangci/golangci-lint-action@v7\n with:\n version: v2.11\n - name: Format, Lint\n shell: bash\n run: ./ci/build.sh\n - name: Run go fix across all platforms and tags\n shell: bash\n run: ./ci/gofix.sh\n build-test-linux:\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n matrix:\n cloud: [ 'AWS', 'AZURE', 'GCP' ]\n go: [ '1.24', '1.25', '1.26' ]\n name: ${{ matrix.cloud }} Go ${{ matrix.go }} on Ubuntu\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: ${{ matrix.go }}\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: ${{ matrix.cloud }}\n GORACE: history_size=7\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n SEQUENTIAL_TESTS: ${{ inputs.sequentialTests }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ./ci/test.sh\n - name: Upload test results to Codecov\n if: ${{!cancelled()}}\n uses: codecov/test-results-action@v1\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n - name: Upload coverage to Codecov\n uses: codecov/codecov-action@v5\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n build-test-linux-no-home:\n runs-on: ubuntu-latest\n name: Ubuntu - no HOME\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: '1.25'\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: AWS\n GORACE: history_size=7\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n SEQUENTIAL_TESTS: ${{ inputs.sequentialTests }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n HOME_EMPTY: \"yes\"\n run: ./ci/test.sh\n build-test-mac:\n runs-on: macos-latest\n strategy:\n fail-fast: false\n matrix:\n cloud: [ 'AWS', 'AZURE', 'GCP' ]\n go: [ '1.24', '1.25', '1.26' ]\n name: ${{ matrix.cloud }} Go ${{ matrix.go }} on Mac\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: ${{ matrix.go }}\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: ${{ matrix.cloud }}\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ./ci/test.sh\n - name: Upload test results to Codecov\n if: ${{!cancelled()}}\n uses: codecov/test-results-action@v1\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n - name: Upload coverage to Codecov\n uses: codecov/codecov-action@v5\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n build-test-mac-no-home:\n runs-on: macos-latest\n name: Mac - no HOME\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: '1.25'\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: AWS\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n HOME_EMPTY: \"yes\"\n run: ./ci/test.sh\n build-test-windows:\n runs-on: windows-latest\n strategy:\n fail-fast: false\n matrix:\n cloud: [ 'AWS', 'AZURE', 'GCP' ]\n go: [ '1.24', '1.25', '1.26' ]\n name: ${{ matrix.cloud }} Go ${{ matrix.go }} on Windows\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: ${{ matrix.go }}\n - uses: actions/setup-python@v5\n with:\n python-version: '3.x'\n architecture: 'x64'\n - name: Test\n shell: cmd\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: ${{ matrix.cloud }}\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n SEQUENTIAL_TESTS: ${{ inputs.sequentialTests }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ci\\\\test.bat\n - name: Upload test results to Codecov\n if: ${{!cancelled()}}\n uses: codecov/test-results-action@v1\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n - name: Upload coverage to Codecov\n uses: codecov/codecov-action@v5\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n fipsOnly:\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n name: FIPS only mode\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: ${{ matrix.go }}\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: ${{ matrix.cloud }}\n GORACE: history_size=7\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n TEST_GODEBUG: fips140=only\n SEQUENTIAL_TESTS: ${{ inputs.sequentialTests }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ./ci/test.sh\n - name: Upload test results to Codecov\n if: ${{!cancelled()}}\n uses: codecov/test-results-action@v1\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n - name: Upload coverage to Codecov\n uses: codecov/codecov-action@v5\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n build-test-linux-minicore-disabled:\n runs-on: ubuntu-latest\n name: Ubuntu - minicore disabled\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: '1.25'\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: AWS\n GORACE: history_size=7\n GO_TEST_PARAMS: ${{ inputs.goTestParams }} -tags=minicore_disabled\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ./ci/test.sh\n build-test-mac-minicore-disabled:\n runs-on: macos-latest\n name: Mac - minicore disabled\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: '1.25'\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: AWS\n GO_TEST_PARAMS: ${{ inputs.goTestParams }} -tags=minicore_disabled\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ./ci/test.sh\n build-test-windows-minicore-disabled:\n runs-on: windows-latest\n name: Windows - minicore disabled\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: '1.25'\n - uses: actions/setup-python@v5\n with:\n python-version: '3.x'\n architecture: 'x64'\n - name: Test\n shell: cmd\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: AWS\n GO_TEST_PARAMS: ${{ inputs.goTestParams }} -tags=minicore_disabled\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ci\\\\test.bat\n ecc:\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n name: Elliptic curves check\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: '1.25'\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: AWS\n GORACE: history_size=7\n GO_TEST_PARAMS: ${{ inputs.goTestParams }} -run TestQueryViaHttps\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n WIREMOCK_ENABLE_ECDSA: true\n run: ./ci/test.sh\n build-test-rockylinux9:\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n matrix:\n cloud_go:\n - cloud: 'AWS'\n go: '1.24.2'\n - cloud: 'AZURE'\n go: '1.25.0'\n - cloud: 'GCP'\n go: '1.26.0'\n name: ${{ matrix.cloud_go.cloud }} Go ${{ matrix.cloud_go.go }} on Rocky Linux 9\n steps:\n - uses: actions/checkout@v4\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: ${{ matrix.cloud_go.cloud }}\n GORACE: history_size=7\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n SEQUENTIAL_TESTS: ${{ inputs.sequentialTests }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ./ci/test_rockylinux9_docker.sh ${{ matrix.cloud_go.go }}\n build-test-ubuntu-arm:\n runs-on: ubuntu-24.04-arm\n strategy:\n fail-fast: false\n matrix:\n cloud_go:\n - cloud: 'AWS'\n go: '1.24'\n - cloud: 'AZURE'\n go: '1.25'\n - cloud: 'GCP'\n go: '1.26'\n name: ${{ matrix.cloud_go.cloud }} Go ${{ matrix.cloud_go.go }} on Ubuntu ARM\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 17\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: ${{ matrix.cloud_go.go }}\n - name: Test\n shell: bash\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: ${{ matrix.cloud_go.cloud }}\n GORACE: history_size=7\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ./ci/test.sh\n - name: Upload test results to Codecov\n if: ${{!cancelled()}}\n uses: codecov/test-results-action@v1\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n - name: Upload coverage to Codecov\n uses: codecov/codecov-action@v5\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n build-test-windows-arm:\n runs-on: windows-11-arm\n strategy:\n fail-fast: false\n matrix:\n cloud_go:\n - cloud: 'AWS'\n go: '1.24'\n - cloud: 'AZURE'\n go: '1.25'\n - cloud: 'GCP'\n go: '1.26'\n name: ${{ matrix.cloud_go.cloud }} Go ${{ matrix.cloud_go.cloud }} on Windows ARM\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-java@v4 # for wiremock\n with:\n java-version: 21\n distribution: 'temurin'\n - name: Setup go\n uses: actions/setup-go@v5\n with:\n go-version: ${{ matrix.cloud_go.go }}\n - uses: actions/setup-python@v5\n with:\n python-version: '3.x'\n architecture: 'x64'\n - name: Test\n shell: cmd\n env:\n PARAMETERS_SECRET: ${{ secrets.PARAMETERS_SECRET }}\n GOLANG_PRIVATE_KEY_SECRET: ${{ secrets.GOLANG_PRIVATE_KEY_SECRET }}\n CLOUD_PROVIDER: ${{ matrix.cloud_go.cloud }}\n GO_TEST_PARAMS: ${{ inputs.goTestParams }}\n WIREMOCK_PORT: 14335\n WIREMOCK_HTTPS_PORT: 13567\n run: ci\\\\test.bat\n - name: Upload test results to Codecov\n if: ${{!cancelled()}}\n uses: codecov/test-results-action@v1\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n - name: Upload coverage to Codecov\n uses: codecov/codecov-action@v5\n with:\n token: ${{ secrets.CODE_COV_UPLOAD_TOKEN }}\n"
},
{
"path": ".github/workflows/changelog.yml",
"content": "name: Changelog Check\n\non:\n pull_request:\n types: [opened, synchronize, labeled, unlabeled]\n\njobs:\n check_change_log:\n runs-on: ubuntu-latest\n if: ${{!contains(github.event.pull_request.labels.*.name, 'NO-CHANGELOG-UPDATES')}}\n steps:\n - name: Checkout\n uses: actions/checkout@v3\n with:\n fetch-depth: 0\n\n - name: Ensure CHANGELOG.md is updated\n run: git diff --name-only --diff-filter=ACMRT ${{ github.event.pull_request.base.sha }} ${{ github.sha }} | grep -wq \"CHANGELOG.md\"\n"
},
{
"path": ".github/workflows/cla_bot.yml",
"content": "name: \"CLA Assistant\"\non:\n issue_comment:\n types: [created]\n pull_request_target:\n types: [opened,closed,synchronize]\n\njobs:\n CLAAssistant:\n runs-on: ubuntu-latest\n permissions:\n actions: write\n contents: write\n pull-requests: write\n statuses: write\n steps:\n - name: \"CLA Assistant\"\n if: (github.event.comment.body == 'recheck' || github.event.comment.body == 'I have read the CLA Document and I hereby sign the CLA') || github.event_name == 'pull_request_target'\n uses: contributor-assistant/github-action/@master\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n PERSONAL_ACCESS_TOKEN : ${{ secrets.CLA_BOT_TOKEN }}\n with:\n path-to-signatures: 'signatures/version1.json'\n path-to-document: 'https://github.com/snowflakedb/CLA/blob/main/README.md'\n branch: 'main'\n allowlist: 'dependabot[bot],github-actions,Jenkins User,_jenkins,sfc-gh-snyk-sca-sa,snyk-bot'\n remote-organization-name: 'snowflake-eng'\n remote-repository-name: 'cla-db'\n"
},
{
"path": ".github/workflows/jira_close.yml",
"content": "name: Jira closure\n\non:\n issues:\n types: [closed, deleted]\n\njobs:\n close-issue:\n runs-on: ubuntu-latest\n steps:\n - name: Extract issue from title\n id: extract\n env:\n TITLE: \"${{ github.event.issue.title }}\"\n run: |\n jira=$(echo -n $TITLE | awk '{print $1}' | sed -e 's/://')\n echo ::set-output name=jira::$jira\n\n - name: Close Jira Issue\n if: startsWith(steps.extract.outputs.jira, 'SNOW-')\n env:\n ISSUE_KEY: ${{ steps.extract.outputs.jira }}\n JIRA_BASE_URL: ${{ secrets.JIRA_BASE_URL }}\n JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}\n JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}\n run: |\n JIRA_API_URL=\"${JIRA_BASE_URL}/rest/api/2/issue/${ISSUE_KEY}/transitions\"\n curl -X POST \\\n --url \"$JIRA_API_URL\" \\\n --user \"${JIRA_USER_EMAIL}:${JIRA_API_TOKEN}\" \\\n --header \"Content-Type: application/json\" \\\n --data \"{\n \\\"update\\\": {\n \\\"comment\\\": [\n { \\\"add\\\": { \\\"body\\\": \\\"Closed on GitHub\\\" } }\n ]\n },\n \\\"fields\\\": {\n \\\"customfield_12860\\\": { \\\"id\\\": \\\"11506\\\" },\n \\\"customfield_10800\\\": { \\\"id\\\": \\\"-1\\\" },\n \\\"customfield_12500\\\": { \\\"id\\\": \\\"11302\\\" },\n \\\"customfield_12400\\\": { \\\"id\\\": \\\"-1\\\" },\n \\\"resolution\\\": { \\\"name\\\": \\\"Done\\\" }\n },\n \\\"transition\\\": { \\\"id\\\": \\\"71\\\" }\n }\"\n"
},
{
"path": ".github/workflows/jira_comment.yml",
"content": "name: Jira comment\n\non:\n issue_comment:\n types: [created]\n\njobs:\n comment-issue:\n runs-on: ubuntu-latest\n steps:\n - name: Jira login\n uses: atlassian/gajira-login@master\n env:\n JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}\n JIRA_BASE_URL: ${{ secrets.JIRA_BASE_URL }}\n JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}\n - name: Extract issue from title\n id: extract\n env:\n TITLE: \"${{ github.event.issue.title }}\"\n run: |\n jira=$(echo -n $TITLE | awk '{print $1}' | sed -e 's/://')\n echo ::set-output name=jira::$jira\n - name: Comment on issue\n uses: atlassian/gajira-comment@master\n if: startsWith(steps.extract.outputs.jira, 'SNOW-') && github.event.comment.user.login != 'codecov[bot]'\n with:\n issue: \"${{ steps.extract.outputs.jira }}\"\n comment: \"${{ github.event.comment.user.login }} commented:\\n\\n${{ github.event.comment.body }}\\n\\n${{ github.event.comment.html_url }}\"\n"
},
{
"path": ".github/workflows/jira_issue.yml",
"content": "name: Jira creation\n\non:\n issues:\n types: [opened]\n issue_comment:\n types: [created]\n\njobs:\n create-issue:\n runs-on: ubuntu-latest\n permissions:\n issues: write\n if: ((github.event_name == 'issue_comment' && github.event.comment.body == 'recreate jira' && github.event.comment.user.login == 'sfc-gh-mkeller') || (github.event_name == 'issues' && github.event.pull_request.user.login != 'whitesource-for-github-com[bot]'))\n steps:\n - name: Create JIRA Ticket\n id: create\n env:\n JIRA_BASE_URL: ${{ secrets.JIRA_BASE_URL }}\n JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}\n JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}\n ISSUE_TITLE: ${{ github.event.issue.title }}\n ISSUE_BODY: ${{ github.event.issue.body }}\n ISSUE_URL: ${{ github.event.issue.html_url }}\n run: |\n # debug\n #set -x\n TMP_BODY=$(mktemp)\n trap \"rm -f $TMP_BODY\" EXIT\n\n # Escape special characters in title and body\n TITLE=$(echo \"${ISSUE_TITLE//`/\\\\`}\" | sed 's/\"/\\\\\"/g' | sed \"s/'/\\\\\\'/g\")\n echo \"${ISSUE_BODY//`/\\\\`}\" | sed 's/\"/\\\\\"/g' | sed \"s/'/\\\\\\'/g\" > $TMP_BODY\n echo -e \"\\n\\n_Created from GitHub Action_ for $ISSUE_URL\" >> $TMP_BODY\n BODY=$(cat \"$TMP_BODY\")\n\n PAYLOAD=$(jq -n \\\n --arg issuetitle \"$TITLE\" \\\n --arg issuebody \"$BODY\" \\\n '{\n fields: {\n project: { key: \"SNOW\" },\n issuetype: { name: \"Bug\" },\n summary: $issuetitle,\n description: $issuebody,\n customfield_11401: { id: \"14723\" },\n assignee: { id: \"712020:e527ae71-55cc-4e02-9217-1ca4ca8028a2\" },\n components: [{ id: \"19286\" }],\n labels: [\"oss\"],\n priority: { id: \"10001\" }\n }\n }')\n\n # Create JIRA issue using REST API\n RESPONSE=$(curl -s -X POST \\\n -H \"Content-Type: application/json\" \\\n -H \"Accept: application/json\" \\\n -u \"$JIRA_USER_EMAIL:$JIRA_API_TOKEN\" \\\n \"$JIRA_BASE_URL/rest/api/2/issue\" \\\n -d \"$PAYLOAD\")\n\n # Extract JIRA issue key from response\n JIRA_KEY=$(echo \"$RESPONSE\" | jq -r '.key')\n\n if [ \"$JIRA_KEY\" = \"null\" ] || [ -z \"$JIRA_KEY\" ]; then\n echo \"Failed to create JIRA issue\"\n echo \"Response: $RESPONSE\"\n echo \"Request payload: $PAYLOAD\"\n exit 1\n fi\n\n echo \"Created JIRA issue: $JIRA_KEY\"\n echo \"jira_key=$JIRA_KEY\" >> $GITHUB_OUTPUT\n\n - name: Update GitHub Issue\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n REPOSITORY: ${{ github.repository }}\n ISSUE_NUMBER: ${{ github.event.issue.number }}\n JIRA_KEY: ${{ steps.create.outputs.jira_key }}\n ISSUE_TITLE: ${{ github.event.issue.title }}\n run: |\n TITLE=$(echo \"${ISSUE_TITLE//`/\\\\`}\" | sed 's/\"/\\\\\"/g' | sed \"s/'/\\\\\\'/g\")\n PAYLOAD=$(jq -n \\\n --arg issuetitle \"$TITLE\" \\\n --arg jirakey \"$JIRA_KEY\" \\\n '{\n title: ($jirakey + \": \" + $issuetitle)\n }')\n\n # Update Github issue title with jira id\n curl -s \\\n -X PATCH \\\n -H \"Authorization: Bearer $GITHUB_TOKEN\" \\\n -H \"Accept: application/vnd.github+json\" \\\n -H \"X-GitHub-Api-Version: 2022-11-28\" \\\n \"https://api.github.com/repos/$REPOSITORY/issues/$ISSUE_NUMBER\" \\\n -d \"$PAYLOAD\"\n\n if [ \"$?\" != 0 ]; then\n echo \"Failed to update GH issue. Payload was:\"\n echo \"$PAYLOAD\"\n exit 1\n fi\n"
},
{
"path": ".github/workflows/semgrep.yml",
"content": "name: Run semgrep checks\n\non:\n pull_request:\n branches: [main, master]\n\npermissions:\n contents: read\n\njobs:\n run-semgrep-reusable-workflow:\n uses: snowflakedb/reusable-workflows/.github/workflows/semgrep-v2.yml@main\n secrets:\n token: ${{ secrets.SEMGREP_APP_TOKEN }}\n"
},
{
"path": ".gitignore",
"content": "*.DS_Store\n.idea/\n.vscode/\nparameters*.json\nparameters*.bat\n*.p8\ncoverage.txt\nfuzz-*/\n/select1\n/selectmany\n/verifycert\nwss-golang-agent.config\nwss-unified-agent.jar\nwhitesource/\n*.swp\ncp.out\n__debug_bin*\ntest-output.txt\ntest-report.junit.xml\n\n# exclude vendor\nvendor\n\n# SSH private key for WIF tests\nci/wif/parameters/rsa_wif_aws_azure\nci/wif/parameters/rsa_wif_gcp\n"
},
{
"path": ".golangci.yml",
"content": "version: \"2\"\n\nrun:\n tests: true\n\nlinters:\n exclusions:\n rules:\n - path: \"_test.go\"\n linters:\n - errcheck\n - path: \"cmd/\"\n linters:\n - errcheck\n - path: \"_test.go\"\n linters:\n - staticcheck\n text: \"implement StmtQueryContext\"\n - path: \"_test.go\"\n linters:\n - staticcheck\n text: \"implement StmtExecContext\"\n - linters:\n - staticcheck\n text: \"QF1001\"\n - linters:\n - staticcheck\n text: \"SA1019: .+\\\\.(LoginTimeout|RequestTimeout|JWTExpireTimeout|ClientTimeout|JWTClientTimeout|ExternalBrowserTimeout|CloudStorageTimeout|Tracing) is deprecated\""
},
{
"path": ".pre-commit-config.yaml",
"content": "repos:\n- repo: git@github.com:snowflakedb/casec_precommit.git # SSH\n# - repo: https://github.com/snowflakedb/casec_precommit.git # HTTPS\n rev: v1.5\n hooks:\n - id: snapps-secret-scanner\n"
},
{
"path": ".windsurf/rules/go.md",
"content": "---\ntrigger: glob\ndescription: \nglobs: **/*.go\n---\n\n# Go files rules\n\n## General\n\n1. Unless it's necessary or told otherwise, try reusing existing files, both for implementation and tests.\n2. If possible, try running relevant tests.\n\n## Tests\n\n1. Create a test file with the name same as prod code file by default.\n2. For assertions use our test helpers defined in assert_test.go.\n\n## Logging\n\n1. Add reasonable logging - don't repeat logs, but add them when it's meaningful.\n2. Always consider log levels."
},
{
"path": "CHANGELOG.md",
"content": "# Changelog\n\n## Upcoming release\n\nBug fixes:\n\n- Fixed empty `Account` when connecting with programmatic `Config` and `database/sql.Connector` by deriving `Account` from the first DNS label of `Host` in `FillMissingConfigParameters` when `Host` matches the Snowflake hostname pattern (snowflakedb/gosnowflake#1772).\n\n## 2.0.1\n\nBug fixes:\n\n- Fixed default `CrlDownloadMaxSize` to be 20MB instead of 200MB, as the previous value was set too high and could cause out-of-memory issues (snowflakedb/gosnowflake#1735).\n- Replaced global `paramsMutex` with per-connection `syncParams` to encapsulate parameter synchronization and avoid cross-connection contention (snowflakedb/gosnoflake#1747).\n- `Config.Params` map is not modified anymore, to avoid changing parameter values across connections of the same connection pool (snowflakedb/gosnowflake#1747).\n- Set `BlobContentMD5` on Azure uploads so that multi-part uploads have the blob content-MD5 property populated (snowflakedb/gosnowflake#1757).\n- Fixed 403 errors from Google/GCP/GCS PUT queries on versioned stages (snowflakedb/gosnowflake#1760).\n- Fixed not updating query context cache for failed queries (snowflakedb/gosnowflake#1763).\n\nInternal changes:\n\n- Moved configuration to a dedicated internal package (snowflakedb/gosnowflake#1720).\n- Modernized Go syntax idioms throughout the codebase.\n- Added libc family, version and dynamic linking marker to client environment telemetry (snowflakedb/gosnowflake#1750).\n- Bumped a few libraries to fix vulnerabilities (snowflakedb/gosnowflake#1751, snowflakedb/gosnowflake#1756).\n- Depointerised query context cache in `snowflakeConn` (snowflakedb/gosnowflake#1763).\n\n## 2.0.0\n\nBreaking changes:\n\n- Removed `RaisePutGetError` from `SnowflakeFileTransferOptions` - current behaviour is aligned to always raise errors for PUT/GET operations (snowflakedb/gosnowflake#1690).\n- Removed `GetFileToStream` from `SnowflakeFileTransferOptions` - using `WithFileGetStream` automatically enables file streaming for GETs (snowflakedb/gosnowflake#1690).\n- Renamed `WithFileStream` to `WithFilePutStream` for consistency (snowflakedb/gosnowflake#1690).\n- `Array` function now returns error for unsupported types (snowflakedb/gosnowflake#1693).\n- `WithMultiStatement` does not return error anymore (snowflakedb/gosnowflake#1693).\n- `WithOriginalTimestamp` is removed, use `WithArrowBatchesTimestampOption(UseOriginalTimestamp)` instead (snowflakedb/gosnowflake#1693).\n- `WithMapValuesNullable` and `WithArrayValuesNullable` combined into one option `WithEmbeddedValuesNullable` (snowflakedb/gosnowflake#1693).\n- Hid streaming chunk downloader. It will be removed completely in the future (snowflakedb/gosnowflake#1696).\n- Maximum number of chunk download goroutines is now configured with `CLIENT_PREFETCH_THREADS` session parameter (snowflakedb/gosnowflake#1696)\n and default to 4.\n- Fixed typo in `GOSNOWFLAKE_SKIP_REGISTRATION` env variable (snowflakedb/gosnowflake#1696).\n- Removed `ClientIP` field from `Config` struct. This field was never used and is not needed for any functionality (snowflakedb/gosnowflake#1692).\n- Unexported MfaToken and IdToken (snowflakedb/gosnowflake#1692).\n- Removed `InsecureMode` field from `Config` struct. Use `DisableOCSPChecks` instead (snowflakedb/gosnowflake#1692).\n- Renamed `KeepSessionAlive` field in `Config` struct to `ServerSessionKeepAlive` to adjust with the remaining drivers (snowflakedb/gosnowflake#1692).\n- Removed `DisableTelemetry` field from `Config` struct. Use `CLIENT_TELEMETRY_ENABLED` session parameter instead (snowflakedb/gosnowflake#1692).\n- Removed stream chunk downloader. Use a regular, default downloader instead. (snowflakedb/gosnowflake#1702).\n- Removed `SnowflakeTransport`. Use `Config.Transporter` or simply register your own TLS config with `RegisterTLSConfig` if you just need a custom root certificates set (snowflakedb/gosnowflake#1703).\n- Arrow batches changes (snowflakedb/gosnowflake#1706):\n - Arrow batches have been extracted to a separate package. It should significantly drop the compilation size for those who don't need arrow batches (~34MB -> ~18MB).\n - Removed `GetArrowBatches` from `SnowflakeRows` and `SnowflakeResult`. Use `arrowbatches.GetArrowBatches(rows.(SnowflakeRows))` instead.\n - Migrated functions:\n - `sf.WithArrowBatchesTimestampOption` -> `arrowbatches.WithTimstampOption`\n - `sf.WithArrowBatchesUtf8Validation` -> `arrowbatches.WithUtf8Validation`\n - `sf.ArrowSnowflakeTimestampToTime` -> `arrowbatches.ArrowSnowflakeTimestampToTime`\n- Logging changes (snowflakedb/gosnowflake#1710):\n - Removed Logrus logger and migrated to slog.\n - Simplified `SFLogger` interface.\n - Added `SFSlogLogger` interface for setting custom slog handler.\n\nBug fixes:\n\n- The query `context.Context` is now propagated to cloud storage operations for PUT and GET queries, allowing for better cancellation handling (snowflakedb/gosnowflake#1690).\n\nNew features:\n\n- Added support for Go 1.26, dropped support for Go 1.23 (snowflakedb/gosnowflake#1707).\n- Added support for FIPS-only mode (snowflakedb/gosnowflake#1496).\n\nBug fixes:\n\n- Added panic recovery block for stage file uploads and downloads operation (snowflakedb/gosnowflake#1687).\n- Fixed WIF metadata request from Azure container, manifested with HTTP 400 error (snowflakedb/gosnowflake#1701).\n- Fixed SAML authentication port validation bypass in `isPrefixEqual` where the second URL's port was never checked (snowflakedb/gosnowflake#1712).\n- Fixed a race condition in OCSP cache clearer (snowflakedb/gosnowflake#1704).\n- The query `context.Context` is now propagated to cloud storage operations for PUT and GET queries, allowing for better cancellation handling (snowflakedb/gosnowflake#1690).\n- Fixed `tokenFilePath` DSN parameter triggering false validation error claiming both `token` and `tokenFilePath` were specified when only `tokenFilePath` was provided in the DSN string (snowflakedb/gosnowflake#1715).\n- Fixed minicore crash (SIGFPE) on fully statically linked Linux binaries by detecting static linking via ELF PT_INTERP inspection and skipping `dlopen` gracefully (snowflakedb/gosnowflake#1721).\n\nInternal changes:\n\n- Moved configuration to a dedicated internal package (snowflakedb/gosnowflake#1720).\n\n## 1.19.0\n\nNew features:\n\n- Added ability to disable minicore loading at compile time (snowflakedb/gosnowflake#1679).\n- Exposed `tokenFilePath` in `Config` (snowflakedb/gosnowflake#1666).\n- `tokenFilePath` is now read for every new connection (snowflakedb/gosnowflake#1666).\n- Added support for identity impersonation when using workload identity federation (snowflakedb/gosnowflake#1652, snowflakedb/gosnowflake#1660).\n\nBug fixes:\n\n- Fixed getting file from an unencrypted stage (snowflakedb/gosnowflake#1672).\n- Fixed minicore file name gathering in client environment (snowflakedb/gosnowflake#1661).\n- Fixed file descriptor leaks in cloud storage calls (snowflakedb/gosnowflake#1682)\n- Fixed path escaping for GCS urls (snowflakedb/gosnowflake#1678).\n\nInternal changes:\n\n- Improved Linux telemetry gathering (snowflakedb/gosnowflake#1677).\n- Improved some logs returned from cloud storage clients (snowflakedb/gosnowflake#1665).\n\n## 1.18.1\n\nBug fixes:\n\n- Handle HTTP307 & 308 in drivers to achieve better resiliency to backend errors (snowflakedb/gosnowflake#1616).\n- Create temp directory only if needed during file transfer (snowflakedb/gosnowflake#1647)\n- Fix unnecessary user expansion for file paths (snowflakedb/gosnowflake#1646).\n\nInternal changes:\n- Remove spammy \"telemetry disabled\" log messages (snowflakedb/gosnowflake#1638).\n- Introduced shared library ([source code](https://github.com/snowflakedb/universal-driver/tree/main/sf_mini_core)) for extended telemetry to identify and prepare testing platform for native rust extensions (snowflakedb/gosnowflake#1629)\n\n## 1.18.0\n\nNew features:\n\n- Added validation of CRL `NextUpdate` for freshly downloaded CRLs (snowflakedb/gosnowflake#1617)\n- Exposed function to send arbitrary telemetry data (snowflakedb/gosnowflake#1627)\n- Added logging of query text and parameters (snowflakedb/gosnowflake#1625)\n\nBug fixes:\n\n- Fixed a data race error in tests caused by platform_detection init() function (snowflakedb/gosnowflake#1618)\n- Make secrets detector initialization thread safe and more maintainable (snowflakedb/gosnowflake#1621)\n\nInternal changes:\n\n- Added ISA to login request telemetry (snowflakedb/gosnowflake#1620)\n\n## 1.17.1\n\n- Fix unsafe reflection of nil pointer on DECFLOAT func in bind uploader (snowflakedb/gosnowflake#1604).\n- Added temporary download files cleanup (snowflakedb/gosnowflake#1577)\n- Marked fields as deprecated (snowflakedb/gosnowflake#1556)\n- Exposed `QueryStatus` from `SnowflakeResult` and `SnowflakeRows` in `GetStatus()` function (snowflakedb/gosnowflake#1556)\n- Split timeout settings into separate groups based on target service types (snowflakedb/gosnowflake#1531)\n- Added small clarification in oauth.go example on token escaping (snowflakedb/gosnowflake#1574)\n- Ensured proper permissions for CRL cache directory (snowflakedb/gosnowflake#1588)\n- Added `CrlDownloadMaxSize` to limit the size of CRL downloads (snowflakedb/gosnowflake#1588)\n- Added platform telemetry to login requests. Can be disabled with `SNOWFLAKE_DISABLE_PLATFORM_DETECTION` environment variable (snowflakedb/gosnowflake#1601)\n- Bypassed proxy settings for WIF metadata requests (snowflakedb/gosnowflake#1593)\n- Fixed a bug where GCP PUT/GET operations would fail when the connection context was cancelled (snowflakedb/gosnowflake#1584)\n- Fixed nil pointer dereference while calling long-running queries (snowflakedb/gosnowflake#1592) (snowflakedb/gosnowflake#1596)\n- Moved keyring-based secure storage manager into separate file to avoid the need to initialize keyring on Linux. (snowflakedb/gosnowflake#1595)\n- Enabling official support for RHEL9 by testing and enabling CI/CD checks for Rocky Linux in CICD, (snowflakedb/gosnowflake#1597)\n- Improve logging (snowflakedb/gosnowflake#1570)\n\n## 1.17.0\n\n- Added ability to configure OCSP per connection (snowflakedb/gosnowflake#1528)\n- Added `DECFLOAT` support, see details in `doc.go` (snowflakedb/gosnowflake#1504, snowflakedb/gosnowflake#1506)\n- Added support for Go 1.25, dropped support for Go 1.22 (snowflakedb/gosnowflake#1544)\n- Added proxy options to connection parameters (snowflakedb/gosnowflake#1511)\n- Added `client_session_keep_alive_heartbeat_frequency` connection param (snowflakedb/gosnowflake#1576)\n- Added support for multi-part downloads for S3, Azure and GCP (snowflakedb/gosnowflake#1549)\n- Added `singleAuthenticationPrompt` to control whether only one authentication should be performed at the same time for authentications that need human interactions (like MFA or OAuth authorization code). Default is true. (snowflakedb/gosnowflake#1561)\n- Fixed missing `DisableTelemetry` option in connection parameters (snowflakedb/gosnowflake#1520)\n- Fixed multistatements in large result sets (snowflakedb/gosnowflake#1539, snowflakedb/gosnowflake#1543, snowflakedb/gosnowflake#1547)\n- Fixed unnecessary retries when context is cancelled (snowflakedb/gosnowflake#1540)\n- Fixed regression in TOML connection file (snowflakedb/gosnowflake#1530)\n\n## Prior Releases\n\nRelease notes available at https://docs.snowflake.com/en/release-notes/clients-drivers/golang\n"
},
{
"path": "CONTRIBUTING.md",
"content": "# Contributing Guidelines\n\n## Reporting Issues\n\nBefore creating a new Issue, please check first if a similar Issue [already exists](https://github.com/snowflakedb/gosnowflake/issues?state=open) or was [recently closed](https://github.com/snowflakedb/gosnowflake/issues?direction=desc&page=1&sort=updated&state=closed).\n\n## Contributing Code\n\nBy contributing to this project, you share your code under the Apache License 2, as specified in the LICENSE file.\n\n### Code Review\n\nEveryone is invited to review and comment on pull requests.\nIf it looks fine to you, comment with \"LGTM\" (Looks good to me).\n\nIf changes are required, notice the reviewers with \"PTAL\" (Please take another look) after committing the fixes.\n\nBefore merging the Pull Request, at least one Snowflake team member must have commented with \"LGTM\".\n"
},
{
"path": "Jenkinsfile",
"content": "@Library('pipeline-utils')\nimport com.snowflake.DevEnvUtils\nimport groovy.json.JsonOutput\n\n\ntimestamps {\n node('high-memory-node') {\n stage('checkout') {\n scmInfo = checkout scm\n println(\"${scmInfo}\")\n env.GIT_BRANCH = scmInfo.GIT_BRANCH\n env.GIT_COMMIT = scmInfo.GIT_COMMIT\n }\n params = [\n string(name: 'svn_revision', value: 'temptest-deployed'),\n string(name: 'branch', value: 'main'),\n string(name: 'client_git_commit', value: scmInfo.GIT_COMMIT),\n string(name: 'client_git_branch', value: scmInfo.GIT_BRANCH),\n string(name: 'TARGET_DOCKER_TEST_IMAGE', value: 'go-chainguard-go1_24'),\n string(name: 'parent_job', value: env.JOB_NAME),\n string(name: 'parent_build_number', value: env.BUILD_NUMBER)\n ]\n \n stage('Authenticate Artifactory') {\n script {\n new DevEnvUtils().withSfCli {\n sh \"sf artifact oci auth\"\n }\n }\n }\n\n parallel(\n 'Test': {\n stage('Test') {\n build job: 'RT-LanguageGo-PC', parameters: params\n }\n },\n 'Test Authentication': {\n stage('Test Authentication') {\n withCredentials([\n string(credentialsId: 'sfctest0-parameters-secret', variable: 'PARAMETERS_SECRET')\n ]) {\n sh '''\\\n |#!/bin/bash -e\n |$WORKSPACE/ci/test_authentication.sh\n '''.stripMargin()\n }\n }\n },\n 'Test WIF Auth': {\n stage('Test WIF Auth') {\n withCredentials([\n string(credentialsId: 'sfctest0-parameters-secret', variable: 'PARAMETERS_SECRET'),\n ]) {\n sh '''\\\n |#!/bin/bash -e\n |$WORKSPACE/ci/test_wif.sh\n '''.stripMargin()\n }\n }\n },\n 'Test Revocation Validation': {\n stage('Test Revocation Validation') {\n withCredentials([\n usernamePassword(credentialsId: 'jenkins-snowflakedb-github-app',\n usernameVariable: 'GITHUB_USER',\n passwordVariable: 'GITHUB_TOKEN')\n ]) {\n try {\n sh '''\\\n |#!/bin/bash -e\n |chmod +x $WORKSPACE/ci/test_revocation.sh\n |$WORKSPACE/ci/test_revocation.sh\n '''.stripMargin()\n } finally {\n archiveArtifacts artifacts: 'revocation-results.json,revocation-report.html', allowEmptyArchive: true\n publishHTML(target: [\n allowMissing: true,\n alwaysLinkToLastBuild: true,\n keepAll: true,\n reportDir: '.',\n reportFiles: 'revocation-report.html',\n reportName: 'Revocation Validation Report'\n ])\n }\n }\n }\n }\n )\n }\n}\n\n\npipeline {\n agent { label 'high-memory-node' }\n options { timestamps() }\n environment {\n COMMIT_SHA_LONG = sh(returnStdout: true, script: \"echo \\$(git rev-parse \" + \"HEAD)\").trim()\n\n // environment variables for semgrep_agent (for findings / analytics page)\n // remove .git at the end\n // remove SCM URL + .git at the end\n\n BASELINE_BRANCH = \"${env.CHANGE_TARGET}\"\n }\n stages {\n stage('Checkout') {\n steps {\n checkout scm\n }\n }\n }\n}\n\ndef wgetUpdateGithub(String state, String folder, String targetUrl, String seconds) {\n def ghURL = \"https://api.github.com/repos/snowflakedb/gosnowflake/statuses/$COMMIT_SHA_LONG\"\n def data = JsonOutput.toJson([state: \"${state}\", context: \"jenkins/${folder}\",target_url: \"${targetUrl}\"])\n sh \"wget ${ghURL} --spider -q --header='Authorization: token $GIT_PASSWORD' --post-data='${data}'\"\n}\n"
},
{
"path": "LICENSE",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"{}\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright (c) 2017-2022 Snowflake Computing Inc. All rights reserved.\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "Makefile",
"content": "NAME:=gosnowflake\nVERSION:=$(shell git describe --tags --abbrev=0)\nREVISION:=$(shell git rev-parse --short HEAD)\nCOVFLAGS:=\n\n## Run fmt, lint and test\nall: fmt lint cov\n\ninclude gosnowflake.mak\n\n## Run tests\ntest_setup: test_teardown\n\tpython3 ci/scripts/hang_webserver.py 12345 &\n\ntest_teardown:\n\tpkill -9 hang_webserver || true\n\ntest: deps test_setup\n\t./ci/scripts/execute_tests.sh\n\n## Run Coverage tests\ncov:\n\tmake test COVFLAGS=\"-coverprofile=coverage.txt -covermode=atomic\"\n\n\n\n## Lint\nlint: clint\n\n## Format source codes\nfmt: cfmt\n\t@for c in $$(ls cmd); do \\\n\t\t(cd cmd/$$c; make fmt); \\\n\tdone\n\n## Install sample programs\ninstall:\n\tfor c in $$(ls cmd); do \\\n\t\t(cd cmd/$$c; GOBIN=$$GOPATH/bin go install $$c.go); \\\n\tdone\n\n## Build fuzz tests\nfuzz-build:\n\tfor c in $$(ls | grep -E \"fuzz-*\"); do \\\n\t\t(cd $$c; make fuzz-build); \\\n\tdone\n\n## Run fuzz-dsn\nfuzz-dsn:\n\t(cd fuzz-dsn; go-fuzz -bin=./dsn-fuzz.zip -workdir=.)\n\n.PHONY: setup deps update test lint help fuzz-dsn\n"
},
{
"path": "README.md",
"content": "## Migrating to v2\n\n**Version 2.0.0 of the Go Snowflake Driver was released on March 3rd, 2026.** This major version includes breaking changes that require code updates when migrating from v1.x.\n\n### Key Changes and Migration Steps\n\n#### 1. Update Import Paths\n\nUpdate your `go.mod` to use v2:\n\n```sh\ngo get -u github.com/snowflakedb/gosnowflake/v2\n```\n\nUpdate imports in your code:\n\n```go\n// Old (v1)\nimport \"github.com/snowflakedb/gosnowflake\"\n\n// New (v2)\nimport \"github.com/snowflakedb/gosnowflake/v2\"\n```\n\n#### 2. Arrow Batches Moved to Separate Package\n\nThe public Arrow batches API now lives in `github.com/snowflakedb/gosnowflake/v2/arrowbatches`.\nImporting that sub-package pulls in the additional Arrow compute dependency only for applications\nthat use Arrow batches directly.\n\n**Migration:**\n\n```go\nimport (\n \"context\"\n \"database/sql/driver\"\n\n sf \"github.com/snowflakedb/gosnowflake/v2\"\n \"github.com/snowflakedb/gosnowflake/v2/arrowbatches\"\n)\n\nctx := arrowbatches.WithArrowBatches(context.Background())\n\nvar rows driver.Rows\nerr := conn.Raw(func(x any) error {\n rows, err = x.(driver.QueryerContext).QueryContext(ctx, query, nil)\n return err\n})\nif err != nil {\n // handle error\n}\n\nbatches, err := arrowbatches.GetArrowBatches(rows.(sf.SnowflakeRows))\nif err != nil {\n // handle error\n}\n```\n\n**Optional helper mapping:**\n- `sf.WithArrowBatchesTimestampOption` → `arrowbatches.WithTimestampOption`\n- `sf.WithArrowBatchesUtf8Validation` → `arrowbatches.WithUtf8Validation`\n- `sf.ArrowSnowflakeTimestampToTime` → `arrowbatches.ArrowSnowflakeTimestampToTime`\n- `sf.WithOriginalTimestamp` → `arrowbatches.WithTimestampOption(ctx, arrowbatches.UseOriginalTimestamp)`\n\n#### 3. Configuration Struct Changes\n\n**Renamed fields:**\n```go\n// Old (v1)\nconfig := &gosnowflake.Config{\n KeepSessionAlive: true,\n InsecureMode: true,\n DisableTelemetry: true,\n}\n\n// New (v2)\nconfig := &gosnowflake.Config{\n ServerSessionKeepAlive: true, // Renamed for consistency with other drivers\n DisableOCSPChecks: true, // Replaces InsecureMode\n // DisableTelemetry removed - use CLIENT_TELEMETRY_ENABLED session parameter\n}\n```\n\n**Removed fields:**\n- `ClientIP` - No longer used\n- `MfaToken` and `IdToken` - Now unexported\n- `DisableTelemetry` - Use `CLIENT_TELEMETRY_ENABLED` session parameter instead\n\n#### 4. Logger Changes\n\nThe built-in logger is now based on Go's standard `log/slog`:\n\n```go\nlogger := gosnowflake.GetLogger()\n_ = logger.SetLogLevel(\"debug\")\n```\n\nFor custom logging, continue implementing `SFLogger`.\nIf you want to customize the built-in slog handler, type-assert `GetLogger()` to `SFSlogLogger`\nand call `SetHandler`.\n\n#### 5. File Transfer Changes\n\n**Configuration options:**\n\n```go\n// Old (v1)\noptions := &gosnowflake.SnowflakeFileTransferOptions{\n RaisePutGetError: true,\n GetFileToStream: true,\n}\nctx = gosnowflake.WithFileStream(ctx, stream)\n\n// New (v2)\n// RaisePutGetError removed - errors always raised\n// GetFileToStream removed - use WithFileGetStream instead\nctx = gosnowflake.WithFilePutStream(ctx, stream) // Renamed from WithFileStream\nctx = gosnowflake.WithFileGetStream(ctx, stream) // For GET operations\n```\n\n#### 6. Context and Function Changes\n\n```go\n// Old (v1)\nctx, err := gosnowflake.WithMultiStatement(ctx, 0)\nif err != nil {\n // handle error\n}\n\n// New (v2)\nctx = gosnowflake.WithMultiStatement(ctx, 0) // No error returned\n```\n\n```go\n// Old (v1)\nvalues := gosnowflake.Array(data)\n\n// New (v2)\nvalues, err := gosnowflake.Array(data) // Now returns error for unsupported types\nif err != nil {\n // handle error\n}\n```\n\n#### 7. Nullable Options Combined\n\n```go\n// Old (v1)\nctx = gosnowflake.WithMapValuesNullable(ctx)\nctx = gosnowflake.WithArrayValuesNullable(ctx)\n\n// New (v2)\nctx = gosnowflake.WithEmbeddedValuesNullable(ctx) // Handles both maps and arrays\n```\n\n#### 8. Session Parameter Changes\n\n**Chunk download workers:**\n\n```go\n// Old (v1)\ngosnowflake.MaxChunkDownloadWorkers = 10 // Global variable\n\n// New (v2)\n// Configure via CLIENT_PREFETCH_THREADS session parameter.\n// NOTE: The default is 4.\ndb.Exec(\"ALTER SESSION SET CLIENT_PREFETCH_THREADS = 10\")\n```\n\n#### 9. Transport Configuration\n\n```go\nimport \"crypto/tls\"\n\n// Old (v1)\ngosnowflake.SnowflakeTransport = yourTransport\n\n// New (v2)\nconfig := &gosnowflake.Config{\n Transporter: yourCustomTransport,\n}\n\n// Or, if you only need custom TLS settings/certificates:\ntlsConfig := &tls.Config{\n // ...\n}\n_ = gosnowflake.RegisterTLSConfig(\"custom\", tlsConfig)\nconfig.TLSConfigName = \"custom\"\n```\n\n#### 10. Environment Variable Fix\n\nIf you use the skip registration environment variable:\n\n```sh\n# Old (v1)\nGOSNOWFLAKE_SKIP_REGISTERATION=true # Note the typo\n\n# New (v2)\nGOSNOWFLAKE_SKIP_REGISTRATION=true # Typo fixed\n```\n\n### Additional Resources\n\n- Full list of changes: See [CHANGELOG.md](./CHANGELOG.md)\n- Questions or issues: [GitHub Issues](https://github.com/snowflakedb/gosnowflake/issues)\n\n\n## Support\n\nFor official support and urgent, production-impacting issues, please [contact Snowflake Support](https://community.snowflake.com/s/article/How-To-Submit-a-Support-Case-in-Snowflake-Lodge).\n\n# Go Snowflake Driver\n\n\n ![\"Coverage\"](\"https://codecov.io/github/snowflakedb/gosnowflake/coverage.svg?branch=master\")
\n\n\n ![](\"https://github.com/snowflakedb/gosnowflake/workflows/Build%20and%20Test/badge.svg?branch=master\")
\n\n\n ![](\"http://img.shields.io/:license-Apache%202-brightgreen.svg\")
\n\n\n ![](\"https://goreportcard.com/badge/github.com/snowflakedb/gosnowflake\")
\n\n\nThis topic provides instructions for installing, running, and modifying the Go Snowflake Driver. The driver supports Go's [database/sql](https://golang.org/pkg/database/sql/) package.\n\n# Prerequisites\n\nThe following software packages are required to use the Go Snowflake Driver.\n\n## Go\n\nThe latest driver requires the [Go language](https://golang.org/) 1.24 or higher. The supported operating systems are 64-bits Linux, Mac OS, and Windows, but you may run the driver on other platforms if the Go language works correctly on those platforms.\n\n# Installation\n\nIf you don't have a project initialized, set it up.\n\n```sh\ngo mod init example.com/snowflake\n```\n\nGet Gosnowflake source code, if not installed.\n\n```sh\ngo get -u github.com/snowflakedb/gosnowflake/v2\n```\n\n# Docs\n\nFor detailed documentation and basic usage examples, please see the documentation at\n[godoc.org](https://godoc.org/github.com/snowflakedb/gosnowflake/v2).\n\n## Notes\n\nThis driver currently does not support GCP regional endpoints. Please ensure that any workloads using through this driver do not require support for regional endpoints on GCP. If you have questions about this, please contact Snowflake Support.\n\nThe driver uses Rust library called sf_mini_core, you can find its source code [here](https://github.com/snowflakedb/universal-driver/tree/main/sf_mini_core)\n\n# Sample Programs\n\nSnowflake provides a set of sample programs to test with. Set the environment variable ``$GOPATH`` to the top directory of your workspace, e.g., ``~/go`` and make certain to\ninclude ``$GOPATH/bin`` in the environment variable ``$PATH``. Run the ``make`` command to build all sample programs.\n\n```sh\nmake install\n```\n\nIn the following example, the program ``select1.go`` is built and installed in ``$GOPATH/bin`` and can be run from the command line:\n\n```sh\nSNOWFLAKE_TEST_ACCOUNT= \\\nSNOWFLAKE_TEST_USER= \\\nSNOWFLAKE_TEST_PASSWORD= \\\nselect1\nCongrats! You have successfully run SELECT 1 with Snowflake DB!\n```\n\n# Development\n\nThe developer notes are hosted with the source code on [GitHub](https://github.com/snowflakedb/gosnowflake/v2).\n\n## Testing Code\n\n\nSet the Snowflake connection info in ``parameters.json``:\n\n```json\n{\n \"testconnection\": {\n \"SNOWFLAKE_TEST_USER\": \"\",\n \"SNOWFLAKE_TEST_PASSWORD\": \"\",\n \"SNOWFLAKE_TEST_ACCOUNT\": \"\",\n \"SNOWFLAKE_TEST_WAREHOUSE\": \"\",\n \"SNOWFLAKE_TEST_DATABASE\": \"\",\n \"SNOWFLAKE_TEST_SCHEMA\": \"\",\n \"SNOWFLAKE_TEST_ROLE\": \"\",\n \"SNOWFLAKE_TEST_DEBUG\": \"false\"\n }\n}\n```\n\nInstall [jq](https://stedolan.github.io/jq) so that the parameters can get parsed correctly, and run ``make test`` in your Go development environment:\n\n```sh\nmake test\n```\n\n### Setting debug mode during tests\nThis is for debugging Large SQL statements (greater than 300 characters). If you want to enable debug mode, set `SNOWFLAKE_TEST_DEBUG` to `true` in `parameters.json`, or export it in your shell instance.\n\n## customizing Logging Tags\n\nIf you would like to ensure that certain tags are always present in the logs, `RegisterClientLogContextHook` can be used in your init function. See example below.\n```go\nimport \"github.com/snowflakedb/gosnowflake/v2\"\n\nfunc init() {\n // each time the logger is used, the logs will contain a REQUEST_ID field with requestID the value extracted \n // from the context\n\tgosnowflake.RegisterClientLogContextHook(\"REQUEST_ID\", func(ctx context.Context) interface{} {\n\t\treturn requestIdFromContext(ctx)\n\t})\n}\n```\n\n## Setting Log Level\nIf you want to change the log level, `SetLogLevel` can be used in your init function like this:\n```go\nimport \"github.com/snowflakedb/gosnowflake/v2\"\n\nfunc init() {\n // The following line changes the log level to debug\n\t_ = gosnowflake.GetLogger().SetLogLevel(\"debug\")\n}\n```\nThe following is a list of options you can pass in to set the level from least to most verbose:\n- `\"OFF\"`\n- `\"fatal\"`\n- `\"error\"`\n- `\"warn\"`\n- `\"info\"`\n- `\"debug\"`\n- `\"trace\"`\n\n\n## Capturing Code Coverage\n\nConfigure your testing environment as described above and run ``make cov``. The coverage percentage will be printed on the console when the testing completes.\n\n```sh\nmake cov\n```\n\nFor more detailed analysis, results are printed to ``coverage.txt`` in the project directory.\n\nTo read the coverage report, run:\n\n```sh\ngo tool cover -html=coverage.txt\n```\n\n## Submitting Pull Requests\n\nYou may use your preferred editor to edit the driver code. Make certain to run ``make fmt lint`` before submitting any pull request to Snowflake. This command formats your source code according to the standard Go style and detects any coding style issues.\n"
},
{
"path": "SECURITY.md",
"content": "# Security Policy\n\nPlease refer to the Snowflake [HackerOne program](https://hackerone.com/snowflake?type=team) for our security policies and for reporting any security vulnerabilities.\n\nFor other security related questions and concerns, please contact the Snowflake security team at security@snowflake.com\n"
},
{
"path": "aaa_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"testing\"\n)\n\nfunc TestShowServerVersion(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\trows := dbt.mustQuery(\"SELECT CURRENT_VERSION()\")\n\t\tdefer func() {\n\t\t\tassertNilF(t, rows.Close())\n\t\t}()\n\n\t\tvar version string\n\t\trows.Next()\n\t\tassertNilF(t, rows.Scan(&version))\n\t\tprintln(version)\n\t})\n}\n"
},
{
"path": "arrow_chunk.go",
"content": "package gosnowflake\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/base64\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/query\"\n\t\"time\"\n\n\t\"github.com/apache/arrow-go/v18/arrow\"\n\t\"github.com/apache/arrow-go/v18/arrow/ipc\"\n\t\"github.com/apache/arrow-go/v18/arrow/memory\"\n)\n\ntype arrowResultChunk struct {\n\treader *ipc.Reader\n\trowCount int\n\tloc *time.Location\n\tallocator memory.Allocator\n}\n\nfunc (arc *arrowResultChunk) decodeArrowChunk(ctx context.Context, rowType []query.ExecResponseRowType, highPrec bool, params *syncParams) ([]chunkRowType, error) {\n\tdefer arc.reader.Release()\n\tlogger.Debug(\"Arrow Decoder\")\n\tvar chunkRows []chunkRowType\n\n\tfor arc.reader.Next() {\n\t\trecord := arc.reader.Record()\n\n\t\tstart := len(chunkRows)\n\t\tnumRows := int(record.NumRows())\n\t\tlogger.Debugf(\"rows in current record: %v\", numRows)\n\t\tcolumns := record.Columns()\n\t\tchunkRows = append(chunkRows, make([]chunkRowType, numRows)...)\n\t\tfor i := start; i < start+numRows; i++ {\n\t\t\tchunkRows[i].ArrowRow = make([]snowflakeValue, len(columns))\n\t\t}\n\n\t\tfor colIdx, col := range columns {\n\t\t\tvalues := make([]snowflakeValue, numRows)\n\t\t\tif err := arrowToValues(ctx, values, rowType[colIdx], col, arc.loc, highPrec, params); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tfor i := range values {\n\t\t\t\tchunkRows[start+i].ArrowRow[colIdx] = values[i]\n\t\t\t}\n\t\t}\n\t\tarc.rowCount += numRows\n\t}\n\tlogger.Debugf(\"The number of chunk rows: %v\", len(chunkRows))\n\n\treturn chunkRows, arc.reader.Err()\n}\n\n// decodeArrowBatchRaw reads raw (untransformed) arrow records from the IPC reader.\n// The records are not transformed with arrow-compute; the arrowbatches sub-package\n// handles transformation when the user calls ArrowBatch.Fetch().\nfunc (arc *arrowResultChunk) decodeArrowBatchRaw() (*[]arrow.Record, error) {\n\tvar records []arrow.Record\n\tdefer arc.reader.Release()\n\n\tfor arc.reader.Next() {\n\t\trecord := arc.reader.Record()\n\t\trecord.Retain()\n\t\trecords = append(records, record)\n\t}\n\n\treturn &records, arc.reader.Err()\n}\n\n// Build arrow chunk based on RowSet of base64\nfunc buildFirstArrowChunk(rowsetBase64 string, loc *time.Location, alloc memory.Allocator) (arrowResultChunk, error) {\n\trowSetBytes, err := base64.StdEncoding.DecodeString(rowsetBase64)\n\tif err != nil {\n\t\treturn arrowResultChunk{}, err\n\t}\n\trr, err := ipc.NewReader(bytes.NewReader(rowSetBytes), ipc.WithAllocator(alloc))\n\tif err != nil {\n\t\treturn arrowResultChunk{}, err\n\t}\n\n\treturn arrowResultChunk{rr, 0, loc, alloc}, nil\n}\n"
},
{
"path": "arrow_stream.go",
"content": "package gosnowflake\n\nimport (\n\t\"bufio\"\n\t\"bytes\"\n\t\"compress/gzip\"\n\t\"context\"\n\t\"encoding/base64\"\n\t\"fmt\"\n\t\"io\"\n\t\"maps\"\n\t\"net/http\"\n\t\"strconv\"\n\t\"time\"\n\n\t\"github.com/apache/arrow-go/v18/arrow/ipc\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/query\"\n)\n\n// ArrowStreamLoader is a convenience interface for downloading\n// Snowflake results via multiple Arrow Record Batch streams.\n//\n// Some queries from Snowflake do not return Arrow data regardless\n// of the settings, such as \"SHOW WAREHOUSES\". In these cases,\n// you'll find TotalRows() > 0 but GetBatches returns no batches\n// and no errors. In this case, the data is accessible via JSONData\n// with the actual types matching up to the metadata in RowTypes.\ntype ArrowStreamLoader interface {\n\tGetBatches() ([]ArrowStreamBatch, error)\n\tNextResultSet(ctx context.Context) error\n\tTotalRows() int64\n\tRowTypes() []query.ExecResponseRowType\n\tLocation() *time.Location\n\tJSONData() [][]*string\n}\n\n// ArrowStreamBatch is a type describing a potentially yet-to-be-downloaded\n// Arrow IPC stream. Call GetStream to download and retrieve an io.Reader\n// that can be used with ipc.NewReader to get record batch results.\ntype ArrowStreamBatch struct {\n\tidx int\n\tnumrows int64\n\tscd *snowflakeArrowStreamChunkDownloader\n\tLoc *time.Location\n\trr io.ReadCloser\n}\n\n// NumRows returns the total number of rows that the metadata stated should\n// be in this stream of record batches.\nfunc (asb *ArrowStreamBatch) NumRows() int64 { return asb.numrows }\n\n// GetStream returns a stream of bytes consisting of an Arrow IPC Record\n// batch stream. Close should be called on the returned stream when done\n// to ensure no leaked memory.\nfunc (asb *ArrowStreamBatch) GetStream(ctx context.Context) (io.ReadCloser, error) {\n\tif asb.rr == nil {\n\t\tif err := asb.downloadChunkStreamHelper(ctx); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\treturn asb.rr, nil\n}\n\n// streamWrapReader wraps an io.Reader so that Close closes the underlying body.\ntype streamWrapReader struct {\n\tio.Reader\n\twrapped io.ReadCloser\n}\n\nfunc (w *streamWrapReader) Close() error {\n\tif cl, ok := w.Reader.(io.ReadCloser); ok {\n\t\tif err := cl.Close(); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\treturn w.wrapped.Close()\n}\n\nfunc (asb *ArrowStreamBatch) downloadChunkStreamHelper(ctx context.Context) error {\n\theaders := make(map[string]string)\n\tif len(asb.scd.ChunkHeader) > 0 {\n\t\tmaps.Copy(headers, asb.scd.ChunkHeader)\n\t} else {\n\t\theaders[headerSseCAlgorithm] = headerSseCAes\n\t\theaders[headerSseCKey] = asb.scd.Qrmk\n\t}\n\n\tresp, err := asb.scd.FuncGet(ctx, asb.scd.sc, asb.scd.ChunkMetas[asb.idx].URL, headers, asb.scd.sc.rest.RequestTimeout)\n\tif err != nil {\n\t\treturn err\n\t}\n\tif resp.StatusCode != http.StatusOK {\n\t\tdefer func() {\n\t\t\t_ = resp.Body.Close()\n\t\t}()\n\t\tb, err := io.ReadAll(resp.Body)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\t_ = b\n\t\treturn &SnowflakeError{\n\t\t\tNumber: ErrFailedToGetChunk,\n\t\t\tSQLState: SQLStateConnectionFailure,\n\t\t\tMessage: fmt.Sprintf(\"failed to get chunk. idx: %v\", asb.idx),\n\t\t\tMessageArgs: []any{asb.idx},\n\t\t}\n\t}\n\n\tdefer func() {\n\t\tif asb.rr == nil {\n\t\t\t_ = resp.Body.Close()\n\t\t}\n\t}()\n\n\tbufStream := bufio.NewReader(resp.Body)\n\tgzipMagic, err := bufStream.Peek(2)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tif gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b {\n\t\tbufStream0, err := gzip.NewReader(bufStream)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tasb.rr = &streamWrapReader{Reader: bufStream0, wrapped: resp.Body}\n\t} else {\n\t\tasb.rr = &streamWrapReader{Reader: bufStream, wrapped: resp.Body}\n\t}\n\treturn nil\n}\n\ntype snowflakeArrowStreamChunkDownloader struct {\n\tsc *snowflakeConn\n\tChunkMetas []query.ExecResponseChunk\n\tTotal int64\n\tQrmk string\n\tChunkHeader map[string]string\n\tFuncGet func(context.Context, *snowflakeConn, string, map[string]string, time.Duration) (*http.Response, error)\n\tRowSet rowSetType\n\tresultIDs []string\n}\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) Location() *time.Location {\n\tif scd.sc != nil {\n\t\treturn getCurrentLocation(&scd.sc.syncParams)\n\t}\n\treturn nil\n}\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) TotalRows() int64 { return scd.Total }\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) RowTypes() []query.ExecResponseRowType {\n\treturn scd.RowSet.RowType\n}\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) JSONData() [][]*string {\n\treturn scd.RowSet.JSON\n}\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) maybeFirstBatch() ([]byte, error) {\n\tif scd.RowSet.RowSetBase64 == \"\" {\n\t\treturn nil, nil\n\t}\n\n\trowSetBytes, err := base64.StdEncoding.DecodeString(scd.RowSet.RowSetBase64)\n\tif err != nil {\n\t\tlogger.Warnf(\"skipping first batch as it is not a valid base64 response. %v\", err)\n\t\treturn nil, err\n\t}\n\n\trr, err := ipc.NewReader(bytes.NewReader(rowSetBytes))\n\tif err != nil {\n\t\tlogger.Warnf(\"skipping first batch as it is not a valid IPC stream. %v\", err)\n\t\treturn nil, err\n\t}\n\trr.Release()\n\n\treturn rowSetBytes, nil\n}\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) GetBatches() (out []ArrowStreamBatch, err error) {\n\tchunkMetaLen := len(scd.ChunkMetas)\n\tloc := scd.Location()\n\n\tout = make([]ArrowStreamBatch, chunkMetaLen, chunkMetaLen+1)\n\ttoFill := out\n\trowSetBytes, err := scd.maybeFirstBatch()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif len(rowSetBytes) > 0 {\n\t\tout = out[:chunkMetaLen+1]\n\t\tout[0] = ArrowStreamBatch{\n\t\t\tscd: scd,\n\t\t\tLoc: loc,\n\t\t\trr: io.NopCloser(bytes.NewReader(rowSetBytes)),\n\t\t}\n\t\ttoFill = out[1:]\n\t}\n\n\tvar totalCounted int64\n\tfor i := range toFill {\n\t\ttoFill[i] = ArrowStreamBatch{\n\t\t\tidx: i,\n\t\t\tnumrows: int64(scd.ChunkMetas[i].RowCount),\n\t\t\tLoc: loc,\n\t\t\tscd: scd,\n\t\t}\n\t\ttotalCounted += int64(scd.ChunkMetas[i].RowCount)\n\t}\n\n\tif len(rowSetBytes) > 0 {\n\t\tout[0].numrows = scd.Total - totalCounted\n\t}\n\treturn\n}\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) NextResultSet(ctx context.Context) error {\n\tif !scd.hasNextResultSet() {\n\t\treturn io.EOF\n\t}\n\tresultID := scd.resultIDs[0]\n\tscd.resultIDs = scd.resultIDs[1:]\n\tresultPath := fmt.Sprintf(urlQueriesResultFmt, resultID)\n\tresp, err := scd.sc.getQueryResultResp(ctx, resultPath)\n\tif err != nil {\n\t\treturn err\n\t}\n\tif !resp.Success {\n\t\tcode, err := strconv.Atoi(resp.Code)\n\t\tif err != nil {\n\t\t\tlogger.WithContext(ctx).Errorf(\"error while parsing code: %v\", err)\n\t\t}\n\t\treturn exceptionTelemetry(&SnowflakeError{\n\t\t\tNumber: code,\n\t\t\tSQLState: resp.Data.SQLState,\n\t\t\tMessage: resp.Message,\n\t\t\tQueryID: resp.Data.QueryID,\n\t\t}, scd.sc)\n\t}\n\tscd.ChunkMetas = resp.Data.Chunks\n\tscd.Total = resp.Data.Total\n\tscd.Qrmk = resp.Data.Qrmk\n\tscd.ChunkHeader = resp.Data.ChunkHeaders\n\tscd.RowSet = rowSetType{\n\t\tRowType: resp.Data.RowType,\n\t\tJSON: resp.Data.RowSet,\n\t\tRowSetBase64: resp.Data.RowSetBase64,\n\t}\n\treturn nil\n}\n\nfunc (scd *snowflakeArrowStreamChunkDownloader) hasNextResultSet() bool {\n\treturn len(scd.resultIDs) > 0\n}\n"
},
{
"path": "arrow_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"fmt\"\n\t\"math/big\"\n\t\"reflect\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/apache/arrow-go/v18/arrow/memory\"\n\tia \"github.com/snowflakedb/gosnowflake/v2/internal/arrow\"\n\n\t\"database/sql/driver\"\n)\n\nfunc TestArrowBatchDataProvider(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tctx := ia.EnableArrowBatches(context.Background())\n\t\tquery := \"select '0.1':: DECIMAL(38, 19) as c\"\n\n\t\tvar rows driver.Rows\n\t\tvar err error\n\n\t\terr = dbt.conn.Raw(func(x any) error {\n\t\t\tqueryer, implementsQueryContext := x.(driver.QueryerContext)\n\t\t\tassertTrueF(t, implementsQueryContext, \"snowflake connection driver does not implement queryerContext\")\n\n\t\t\trows, err = queryer.QueryContext(ctx, query, nil)\n\t\t\treturn err\n\t\t})\n\n\t\tassertNilF(t, err, \"error running select query\")\n\n\t\tsfRows, isSfRows := rows.(SnowflakeRows)\n\t\tassertTrueF(t, isSfRows, \"rows should be snowflakeRows\")\n\n\t\tprovider, isProvider := sfRows.(ia.BatchDataProvider)\n\t\tassertTrueF(t, isProvider, \"rows should implement BatchDataProvider\")\n\n\t\tinfo, err := provider.GetArrowBatches()\n\t\tassertNilF(t, err, \"error getting arrow batch data\")\n\t\tassertNotEqualF(t, len(info.Batches), 0, \"should have at least one batch\")\n\n\t\t// Verify raw records are available for the first batch\n\t\tbatch := info.Batches[0]\n\t\tassertNotNilF(t, batch.Records, \"first batch should have pre-decoded records\")\n\n\t\trecords := *batch.Records\n\t\tassertNotEqualF(t, len(records), 0, \"should have at least one record\")\n\n\t\t// Verify column 0 has data (raw decimal value)\n\t\tstrVal := records[0].Column(0).ValueStr(0)\n\t\tassertTrueF(t, len(strVal) > 0, fmt.Sprintf(\"column should have a value, got: %s\", strVal))\n\t})\n}\n\nfunc TestArrowBigInt(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\ttestcases := []struct {\n\t\t\tnum string\n\t\t\tprec int\n\t\t\tsc int\n\t\t}{\n\t\t\t{\"10000000000000000000000000000000000000\", 38, 0},\n\t\t\t{\"-10000000000000000000000000000000000000\", 38, 0},\n\t\t\t{\"12345678901234567890123456789012345678\", 38, 0}, // #pragma: allowlist secret\n\t\t\t{\"-12345678901234567890123456789012345678\", 38, 0},\n\t\t\t{\"99999999999999999999999999999999999999\", 38, 0},\n\t\t\t{\"-99999999999999999999999999999999999999\", 38, 0},\n\t\t}\n\n\t\tfor _, tc := range testcases {\n\t\t\trows := dbt.mustQueryContext(WithHigherPrecision(context.Background()),\n\t\t\t\tfmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\tif !rows.Next() {\n\t\t\t\tdbt.Error(\"failed to query\")\n\t\t\t}\n\t\t\tdefer rows.Close()\n\t\t\tvar v *big.Int\n\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\tdbt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t}\n\n\t\t\tb, ok := new(big.Int).SetString(tc.num, 10)\n\t\t\tif !ok {\n\t\t\t\tdbt.Errorf(\"failed to convert %v big.Int.\", tc.num)\n\t\t\t}\n\t\t\tif v.Cmp(b) != 0 {\n\t\t\t\tdbt.Errorf(\"big.Int value mismatch: expected %v, got %v\", b, v)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestArrowBigFloat(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\ttestcases := []struct {\n\t\t\tnum string\n\t\t\tprec int\n\t\t\tsc int\n\t\t}{\n\t\t\t{\"1.23\", 30, 2},\n\t\t\t{\"1.0000000000000000000000000000000000000\", 38, 37},\n\t\t\t{\"-1.0000000000000000000000000000000000000\", 38, 37},\n\t\t\t{\"1.2345678901234567890123456789012345678\", 38, 37},\n\t\t\t{\"-1.2345678901234567890123456789012345678\", 38, 37},\n\t\t\t{\"9.9999999999999999999999999999999999999\", 38, 37},\n\t\t\t{\"-9.9999999999999999999999999999999999999\", 38, 37},\n\t\t}\n\n\t\tfor _, tc := range testcases {\n\t\t\trows := dbt.mustQueryContext(WithHigherPrecision(context.Background()),\n\t\t\t\tfmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\tif !rows.Next() {\n\t\t\t\tdbt.Error(\"failed to query\")\n\t\t\t}\n\t\t\tdefer rows.Close()\n\t\t\tvar v *big.Float\n\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\tdbt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t}\n\n\t\t\tprec := v.Prec()\n\t\t\tb, ok := new(big.Float).SetPrec(prec).SetString(tc.num)\n\t\t\tif !ok {\n\t\t\t\tdbt.Errorf(\"failed to convert %v to big.Float.\", tc.num)\n\t\t\t}\n\t\t\tif v.Cmp(b) != 0 {\n\t\t\t\tdbt.Errorf(\"big.Float value mismatch: expected %v, got %v\", b, v)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestArrowIntPrecision(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tdbt.mustExec(forceJSON)\n\n\t\tintTestcases := []struct {\n\t\t\tnum string\n\t\t\tprec int\n\t\t\tsc int\n\t\t}{\n\t\t\t{\"10000000000000000000000000000000000000\", 38, 0},\n\t\t\t{\"-10000000000000000000000000000000000000\", 38, 0},\n\t\t\t{\"12345678901234567890123456789012345678\", 38, 0}, // pragma: allowlist secret\n\t\t\t{\"-12345678901234567890123456789012345678\", 38, 0},\n\t\t\t{\"99999999999999999999999999999999999999\", 38, 0},\n\t\t\t{\"-99999999999999999999999999999999999999\", 38, 0},\n\t\t}\n\n\t\tt.Run(\"arrow_disabled_scan_int64\", func(t *testing.T) {\n\t\t\tfor _, tc := range intTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v int64\n\t\t\t\tif err := rows.Scan(&v); err == nil {\n\t\t\t\t\tt.Error(\"should fail to scan\")\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t\tt.Run(\"arrow_disabled_scan_string\", func(t *testing.T) {\n\t\t\tfor _, tc := range intTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v string\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t\tif v != tc.num {\n\t\t\t\t\tt.Errorf(\"string value mismatch: expected %v, got %v\", tc.num, v)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\n\t\tdbt.mustExec(forceARROW)\n\n\t\tt.Run(\"arrow_enabled_scan_big_int\", func(t *testing.T) {\n\t\t\tfor _, tc := range intTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v string\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t\tif !strings.EqualFold(v, tc.num) {\n\t\t\t\t\tt.Errorf(\"int value mismatch: expected %v, got %v\", tc.num, v)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t\tt.Run(\"arrow_high_precision_enabled_scan_big_int\", func(t *testing.T) {\n\t\t\tfor _, tc := range intTestcases {\n\t\t\t\trows := dbt.mustQueryContext(WithHigherPrecision(context.Background()), fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v *big.Int\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\n\t\t\t\tb, ok := new(big.Int).SetString(tc.num, 10)\n\t\t\t\tif !ok {\n\t\t\t\t\tt.Errorf(\"failed to convert %v big.Int.\", tc.num)\n\t\t\t\t}\n\t\t\t\tif v.Cmp(b) != 0 {\n\t\t\t\t\tt.Errorf(\"big.Int value mismatch: expected %v, got %v\", b, v)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t})\n}\n\n// TestArrowFloatPrecision tests the different variable types allowed in the\n// rows.Scan() method. Note that for lower precision types we do not attempt\n// to check the value as precision could be lost.\nfunc TestArrowFloatPrecision(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tdbt.mustExec(forceJSON)\n\n\t\tfltTestcases := []struct {\n\t\t\tnum string\n\t\t\tprec int\n\t\t\tsc int\n\t\t}{\n\t\t\t{\"1.23\", 30, 2},\n\t\t\t{\"1.0000000000000000000000000000000000000\", 38, 37},\n\t\t\t{\"-1.0000000000000000000000000000000000000\", 38, 37},\n\t\t\t{\"1.2345678901234567890123456789012345678\", 38, 37},\n\t\t\t{\"-1.2345678901234567890123456789012345678\", 38, 37},\n\t\t\t{\"9.9999999999999999999999999999999999999\", 38, 37},\n\t\t\t{\"-9.9999999999999999999999999999999999999\", 38, 37},\n\t\t}\n\n\t\tt.Run(\"arrow_disabled_scan_float64\", func(t *testing.T) {\n\t\t\tfor _, tc := range fltTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v float64\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t\tt.Run(\"arrow_disabled_scan_float32\", func(t *testing.T) {\n\t\t\tfor _, tc := range fltTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v float32\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t\tt.Run(\"arrow_disabled_scan_string\", func(t *testing.T) {\n\t\t\tfor _, tc := range fltTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v string\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t\tif !strings.EqualFold(v, tc.num) {\n\t\t\t\t\tt.Errorf(\"int value mismatch: expected %v, got %v\", tc.num, v)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\n\t\tdbt.mustExec(forceARROW)\n\n\t\tt.Run(\"arrow_enabled_scan_float64\", func(t *testing.T) {\n\t\t\tfor _, tc := range fltTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v float64\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t\tt.Run(\"arrow_enabled_scan_float32\", func(t *testing.T) {\n\t\t\tfor _, tc := range fltTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v float32\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t\tt.Run(\"arrow_enabled_scan_string\", func(t *testing.T) {\n\t\t\tfor _, tc := range fltTestcases {\n\t\t\t\trows := dbt.mustQuery(fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v string\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\t\t\t\tif v != tc.num {\n\t\t\t\t\tt.Errorf(\"string value mismatch: expected %v, got %v\", tc.num, v)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t\tt.Run(\"arrow_high_precision_enabled_scan_big_float\", func(t *testing.T) {\n\t\t\tfor _, tc := range fltTestcases {\n\t\t\t\trows := dbt.mustQueryContext(WithHigherPrecision(context.Background()), fmt.Sprintf(selectNumberSQL, tc.num, tc.prec, tc.sc))\n\t\t\t\tdefer rows.Close()\n\t\t\t\tif !rows.Next() {\n\t\t\t\t\tt.Error(\"failed to query\")\n\t\t\t\t}\n\t\t\t\tvar v *big.Float\n\t\t\t\tif err := rows.Scan(&v); err != nil {\n\t\t\t\t\tt.Errorf(\"failed to scan. %#v\", err)\n\t\t\t\t}\n\n\t\t\t\tprec := v.Prec()\n\t\t\t\tb, ok := new(big.Float).SetPrec(prec).SetString(tc.num)\n\t\t\t\tif !ok {\n\t\t\t\t\tt.Errorf(\"failed to convert %v to big.Float.\", tc.num)\n\t\t\t\t}\n\t\t\t\tif v.Cmp(b) != 0 {\n\t\t\t\t\tt.Errorf(\"big.Float value mismatch: expected %v, got %v\", b, v)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t})\n}\n\nfunc TestArrowTimePrecision(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tdbt.mustExec(\"CREATE TABLE t (col5 TIME(5), col6 TIME(6), col7 TIME(7), col8 TIME(8));\")\n\t\tdefer dbt.mustExec(\"DROP TABLE IF EXISTS t\")\n\t\tdbt.mustExec(\"INSERT INTO t VALUES ('23:59:59.99999', '23:59:59.999999', '23:59:59.9999999', '23:59:59.99999999');\")\n\n\t\trows := dbt.mustQuery(\"select * from t\")\n\t\tdefer rows.Close()\n\t\tvar c5, c6, c7, c8 time.Time\n\t\tfor rows.Next() {\n\t\t\tif err := rows.Scan(&c5, &c6, &c7, &c8); err != nil {\n\t\t\t\tt.Errorf(\"values were not scanned: %v\", err)\n\t\t\t}\n\t\t}\n\n\t\tnano := 999999990\n\t\texpected := time.Time{}.Add(23*time.Hour + 59*time.Minute + 59*time.Second + 99*time.Millisecond)\n\t\tif c8.Unix() != expected.Unix() || c8.Nanosecond() != nano {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c8)\n\t\t}\n\t\tif c7.Unix() != expected.Unix() || c7.Nanosecond() != nano-(nano%1e2) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c7)\n\t\t}\n\t\tif c6.Unix() != expected.Unix() || c6.Nanosecond() != nano-(nano%1e3) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c6)\n\t\t}\n\t\tif c5.Unix() != expected.Unix() || c5.Nanosecond() != nano-(nano%1e4) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c5)\n\t\t}\n\n\t\tdbt.mustExec(`CREATE TABLE t_ntz (\n\t\t col1 TIMESTAMP_NTZ(1),\n\t\t col2 TIMESTAMP_NTZ(2),\n\t\t col3 TIMESTAMP_NTZ(3),\n\t\t col4 TIMESTAMP_NTZ(4),\n\t\t col5 TIMESTAMP_NTZ(5),\n\t\t col6 TIMESTAMP_NTZ(6),\n\t\t col7 TIMESTAMP_NTZ(7),\n\t\t col8 TIMESTAMP_NTZ(8)\n\t\t);`)\n\t\tdefer dbt.mustExec(\"DROP TABLE IF EXISTS t_ntz\")\n\t\tdbt.mustExec(`INSERT INTO t_ntz VALUES (\n\t\t '9999-12-31T23:59:59.9',\n\t\t '9999-12-31T23:59:59.99',\n\t\t '9999-12-31T23:59:59.999',\n\t\t '9999-12-31T23:59:59.9999',\n\t\t '9999-12-31T23:59:59.99999',\n\t\t '9999-12-31T23:59:59.999999',\n\t\t '9999-12-31T23:59:59.9999999',\n\t\t '9999-12-31T23:59:59.99999999'\n\t\t);`)\n\n\t\trows2 := dbt.mustQuery(\"select * from t_ntz\")\n\t\tdefer rows2.Close()\n\t\tvar c1, c2, c3, c4 time.Time\n\t\tfor rows2.Next() {\n\t\t\tif err := rows2.Scan(&c1, &c2, &c3, &c4, &c5, &c6, &c7, &c8); err != nil {\n\t\t\t\tt.Errorf(\"values were not scanned: %v\", err)\n\t\t\t}\n\t\t}\n\n\t\texpected = time.Date(9999, 12, 31, 23, 59, 59, 0, time.UTC)\n\t\tif c8.Unix() != expected.Unix() || c8.Nanosecond() != nano {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c8)\n\t\t}\n\t\tif c7.Unix() != expected.Unix() || c7.Nanosecond() != nano-(nano%1e2) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c7)\n\t\t}\n\t\tif c6.Unix() != expected.Unix() || c6.Nanosecond() != nano-(nano%1e3) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c6)\n\t\t}\n\t\tif c5.Unix() != expected.Unix() || c5.Nanosecond() != nano-(nano%1e4) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c5)\n\t\t}\n\t\tif c4.Unix() != expected.Unix() || c4.Nanosecond() != nano-(nano%1e5) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c4)\n\t\t}\n\t\tif c3.Unix() != expected.Unix() || c3.Nanosecond() != nano-(nano%1e6) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c3)\n\t\t}\n\t\tif c2.Unix() != expected.Unix() || c2.Nanosecond() != nano-(nano%1e7) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c2)\n\t\t}\n\t\tif c1.Unix() != expected.Unix() || c1.Nanosecond() != nano-(nano%1e8) {\n\t\t\tt.Errorf(\"the value did not match. expected: %v, got: %v\", expected, c1)\n\t\t}\n\t})\n}\n\nfunc TestArrowVariousTypes(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\trows := dbt.mustQueryContext(\n\t\t\tWithHigherPrecision(context.Background()), selectVariousTypes)\n\t\tdefer rows.Close()\n\t\tif !rows.Next() {\n\t\t\tdbt.Error(\"failed to query\")\n\t\t}\n\t\tcc, err := rows.Columns()\n\t\tif err != nil {\n\t\t\tdbt.Errorf(\"columns: %v\", cc)\n\t\t}\n\t\tct, err := rows.ColumnTypes()\n\t\tif err != nil {\n\t\t\tdbt.Errorf(\"column types: %v\", ct)\n\t\t}\n\t\tvar v1 *big.Float\n\t\tvar v2, v2a int\n\t\tvar v3 string\n\t\tvar v4 float64\n\t\tvar v5 []byte\n\t\tvar v6 bool\n\t\tif err = rows.Scan(&v1, &v2, &v2a, &v3, &v4, &v5, &v6); err != nil {\n\t\t\tdbt.Errorf(\"failed to scan: %#v\", err)\n\t\t}\n\t\tif v1.Cmp(big.NewFloat(1.0)) != 0 {\n\t\t\tdbt.Errorf(\"failed to scan. %#v\", *v1)\n\t\t}\n\t\tif ct[0].Name() != \"C1\" || ct[1].Name() != \"C2\" || ct[2].Name() != \"C2A\" || ct[3].Name() != \"C3\" || ct[4].Name() != \"C4\" || ct[5].Name() != \"C5\" || ct[6].Name() != \"C6\" {\n\t\t\tdbt.Errorf(\"failed to get column names: %#v\", ct)\n\t\t}\n\t\tif ct[0].ScanType() != reflect.TypeFor[*big.Float]() {\n\t\t\tdbt.Errorf(\"failed to get scan type. expected: %v, got: %v\", reflect.TypeFor[float64](), ct[0].ScanType())\n\t\t}\n\t\tif ct[1].ScanType() != reflect.TypeFor[int64]() {\n\t\t\tdbt.Errorf(\"failed to get scan type. expected: %v, got: %v\", reflect.TypeFor[int64](), ct[1].ScanType())\n\t\t}\n\t\tif ct[2].ScanType() != reflect.TypeFor[*big.Int]() {\n\t\t\tdbt.Errorf(\"failed to get scan type. expected: %v, got: %v\", reflect.TypeFor[*big.Int](), ct[2].ScanType())\n\t\t}\n\t\tvar pr, sc int64\n\t\tvar cLen int64\n\t\tpr, sc = dbt.mustDecimalSize(ct[0])\n\t\tif pr != 30 || sc != 2 {\n\t\t\tdbt.Errorf(\"failed to get precision and scale. %#v\", ct[0])\n\t\t}\n\t\tdbt.mustFailLength(ct[0])\n\t\tif canNull := dbt.mustNullable(ct[0]); canNull {\n\t\t\tdbt.Errorf(\"failed to get nullable. %#v\", ct[0])\n\t\t}\n\t\tif cLen != 0 {\n\t\t\tdbt.Errorf(\"failed to get length. %#v\", ct[0])\n\t\t}\n\t\tif v2 != 2 {\n\t\t\tdbt.Errorf(\"failed to scan. %#v\", v2)\n\t\t}\n\t\tpr, sc = dbt.mustDecimalSize(ct[1])\n\t\tif pr != 18 || sc != 0 {\n\t\t\tdbt.Errorf(\"failed to get precision and scale. %#v\", ct[1])\n\t\t}\n\t\tdbt.mustFailLength(ct[1])\n\t\tif canNull := dbt.mustNullable(ct[1]); canNull {\n\t\t\tdbt.Errorf(\"failed to get nullable. %#v\", ct[1])\n\t\t}\n\t\tif v2a != 22 {\n\t\t\tdbt.Errorf(\"failed to scan. %#v\", v2a)\n\t\t}\n\t\tdbt.mustFailLength(ct[2])\n\t\tif canNull := dbt.mustNullable(ct[2]); canNull {\n\t\t\tdbt.Errorf(\"failed to get nullable. %#v\", ct[2])\n\t\t}\n\t\tif v3 != \"t3\" {\n\t\t\tdbt.Errorf(\"failed to scan. %#v\", v3)\n\t\t}\n\t\tdbt.mustFailDecimalSize(ct[3])\n\t\tif cLen = dbt.mustLength(ct[3]); cLen != 2 {\n\t\t\tdbt.Errorf(\"failed to get length. %#v\", ct[3])\n\t\t}\n\t\tif canNull := dbt.mustNullable(ct[3]); canNull {\n\t\t\tdbt.Errorf(\"failed to get nullable. %#v\", ct[3])\n\t\t}\n\t\tif v4 != 4.2 {\n\t\t\tdbt.Errorf(\"failed to scan. %#v\", v4)\n\t\t}\n\t\tdbt.mustFailDecimalSize(ct[4])\n\t\tdbt.mustFailLength(ct[4])\n\t\tif canNull := dbt.mustNullable(ct[4]); canNull {\n\t\t\tdbt.Errorf(\"failed to get nullable. %#v\", ct[4])\n\t\t}\n\t\tif !bytes.Equal(v5, []byte{0xab, 0xcd}) {\n\t\t\tdbt.Errorf(\"failed to scan. %#v\", v5)\n\t\t}\n\t\tdbt.mustFailDecimalSize(ct[5])\n\t\tif cLen = dbt.mustLength(ct[5]); cLen != 8388608 { // BINARY\n\t\t\tdbt.Errorf(\"failed to get length. %#v\", ct[5])\n\t\t}\n\t\tif canNull := dbt.mustNullable(ct[5]); canNull {\n\t\t\tdbt.Errorf(\"failed to get nullable. %#v\", ct[5])\n\t\t}\n\t\tif !v6 {\n\t\t\tdbt.Errorf(\"failed to scan. %#v\", v6)\n\t\t}\n\t\tdbt.mustFailDecimalSize(ct[6])\n\t\tdbt.mustFailLength(ct[6])\n\t})\n}\n\nfunc TestArrowMemoryCleanedUp(t *testing.T) {\n\tmem := memory.NewCheckedAllocator(memory.NewGoAllocator())\n\tdefer mem.AssertSize(t, 0)\n\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tctx := WithArrowAllocator(\n\t\t\tcontext.Background(),\n\t\t\tmem,\n\t\t)\n\n\t\trows := dbt.mustQueryContext(ctx, \"select 1 UNION select 2 ORDER BY 1\")\n\t\tdefer rows.Close()\n\t\tvar v int\n\t\tassertTrueF(t, rows.Next())\n\t\tassertNilF(t, rows.Scan(&v))\n\t\tassertEqualE(t, v, 1)\n\t\tassertTrueF(t, rows.Next())\n\t\tassertNilF(t, rows.Scan(&v))\n\t\tassertEqualE(t, v, 2)\n\t\tassertFalseE(t, rows.Next())\n\t})\n}\n"
},
{
"path": "arrowbatches/batches.go",
"content": "package arrowbatches\n\nimport (\n\t\"cmp\"\n\t\"context\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/query\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/types\"\n\t\"time\"\n\n\tsf \"github.com/snowflakedb/gosnowflake/v2\"\n\tia \"github.com/snowflakedb/gosnowflake/v2/internal/arrow\"\n\n\t\"github.com/apache/arrow-go/v18/arrow\"\n\t\"github.com/apache/arrow-go/v18/arrow/memory\"\n)\n\n// ArrowBatch represents a chunk of data retrievable in arrow.Record format.\ntype ArrowBatch struct {\n\traw ia.BatchRaw\n\trowTypes []query.ExecResponseRowType\n\tallocator memory.Allocator\n\tctx context.Context\n}\n\n// WithContext sets the context for subsequent Fetch calls on this batch.\nfunc (rb *ArrowBatch) WithContext(ctx context.Context) *ArrowBatch {\n\trb.ctx = ctx\n\treturn rb\n}\n\n// Fetch returns an array of arrow.Record representing this batch's data.\n// Records are transformed from Snowflake's internal format to standard Arrow types.\nfunc (rb *ArrowBatch) Fetch() (*[]arrow.Record, error) {\n\tvar rawRecords *[]arrow.Record\n\tctx := cmp.Or(rb.ctx, context.Background())\n\n\tif rb.raw.Records != nil {\n\t\trawRecords = rb.raw.Records\n\t} else if rb.raw.Download != nil {\n\t\trecs, rowCount, err := rb.raw.Download(ctx)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\trawRecords = recs\n\t\trb.raw.Records = recs\n\t\trb.raw.RowCount = rowCount\n\t}\n\n\tif rawRecords == nil || len(*rawRecords) == 0 {\n\t\tempty := make([]arrow.Record, 0)\n\t\treturn &empty, nil\n\t}\n\n\tvar transformed []arrow.Record\n\tfor i, rec := range *rawRecords {\n\t\tnewRec, err := arrowToRecord(ctx, rec, rb.allocator, rb.rowTypes, rb.raw.Location)\n\t\tif err != nil {\n\t\t\tfor _, t := range transformed {\n\t\t\t\tt.Release()\n\t\t\t}\n\t\t\tfor _, r := range (*rawRecords)[i:] {\n\t\t\t\tr.Release()\n\t\t\t}\n\t\t\trb.raw.Records = nil\n\t\t\treturn nil, err\n\t\t}\n\t\ttransformed = append(transformed, newRec)\n\t\trec.Release()\n\t}\n\trb.raw.Records = nil\n\trb.raw.RowCount = countArrowBatchRows(&transformed)\n\treturn &transformed, nil\n}\n\n// GetRowCount returns the number of rows in this batch.\nfunc (rb *ArrowBatch) GetRowCount() int {\n\treturn rb.raw.RowCount\n}\n\n// GetLocation returns the timezone location for this batch.\nfunc (rb *ArrowBatch) GetLocation() *time.Location {\n\treturn rb.raw.Location\n}\n\n// GetRowTypes returns the column metadata for this batch.\nfunc (rb *ArrowBatch) GetRowTypes() []query.ExecResponseRowType {\n\treturn rb.rowTypes\n}\n\n// ArrowSnowflakeTimestampToTime converts an original Snowflake timestamp to time.Time.\nfunc (rb *ArrowBatch) ArrowSnowflakeTimestampToTime(rec arrow.Record, colIdx int, recIdx int) *time.Time {\n\tscale := int(rb.rowTypes[colIdx].Scale)\n\tdbType := rb.rowTypes[colIdx].Type\n\treturn ArrowSnowflakeTimestampToTime(rec.Column(colIdx), types.GetSnowflakeType(dbType), scale, recIdx, rb.raw.Location)\n}\n\n// GetArrowBatches retrieves arrow batches from SnowflakeRows.\n// The rows must have been queried with arrowbatches.WithArrowBatches(ctx).\nfunc GetArrowBatches(rows sf.SnowflakeRows) ([]*ArrowBatch, error) {\n\tprovider, ok := rows.(ia.BatchDataProvider)\n\tif !ok {\n\t\treturn nil, &sf.SnowflakeError{\n\t\t\tNumber: sf.ErrNotImplemented,\n\t\t\tMessage: \"rows do not support arrow batch data\",\n\t\t}\n\t}\n\n\tinfo, err := provider.GetArrowBatches()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tbatches := make([]*ArrowBatch, len(info.Batches))\n\tfor i, raw := range info.Batches {\n\t\tbatches[i] = &ArrowBatch{\n\t\t\traw: raw,\n\t\t\trowTypes: info.RowTypes,\n\t\t\tallocator: info.Allocator,\n\t\t\tctx: info.Ctx,\n\t\t}\n\t}\n\treturn batches, nil\n}\n\nfunc countArrowBatchRows(recs *[]arrow.Record) (cnt int) {\n\tfor _, r := range *recs {\n\t\tcnt += int(r.NumRows())\n\t}\n\treturn\n}\n\n// GetAllocator returns the memory allocator for this batch.\nfunc (rb *ArrowBatch) GetAllocator() memory.Allocator {\n\treturn rb.allocator\n}\n"
},
{
"path": "arrowbatches/batches_test.go",
"content": "package arrowbatches\n\nimport (\n\t\"context\"\n\t\"crypto/rsa\"\n\t\"crypto/x509\"\n\t\"database/sql\"\n\t\"database/sql/driver\"\n\t\"encoding/pem\"\n\t\"errors\"\n\t\"fmt\"\n\t\"math\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"strconv\"\n\t\"strings\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/apache/arrow-go/v18/arrow\"\n\t\"github.com/apache/arrow-go/v18/arrow/array\"\n\t\"github.com/apache/arrow-go/v18/arrow/memory\"\n\n\tsf \"github.com/snowflakedb/gosnowflake/v2\"\n\tia \"github.com/snowflakedb/gosnowflake/v2/internal/arrow\"\n)\n\n// testConn holds a reusable database connection for running multiple queries.\ntype testConn struct {\n\tdb *sql.DB\n\tconn *sql.Conn\n}\n\n// repoRoot walks up from the current working directory to find the directory\n// containing go.mod, which is the repository root.\nfunc repoRoot(t *testing.T) string {\n\tt.Helper()\n\tdir, err := os.Getwd()\n\tif err != nil {\n\t\tt.Fatalf(\"failed to get working directory: %v\", err)\n\t}\n\tfor {\n\t\tif _, err = os.Stat(filepath.Join(dir, \"go.mod\")); err == nil {\n\t\t\treturn dir\n\t\t}\n\t\tif !os.IsNotExist(err) {\n\t\t\tt.Fatalf(\"failed to stat go.mod in %q: %v\", dir, err)\n\t\t}\n\t\tparent := filepath.Dir(dir)\n\t\tif parent == dir {\n\t\t\tt.Fatal(\"could not find repository root (no go.mod found)\")\n\t\t}\n\t\tdir = parent\n\t}\n}\n\n// readPrivateKey reads an RSA private key from a PEM file. If the path is\n// relative it is resolved against the repository root so that tests in\n// sub-packages work with repo-root-relative paths.\nfunc readPrivateKey(t *testing.T, path string) *rsa.PrivateKey {\n\tt.Helper()\n\tif !filepath.IsAbs(path) {\n\t\tpath = filepath.Join(repoRoot(t), path)\n\t}\n\tdata, err := os.ReadFile(path)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to read private key file %q: %v\", path, err)\n\t}\n\tblock, _ := pem.Decode(data)\n\tif block == nil {\n\t\tt.Fatalf(\"failed to decode PEM block from %q\", path)\n\t}\n\tkey, err := x509.ParsePKCS8PrivateKey(block.Bytes)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to parse private key from %q: %v\", path, err)\n\t}\n\trsaKey, ok := key.(*rsa.PrivateKey)\n\tif !ok {\n\t\tt.Fatalf(\"private key in %q is not RSA (got %T)\", path, key)\n\t}\n\treturn rsaKey\n}\n\nfunc testConfig(t *testing.T) *sf.Config {\n\tt.Helper()\n\tconfigParams := []*sf.ConfigParam{\n\t\t{Name: \"Account\", EnvName: \"SNOWFLAKE_TEST_ACCOUNT\", FailOnMissing: true},\n\t\t{Name: \"User\", EnvName: \"SNOWFLAKE_TEST_USER\", FailOnMissing: true},\n\t\t{Name: \"Host\", EnvName: \"SNOWFLAKE_TEST_HOST\", FailOnMissing: false},\n\t\t{Name: \"Port\", EnvName: \"SNOWFLAKE_TEST_PORT\", FailOnMissing: false},\n\t\t{Name: \"Protocol\", EnvName: \"SNOWFLAKE_TEST_PROTOCOL\", FailOnMissing: false},\n\t\t{Name: \"Warehouse\", EnvName: \"SNOWFLAKE_TEST_WAREHOUSE\", FailOnMissing: false},\n\t}\n\tisJWT := os.Getenv(\"SNOWFLAKE_TEST_AUTHENTICATOR\") == \"SNOWFLAKE_JWT\"\n\tif !isJWT {\n\t\tconfigParams = append(configParams,\n\t\t\t&sf.ConfigParam{Name: \"Password\", EnvName: \"SNOWFLAKE_TEST_PASSWORD\", FailOnMissing: true},\n\t\t)\n\t}\n\tcfg, err := sf.GetConfigFromEnv(configParams)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to get config from environment: %v\", err)\n\t}\n\tif isJWT {\n\t\tprivKeyPath := os.Getenv(\"SNOWFLAKE_TEST_PRIVATE_KEY\")\n\t\tif privKeyPath == \"\" {\n\t\t\tt.Fatal(\"SNOWFLAKE_TEST_PRIVATE_KEY must be set for JWT authentication\")\n\t\t}\n\t\tcfg.PrivateKey = readPrivateKey(t, privKeyPath)\n\t\tcfg.Authenticator = sf.AuthTypeJwt\n\t}\n\ttz := \"UTC\"\n\tif cfg.Params == nil {\n\t\tcfg.Params = make(map[string]*string)\n\t}\n\tcfg.Params[\"timezone\"] = &tz\n\treturn cfg\n}\n\nfunc openTestConn(ctx context.Context, t *testing.T) *testConn {\n\tt.Helper()\n\tcfg := testConfig(t)\n\tdsn, err := sf.DSN(cfg)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to create DSN: %v\", err)\n\t}\n\tdb, err := sql.Open(\"snowflake\", dsn)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to open db: %v\", err)\n\t}\n\tconn, err := db.Conn(ctx)\n\tif err != nil {\n\t\tdb.Close()\n\t\tt.Fatalf(\"failed to get connection: %v\", err)\n\t}\n\treturn &testConn{db: db, conn: conn}\n}\n\nfunc (tc *testConn) close() {\n\ttc.conn.Close()\n\ttc.db.Close()\n}\n\n// queryRows executes a query on the existing connection and returns\n// SnowflakeRows plus a function to close just the rows.\nfunc (tc *testConn) queryRows(ctx context.Context, t *testing.T, query string) (sf.SnowflakeRows, func()) {\n\tt.Helper()\n\tvar rows driver.Rows\n\tvar err error\n\terr = tc.conn.Raw(func(x any) error {\n\t\tqueryer, ok := x.(driver.QueryerContext)\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"connection does not implement QueryerContext\")\n\t\t}\n\t\trows, err = queryer.QueryContext(ctx, query, nil)\n\t\treturn err\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to execute query: %v\", err)\n\t}\n\tsfRows, ok := rows.(sf.SnowflakeRows)\n\tif !ok {\n\t\trows.Close()\n\t\tt.Fatalf(\"rows do not implement SnowflakeRows\")\n\t}\n\treturn sfRows, func() { rows.Close() }\n}\n\n// queryRawRows is a convenience wrapper that opens a new connection,\n// runs a single query, and returns SnowflakeRows with a full cleanup.\nfunc queryRawRows(ctx context.Context, t *testing.T, query string) (sf.SnowflakeRows, func()) {\n\tt.Helper()\n\ttc := openTestConn(ctx, t)\n\tsfRows, closeRows := tc.queryRows(ctx, t, query)\n\treturn sfRows, func() {\n\t\tcloseRows()\n\t\ttc.close()\n\t}\n}\n\nfunc TestGetArrowBatches(t *testing.T) {\n\tctx := WithArrowBatches(context.Background())\n\n\tsfRows, cleanup := queryRawRows(ctx, t, \"SELECT 1 AS num, 'hello' AS str\")\n\tdefer cleanup()\n\n\tbatches, err := GetArrowBatches(sfRows)\n\tif err != nil {\n\t\tt.Fatalf(\"GetArrowBatches failed: %v\", err)\n\t}\n\tif len(batches) == 0 {\n\t\tt.Fatal(\"expected at least one batch\")\n\t}\n\n\trecords, err := batches[0].Fetch()\n\tif err != nil {\n\t\tt.Fatalf(\"Fetch failed: %v\", err)\n\t}\n\tif records == nil || len(*records) == 0 {\n\t\tt.Fatal(\"expected at least one record\")\n\t}\n\n\trec := (*records)[0]\n\tdefer rec.Release()\n\n\tif rec.NumCols() != 2 {\n\t\tt.Fatalf(\"expected 2 columns, got %d\", rec.NumCols())\n\t}\n\tif rec.NumRows() != 1 {\n\t\tt.Fatalf(\"expected 1 row, got %d\", rec.NumRows())\n\t}\n}\n\nfunc TestGetArrowBatchesHighPrecision(t *testing.T) {\n\tctx := sf.WithHigherPrecision(WithArrowBatches(context.Background()))\n\n\tsfRows, cleanup := queryRawRows(ctx, t, \"SELECT '0.1'::DECIMAL(38, 19) AS c\")\n\tdefer cleanup()\n\n\tbatches, err := GetArrowBatches(sfRows)\n\tif err != nil {\n\t\tt.Fatalf(\"GetArrowBatches failed: %v\", err)\n\t}\n\tif len(batches) == 0 {\n\t\tt.Fatal(\"expected at least one batch\")\n\t}\n\n\trecords, err := batches[0].Fetch()\n\tif err != nil {\n\t\tt.Fatalf(\"Fetch failed: %v\", err)\n\t}\n\tif records == nil || len(*records) == 0 {\n\t\tt.Fatal(\"expected at least one record\")\n\t}\n\n\trec := (*records)[0]\n\tdefer rec.Release()\n\n\tstrVal := rec.Column(0).ValueStr(0)\n\texpected := \"1000000000000000000\"\n\tif strVal != expected {\n\t\tt.Fatalf(\"expected %q, got %q\", expected, strVal)\n\t}\n}\n\nfunc TestGetArrowBatchesLargeResultSet(t *testing.T) {\n\tnumrows := 3000\n\tpool := memory.NewCheckedAllocator(memory.DefaultAllocator)\n\tdefer pool.AssertSize(t, 0)\n\n\tctx := sf.WithArrowAllocator(WithArrowBatches(context.Background()), pool)\n\n\tquery := fmt.Sprintf(\"SELECT SEQ8(), RANDSTR(1000, RANDOM()) FROM TABLE(GENERATOR(ROWCOUNT=>%v))\", numrows)\n\tsfRows, cleanup := queryRawRows(ctx, t, query)\n\tdefer cleanup()\n\n\tbatches, err := GetArrowBatches(sfRows)\n\tif err != nil {\n\t\tt.Fatalf(\"GetArrowBatches failed: %v\", err)\n\t}\n\tif len(batches) == 0 {\n\t\tt.Fatal(\"expected at least one batch\")\n\t}\n\n\tmaxWorkers := 10\n\ttype count struct {\n\t\tmu sync.Mutex\n\t\tval int\n\t}\n\tcnt := &count{}\n\tvar wg sync.WaitGroup\n\twork := make(chan int, len(batches))\n\n\tfor range maxWorkers {\n\t\twg.Add(1)\n\t\tgo func() {\n\t\t\tdefer wg.Done()\n\t\t\tfor i := range work {\n\t\t\t\trecs, fetchErr := batches[i].Fetch()\n\t\t\t\tif fetchErr != nil {\n\t\t\t\t\tt.Errorf(\"Fetch failed for batch %d: %v\", i, fetchErr)\n\t\t\t\t\treturn\n\t\t\t\t}\n\t\t\t\tfor _, r := range *recs {\n\t\t\t\t\tcnt.mu.Lock()\n\t\t\t\t\tcnt.val += int(r.NumRows())\n\t\t\t\t\tcnt.mu.Unlock()\n\t\t\t\t\tr.Release()\n\t\t\t\t}\n\t\t\t}\n\t\t}()\n\t}\n\tfor i := range batches {\n\t\twork <- i\n\t}\n\tclose(work)\n\twg.Wait()\n\n\tif cnt.val != numrows {\n\t\tt.Fatalf(\"row count mismatch: expected %d, got %d\", numrows, cnt.val)\n\t}\n}\n\nfunc TestGetArrowBatchesWithTimestampOption(t *testing.T) {\n\tctx := WithTimestampOption(\n\t\tWithArrowBatches(context.Background()),\n\t\tUseOriginalTimestamp,\n\t)\n\n\tsfRows, cleanup := queryRawRows(ctx, t, \"SELECT TO_TIMESTAMP_NTZ('2024-01-15 13:45:30.123456789') AS ts\")\n\tdefer cleanup()\n\n\tbatches, err := GetArrowBatches(sfRows)\n\tif err != nil {\n\t\tt.Fatalf(\"GetArrowBatches failed: %v\", err)\n\t}\n\tif len(batches) == 0 {\n\t\tt.Fatal(\"expected at least one batch\")\n\t}\n\n\trecords, err := batches[0].Fetch()\n\tif err != nil {\n\t\tt.Fatalf(\"Fetch failed: %v\", err)\n\t}\n\tif records == nil || len(*records) == 0 {\n\t\tt.Fatal(\"expected at least one record\")\n\t}\n\n\trec := (*records)[0]\n\tdefer rec.Release()\n\n\tif rec.NumRows() != 1 {\n\t\tt.Fatalf(\"expected 1 row, got %d\", rec.NumRows())\n\t}\n\tif rec.NumCols() != 1 {\n\t\tt.Fatalf(\"expected 1 column, got %d\", rec.NumCols())\n\t}\n}\n\nfunc TestGetArrowBatchesJSONResponseError(t *testing.T) {\n\tctx := WithArrowBatches(context.Background())\n\n\tcfg := testConfig(t)\n\n\tdsn, err := sf.DSN(cfg)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to create DSN: %v\", err)\n\t}\n\n\tdb, err := sql.Open(\"snowflake\", dsn)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to open db: %v\", err)\n\t}\n\tdefer db.Close()\n\n\tconn, err := db.Conn(ctx)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to get connection: %v\", err)\n\t}\n\tdefer conn.Close()\n\n\t_, err = conn.ExecContext(ctx, \"ALTER SESSION SET GO_QUERY_RESULT_FORMAT = json\")\n\tif err != nil {\n\t\tt.Fatalf(\"failed to set JSON format: %v\", err)\n\t}\n\n\tvar rows driver.Rows\n\terr = conn.Raw(func(x any) error {\n\t\tqueryer, ok := x.(driver.QueryerContext)\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"connection does not implement QueryerContext\")\n\t\t}\n\t\trows, err = queryer.QueryContext(ctx, \"SELECT 'hello'\", nil)\n\t\treturn err\n\t})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to execute query: %v\", err)\n\t}\n\tdefer rows.Close()\n\n\tsfRows, ok := rows.(sf.SnowflakeRows)\n\tif !ok {\n\t\tt.Fatal(\"rows do not implement SnowflakeRows\")\n\t}\n\n\t_, err = GetArrowBatches(sfRows)\n\tif err == nil {\n\t\tt.Fatal(\"expected error when using arrow batches with JSON response\")\n\t}\n\n\tvar se *sf.SnowflakeError\n\tif !errors.As(err, &se) {\n\t\tt.Fatalf(\"expected SnowflakeError, got %T: %v\", err, err)\n\t}\n\tif se.Number != sf.ErrNonArrowResponseInArrowBatches {\n\t\tt.Fatalf(\"expected error code %d, got %d\", sf.ErrNonArrowResponseInArrowBatches, se.Number)\n\t}\n}\n\n// TestTimestampConversionDistantDates tests all 10 timestamp scales (0-9)\n// because each scale exercises a mathematically distinct code path in\n// extractEpoch/extractFraction (converter.go). Past bugs have been\n// scale-specific: SNOW-526255 (time scale for arrow) and SNOW-2091309\n// (precision loss at scale 0). Do not reduce the scale range.\nfunc TestTimestampConversionDistantDates(t *testing.T) {\n\ttimestamps := [2]string{\n\t\t\"9999-12-12 23:59:59.999999999\",\n\t\t\"0001-01-01 00:00:00.000000000\",\n\t}\n\ttsTypes := [3]string{\"TIMESTAMP_NTZ\", \"TIMESTAMP_LTZ\", \"TIMESTAMP_TZ\"}\n\n\tprecisions := []struct {\n\t\tname string\n\t\toption ia.TimestampOption\n\t\tunit arrow.TimeUnit\n\t\texpectError bool\n\t}{\n\t\t{\"second\", UseSecondTimestamp, arrow.Second, false},\n\t\t{\"millisecond\", UseMillisecondTimestamp, arrow.Millisecond, false},\n\t\t{\"microsecond\", UseMicrosecondTimestamp, arrow.Microsecond, false},\n\t\t{\"nanosecond\", UseNanosecondTimestamp, arrow.Nanosecond, true},\n\t}\n\n\tfor _, prec := range precisions {\n\t\tt.Run(prec.name, func(t *testing.T) {\n\t\t\tt.Parallel()\n\t\t\tpool := memory.NewCheckedAllocator(memory.DefaultAllocator)\n\t\t\tdefer pool.AssertSize(t, 0)\n\n\t\t\tctx := sf.WithArrowAllocator(\n\t\t\t\tWithTimestampOption(WithArrowBatches(context.Background()), prec.option),\n\t\t\t\tpool,\n\t\t\t)\n\n\t\t\ttc := openTestConn(ctx, t)\n\t\t\tdefer tc.close()\n\n\t\t\tfor _, tsStr := range timestamps {\n\t\t\t\tfor _, tp := range tsTypes {\n\t\t\t\t\tfor scale := 0; scale <= 9; scale++ {\n\t\t\t\t\t\tt.Run(tp+\"(\"+strconv.Itoa(scale)+\")_\"+tsStr, func(t *testing.T) {\n\t\t\t\t\t\t\tquery := fmt.Sprintf(\"SELECT '%s'::%s(%v)\", tsStr, tp, scale)\n\t\t\t\t\t\t\tsfRows, closeRows := tc.queryRows(ctx, t, query)\n\t\t\t\t\t\t\tdefer closeRows()\n\n\t\t\t\t\t\t\tbatches, err := GetArrowBatches(sfRows)\n\t\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\t\tt.Fatalf(\"GetArrowBatches failed: %v\", err)\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\tif len(batches) == 0 {\n\t\t\t\t\t\t\t\tt.Fatal(\"expected at least one batch\")\n\t\t\t\t\t\t\t}\n\n\t\t\t\t\t\t\trecords, err := batches[0].Fetch()\n\n\t\t\t\t\t\t\tif prec.expectError {\n\t\t\t\t\t\t\t\texpectedError := \"Cannot convert timestamp\"\n\t\t\t\t\t\t\t\tif err == nil {\n\t\t\t\t\t\t\t\t\tt.Fatalf(\"no error, expected: %v\", expectedError)\n\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t\tif !strings.Contains(err.Error(), expectedError) {\n\t\t\t\t\t\t\t\t\tt.Fatalf(\"improper error, expected: %v, got: %v\", expectedError, err.Error())\n\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t\treturn\n\t\t\t\t\t\t\t}\n\n\t\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\t\tt.Fatalf(\"Fetch failed: %v\", err)\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\tif records == nil || len(*records) == 0 {\n\t\t\t\t\t\t\t\tt.Fatal(\"expected at least one record\")\n\t\t\t\t\t\t\t}\n\n\t\t\t\t\t\t\trec := (*records)[0]\n\t\t\t\t\t\t\tdefer rec.Release()\n\n\t\t\t\t\t\t\tactual := rec.Column(0).(*array.Timestamp).TimestampValues()[0]\n\t\t\t\t\t\t\tactualYear := actual.ToTime(prec.unit).Year()\n\n\t\t\t\t\t\t\tts, err := time.Parse(\"2006-01-02 15:04:05\", tsStr)\n\t\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\t\tt.Fatalf(\"failed to parse time: %v\", err)\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\texp := ts.Truncate(time.Duration(math.Pow10(9 - scale)))\n\n\t\t\t\t\t\t\tif actualYear != exp.Year() {\n\t\t\t\t\t\t\t\tt.Fatalf(\"unexpected year, expected: %v, got: %v\", exp.Year(), actualYear)\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t})\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t}\n}\n\n// TestTimestampConversionWithOriginalTimestamp tests all 10 timestamp scales\n// (0-9) because each scale exercises a mathematically distinct code path in\n// extractEpoch/extractFraction. See TestTimestampConversionDistantDates for\n// rationale on why the full scale range is required.\nfunc TestTimestampConversionWithOriginalTimestamp(t *testing.T) {\n\ttimestamps := [3]string{\n\t\t\"2000-10-10 10:10:10.123456789\",\n\t\t\"9999-12-12 23:59:59.999999999\",\n\t\t\"0001-01-01 00:00:00.000000000\",\n\t}\n\ttsTypes := [3]string{\"TIMESTAMP_NTZ\", \"TIMESTAMP_LTZ\", \"TIMESTAMP_TZ\"}\n\n\tpool := memory.NewCheckedAllocator(memory.DefaultAllocator)\n\tdefer pool.AssertSize(t, 0)\n\n\tctx := sf.WithArrowAllocator(\n\t\tWithTimestampOption(WithArrowBatches(context.Background()), UseOriginalTimestamp),\n\t\tpool,\n\t)\n\n\ttc := openTestConn(ctx, t)\n\tdefer tc.close()\n\n\tfor _, tsStr := range timestamps {\n\t\tts, err := time.Parse(\"2006-01-02 15:04:05\", tsStr)\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"failed to parse time: %v\", err)\n\t\t}\n\t\tfor _, tp := range tsTypes {\n\t\t\tt.Run(tp+\"_\"+tsStr, func(t *testing.T) {\n\t\t\t\t// Batch all 10 scales into a single multi-column query to reduce round trips.\n\t\t\t\tvar cols []string\n\t\t\t\tfor scale := 0; scale <= 9; scale++ {\n\t\t\t\t\tcols = append(cols, fmt.Sprintf(\"'%s'::%s(%v)\", tsStr, tp, scale))\n\t\t\t\t}\n\t\t\t\tquery := \"SELECT \" + strings.Join(cols, \", \")\n\t\t\t\tsfRows, closeRows := tc.queryRows(ctx, t, query)\n\t\t\t\tdefer closeRows()\n\n\t\t\t\tbatches, err := GetArrowBatches(sfRows)\n\t\t\t\tif err != nil {\n\t\t\t\t\tt.Fatalf(\"GetArrowBatches failed: %v\", err)\n\t\t\t\t}\n\t\t\t\tif len(batches) != 1 {\n\t\t\t\t\tt.Fatalf(\"expected 1 batch, got %d\", len(batches))\n\t\t\t\t}\n\n\t\t\t\trecords, err := batches[0].Fetch()\n\t\t\t\tif err != nil {\n\t\t\t\t\tt.Fatalf(\"Fetch failed: %v\", err)\n\t\t\t\t}\n\t\t\t\tif records == nil || len(*records) == 0 {\n\t\t\t\t\tt.Fatal(\"expected at least one record\")\n\t\t\t\t}\n\n\t\t\t\tfor scale := 0; scale <= 9; scale++ {\n\t\t\t\t\texp := ts.Truncate(time.Duration(math.Pow10(9 - scale)))\n\t\t\t\t\tfor _, r := range *records {\n\t\t\t\t\t\tdefer r.Release()\n\t\t\t\t\t\tact := batches[0].ArrowSnowflakeTimestampToTime(r, scale, 0)\n\t\t\t\t\t\tif act == nil {\n\t\t\t\t\t\t\tt.Fatalf(\"scale %d: unexpected nil, expected: %v\", scale, exp)\n\t\t\t\t\t\t} else if !exp.Equal(*act) {\n\t\t\t\t\t\t\tt.Fatalf(\"scale %d: unexpected result, expected: %v, got: %v\", scale, exp, *act)\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t})\n\t\t}\n\t}\n}\n"
},
{
"path": "arrowbatches/context.go",
"content": "package arrowbatches\n\nimport (\n\t\"context\"\n\n\tia \"github.com/snowflakedb/gosnowflake/v2/internal/arrow\"\n)\n\n// Timestamp option constants.\nconst (\n\tUseNanosecondTimestamp = ia.UseNanosecondTimestamp\n\tUseMicrosecondTimestamp = ia.UseMicrosecondTimestamp\n\tUseMillisecondTimestamp = ia.UseMillisecondTimestamp\n\tUseSecondTimestamp = ia.UseSecondTimestamp\n\tUseOriginalTimestamp = ia.UseOriginalTimestamp\n)\n\n// WithArrowBatches returns a context that enables arrow batch mode for queries.\nfunc WithArrowBatches(ctx context.Context) context.Context {\n\treturn ia.EnableArrowBatches(ctx)\n}\n\n// WithTimestampOption returns a context that sets the timestamp conversion option\n// for arrow batches.\nfunc WithTimestampOption(ctx context.Context, option ia.TimestampOption) context.Context {\n\treturn ia.WithTimestampOption(ctx, option)\n}\n\n// WithUtf8Validation returns a context that enables UTF-8 validation for\n// string columns in arrow batches.\nfunc WithUtf8Validation(ctx context.Context) context.Context {\n\treturn ia.EnableUtf8Validation(ctx)\n}\n"
},
{
"path": "arrowbatches/converter.go",
"content": "package arrowbatches\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/query\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/types\"\n\t\"math\"\n\t\"math/big\"\n\t\"strings\"\n\t\"time\"\n\t\"unicode/utf8\"\n\n\tsf \"github.com/snowflakedb/gosnowflake/v2\"\n\tia \"github.com/snowflakedb/gosnowflake/v2/internal/arrow\"\n\n\t\"github.com/apache/arrow-go/v18/arrow\"\n\t\"github.com/apache/arrow-go/v18/arrow/array\"\n\t\"github.com/apache/arrow-go/v18/arrow/compute\"\n\t\"github.com/apache/arrow-go/v18/arrow/memory\"\n)\n\n// arrowToRecord transforms a raw arrow.Record from Snowflake into a record\n// with standard Arrow types (e.g., converting struct-based timestamps to\n// arrow.Timestamp, decimal128 to int64/float64, etc.)\nfunc arrowToRecord(ctx context.Context, record arrow.Record, pool memory.Allocator, rowType []query.ExecResponseRowType, loc *time.Location) (arrow.Record, error) {\n\ttimestampOption := ia.GetTimestampOption(ctx)\n\thigherPrecision := ia.HigherPrecisionEnabled(ctx)\n\n\ts, err := recordToSchema(record.Schema(), rowType, loc, timestampOption, higherPrecision)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tvar cols []arrow.Array\n\tnumRows := record.NumRows()\n\tctxAlloc := compute.WithAllocator(ctx, pool)\n\n\tfor i, col := range record.Columns() {\n\t\tfieldMetadata := rowType[i].ToFieldMetadata()\n\n\t\tnewCol, err := arrowToRecordSingleColumn(ctxAlloc, s.Field(i), col, fieldMetadata, higherPrecision, timestampOption, pool, loc, numRows)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tcols = append(cols, newCol)\n\t\tdefer newCol.Release()\n\t}\n\tnewRecord := array.NewRecord(s, cols, numRows)\n\treturn newRecord, nil\n}\n\nfunc arrowToRecordSingleColumn(ctx context.Context, field arrow.Field, col arrow.Array, fieldMetadata query.FieldMetadata, higherPrecisionEnabled bool, timestampOption ia.TimestampOption, pool memory.Allocator, loc *time.Location, numRows int64) (arrow.Array, error) {\n\tvar err error\n\tnewCol := col\n\tsnowflakeType := types.GetSnowflakeType(fieldMetadata.Type)\n\tswitch snowflakeType {\n\tcase types.FixedType:\n\t\tif higherPrecisionEnabled {\n\t\t\tcol.Retain()\n\t\t} else if col.DataType().ID() == arrow.DECIMAL || col.DataType().ID() == arrow.DECIMAL256 {\n\t\t\tvar toType arrow.DataType\n\t\t\tif fieldMetadata.Scale == 0 {\n\t\t\t\ttoType = arrow.PrimitiveTypes.Int64\n\t\t\t} else {\n\t\t\t\ttoType = arrow.PrimitiveTypes.Float64\n\t\t\t}\n\t\t\tnewCol, err = compute.CastArray(ctx, col, compute.UnsafeCastOptions(toType))\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t} else if fieldMetadata.Scale != 0 && col.DataType().ID() != arrow.INT64 {\n\t\t\tresult, err := compute.Divide(ctx, compute.ArithmeticOptions{NoCheckOverflow: true},\n\t\t\t\t&compute.ArrayDatum{Value: newCol.Data()},\n\t\t\t\tcompute.NewDatum(math.Pow10(int(fieldMetadata.Scale))))\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tdefer result.Release()\n\t\t\tnewCol = result.(*compute.ArrayDatum).MakeArray()\n\t\t} else if fieldMetadata.Scale != 0 && col.DataType().ID() == arrow.INT64 {\n\t\t\tvalues := col.(*array.Int64).Int64Values()\n\t\t\tfloatValues := make([]float64, len(values))\n\t\t\tfor i, val := range values {\n\t\t\t\tfloatValues[i], _ = intToBigFloat(val, int64(fieldMetadata.Scale)).Float64()\n\t\t\t}\n\t\t\tbuilder := array.NewFloat64Builder(pool)\n\t\t\tbuilder.AppendValues(floatValues, nil)\n\t\t\tnewCol = builder.NewArray()\n\t\t\tbuilder.Release()\n\t\t} else {\n\t\t\tcol.Retain()\n\t\t}\n\tcase types.TimeType:\n\t\tnewCol, err = compute.CastArray(ctx, col, compute.SafeCastOptions(arrow.FixedWidthTypes.Time64ns))\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\tcase types.TimestampNtzType, types.TimestampLtzType, types.TimestampTzType:\n\t\tif timestampOption == ia.UseOriginalTimestamp {\n\t\t\tcol.Retain()\n\t\t} else {\n\t\t\tvar unit arrow.TimeUnit\n\t\t\tswitch timestampOption {\n\t\t\tcase ia.UseMicrosecondTimestamp:\n\t\t\t\tunit = arrow.Microsecond\n\t\t\tcase ia.UseMillisecondTimestamp:\n\t\t\t\tunit = arrow.Millisecond\n\t\t\tcase ia.UseSecondTimestamp:\n\t\t\t\tunit = arrow.Second\n\t\t\tcase ia.UseNanosecondTimestamp:\n\t\t\t\tunit = arrow.Nanosecond\n\t\t\t}\n\t\t\tvar tb *array.TimestampBuilder\n\t\t\tif snowflakeType == types.TimestampLtzType {\n\t\t\t\ttb = array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: unit, TimeZone: loc.String()})\n\t\t\t} else {\n\t\t\t\ttb = array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: unit})\n\t\t\t}\n\t\t\tdefer tb.Release()\n\n\t\t\tfor i := 0; i < int(numRows); i++ {\n\t\t\t\tts := ArrowSnowflakeTimestampToTime(col, snowflakeType, int(fieldMetadata.Scale), i, loc)\n\t\t\t\tif ts != nil {\n\t\t\t\t\tvar ar arrow.Timestamp\n\t\t\t\t\tswitch timestampOption {\n\t\t\t\t\tcase ia.UseMicrosecondTimestamp:\n\t\t\t\t\t\tar = arrow.Timestamp(ts.UnixMicro())\n\t\t\t\t\tcase ia.UseMillisecondTimestamp:\n\t\t\t\t\t\tar = arrow.Timestamp(ts.UnixMilli())\n\t\t\t\t\tcase ia.UseSecondTimestamp:\n\t\t\t\t\t\tar = arrow.Timestamp(ts.Unix())\n\t\t\t\t\tcase ia.UseNanosecondTimestamp:\n\t\t\t\t\t\tar = arrow.Timestamp(ts.UnixNano())\n\t\t\t\t\t\tif ts.UTC().Year() != ar.ToTime(arrow.Nanosecond).Year() {\n\t\t\t\t\t\t\treturn nil, &sf.SnowflakeError{\n\t\t\t\t\t\t\t\tNumber: sf.ErrTooHighTimestampPrecision,\n\t\t\t\t\t\t\t\tSQLState: sf.SQLStateInvalidDataTimeFormat,\n\t\t\t\t\t\t\t\tMessage: fmt.Sprintf(\"Cannot convert timestamp %v in column %v to Arrow.Timestamp data type due to too high precision. Please use context with WithOriginalTimestamp.\", ts.UTC(), fieldMetadata.Name),\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\ttb.Append(ar)\n\t\t\t\t} else {\n\t\t\t\t\ttb.AppendNull()\n\t\t\t\t}\n\t\t\t}\n\t\t\tnewCol = tb.NewArray()\n\t\t}\n\tcase types.TextType:\n\t\tif stringCol, ok := col.(*array.String); ok {\n\t\t\tnewCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows)\n\t\t}\n\tcase types.ObjectType:\n\t\tif structCol, ok := col.(*array.Struct); ok {\n\t\t\tvar internalCols []arrow.Array\n\t\t\tfor i := 0; i < structCol.NumField(); i++ {\n\t\t\t\tinternalCol := structCol.Field(i)\n\t\t\t\tnewInternalCol, err := arrowToRecordSingleColumn(ctx, field.Type.(*arrow.StructType).Field(i), internalCol, fieldMetadata.Fields[i], higherPrecisionEnabled, timestampOption, pool, loc, numRows)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, err\n\t\t\t\t}\n\t\t\t\tinternalCols = append(internalCols, newInternalCol)\n\t\t\t\tdefer newInternalCol.Release()\n\t\t\t}\n\t\t\tvar fieldNames []string\n\t\t\tfor _, f := range field.Type.(*arrow.StructType).Fields() {\n\t\t\t\tfieldNames = append(fieldNames, f.Name)\n\t\t\t}\n\t\t\tnullBitmap := memory.NewBufferBytes(structCol.NullBitmapBytes())\n\t\t\tnumberOfNulls := structCol.NullN()\n\t\t\treturn array.NewStructArrayWithNulls(internalCols, fieldNames, nullBitmap, numberOfNulls, 0)\n\t\t} else if stringCol, ok := col.(*array.String); ok {\n\t\t\tnewCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows)\n\t\t}\n\tcase types.ArrayType:\n\t\tif listCol, ok := col.(*array.List); ok {\n\t\t\tnewCol, err = arrowToRecordSingleColumn(ctx, field.Type.(*arrow.ListType).ElemField(), listCol.ListValues(), fieldMetadata.Fields[0], higherPrecisionEnabled, timestampOption, pool, loc, numRows)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tdefer newCol.Release()\n\t\t\tnewData := array.NewData(arrow.ListOf(newCol.DataType()), listCol.Len(), listCol.Data().Buffers(), []arrow.ArrayData{newCol.Data()}, listCol.NullN(), 0)\n\t\t\tdefer newData.Release()\n\t\t\treturn array.NewListData(newData), nil\n\t\t} else if stringCol, ok := col.(*array.String); ok {\n\t\t\tnewCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows)\n\t\t}\n\tcase types.MapType:\n\t\tif mapCol, ok := col.(*array.Map); ok {\n\t\t\tkeyCol, err := arrowToRecordSingleColumn(ctx, field.Type.(*arrow.MapType).KeyField(), mapCol.Keys(), fieldMetadata.Fields[0], higherPrecisionEnabled, timestampOption, pool, loc, numRows)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tdefer keyCol.Release()\n\t\t\tvalueCol, err := arrowToRecordSingleColumn(ctx, field.Type.(*arrow.MapType).ItemField(), mapCol.Items(), fieldMetadata.Fields[1], higherPrecisionEnabled, timestampOption, pool, loc, numRows)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tdefer valueCol.Release()\n\n\t\t\tstructArr, err := array.NewStructArray([]arrow.Array{keyCol, valueCol}, []string{\"k\", \"v\"})\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tdefer structArr.Release()\n\t\t\tnewData := array.NewData(arrow.MapOf(keyCol.DataType(), valueCol.DataType()), mapCol.Len(), mapCol.Data().Buffers(), []arrow.ArrayData{structArr.Data()}, mapCol.NullN(), 0)\n\t\t\tdefer newData.Release()\n\t\t\treturn array.NewMapData(newData), nil\n\t\t} else if stringCol, ok := col.(*array.String); ok {\n\t\t\tnewCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows)\n\t\t}\n\tdefault:\n\t\tcol.Retain()\n\t}\n\treturn newCol, nil\n}\n\nfunc arrowStringRecordToColumn(\n\tctx context.Context,\n\tstringCol *array.String,\n\tmem memory.Allocator,\n\tnumRows int64,\n) arrow.Array {\n\tif ia.Utf8ValidationEnabled(ctx) && stringCol.DataType().ID() == arrow.STRING {\n\t\ttb := array.NewStringBuilder(mem)\n\t\tdefer tb.Release()\n\n\t\tfor i := 0; i < int(numRows); i++ {\n\t\t\tif stringCol.IsValid(i) {\n\t\t\t\tstringValue := stringCol.Value(i)\n\t\t\t\tif !utf8.ValidString(stringValue) {\n\t\t\t\t\tstringValue = strings.ToValidUTF8(stringValue, \"�\")\n\t\t\t\t}\n\t\t\t\ttb.Append(stringValue)\n\t\t\t} else {\n\t\t\t\ttb.AppendNull()\n\t\t\t}\n\t\t}\n\t\tarr := tb.NewArray()\n\t\treturn arr\n\t}\n\tstringCol.Retain()\n\treturn stringCol\n}\n\nfunc intToBigFloat(val int64, scale int64) *big.Float {\n\tf := new(big.Float).SetInt64(val)\n\ts := new(big.Float).SetInt(new(big.Int).Exp(big.NewInt(10), big.NewInt(scale), nil))\n\treturn new(big.Float).Quo(f, s)\n}\n\n// ArrowSnowflakeTimestampToTime converts original timestamp returned by Snowflake to time.Time.\nfunc ArrowSnowflakeTimestampToTime(\n\tcolumn arrow.Array,\n\tsfType types.SnowflakeType,\n\tscale int,\n\trecIdx int,\n\tloc *time.Location) *time.Time {\n\n\tif column.IsNull(recIdx) {\n\t\treturn nil\n\t}\n\tvar ret time.Time\n\tswitch sfType {\n\tcase types.TimestampNtzType:\n\t\tif column.DataType().ID() == arrow.STRUCT {\n\t\t\tstructData := column.(*array.Struct)\n\t\t\tepoch := structData.Field(0).(*array.Int64).Int64Values()\n\t\t\tfraction := structData.Field(1).(*array.Int32).Int32Values()\n\t\t\tret = time.Unix(epoch[recIdx], int64(fraction[recIdx])).UTC()\n\t\t} else {\n\t\t\tintData := column.(*array.Int64)\n\t\t\tvalue := intData.Value(recIdx)\n\t\t\tepoch := extractEpoch(value, scale)\n\t\t\tfraction := extractFraction(value, scale)\n\t\t\tret = time.Unix(epoch, fraction).UTC()\n\t\t}\n\tcase types.TimestampLtzType:\n\t\tif column.DataType().ID() == arrow.STRUCT {\n\t\t\tstructData := column.(*array.Struct)\n\t\t\tepoch := structData.Field(0).(*array.Int64).Int64Values()\n\t\t\tfraction := structData.Field(1).(*array.Int32).Int32Values()\n\t\t\tret = time.Unix(epoch[recIdx], int64(fraction[recIdx])).In(loc)\n\t\t} else {\n\t\t\tintData := column.(*array.Int64)\n\t\t\tvalue := intData.Value(recIdx)\n\t\t\tepoch := extractEpoch(value, scale)\n\t\t\tfraction := extractFraction(value, scale)\n\t\t\tret = time.Unix(epoch, fraction).In(loc)\n\t\t}\n\tcase types.TimestampTzType:\n\t\tstructData := column.(*array.Struct)\n\t\tif structData.NumField() == 2 {\n\t\t\tvalue := structData.Field(0).(*array.Int64).Int64Values()\n\t\t\ttimezone := structData.Field(1).(*array.Int32).Int32Values()\n\t\t\tepoch := extractEpoch(value[recIdx], scale)\n\t\t\tfraction := extractFraction(value[recIdx], scale)\n\t\t\tlocTz := sf.Location(int(timezone[recIdx]) - 1440)\n\t\t\tret = time.Unix(epoch, fraction).In(locTz)\n\t\t} else {\n\t\t\tepoch := structData.Field(0).(*array.Int64).Int64Values()\n\t\t\tfraction := structData.Field(1).(*array.Int32).Int32Values()\n\t\t\ttimezone := structData.Field(2).(*array.Int32).Int32Values()\n\t\t\tlocTz := sf.Location(int(timezone[recIdx]) - 1440)\n\t\t\tret = time.Unix(epoch[recIdx], int64(fraction[recIdx])).In(locTz)\n\t\t}\n\t}\n\treturn &ret\n}\n\nfunc extractEpoch(value int64, scale int) int64 {\n\treturn value / int64(math.Pow10(scale))\n}\n\nfunc extractFraction(value int64, scale int) int64 {\n\treturn (value % int64(math.Pow10(scale))) * int64(math.Pow10(9-scale))\n}\n"
},
{
"path": "arrowbatches/converter_test.go",
"content": "package arrowbatches\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/query\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/types\"\n\t\"math/big\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n\n\tia \"github.com/snowflakedb/gosnowflake/v2/internal/arrow\"\n\n\t\"github.com/apache/arrow-go/v18/arrow\"\n\t\"github.com/apache/arrow-go/v18/arrow/array\"\n\t\"github.com/apache/arrow-go/v18/arrow/decimal128\"\n\t\"github.com/apache/arrow-go/v18/arrow/memory\"\n)\n\nvar decimalShift = new(big.Int).Exp(big.NewInt(2), big.NewInt(64), nil)\n\nfunc stringIntToDecimal(src string) (decimal128.Num, bool) {\n\tb, ok := new(big.Int).SetString(src, 10)\n\tif !ok {\n\t\treturn decimal128.Num{}, ok\n\t}\n\tvar high, low big.Int\n\thigh.QuoRem(b, decimalShift, &low)\n\treturn decimal128.New(high.Int64(), low.Uint64()), true\n}\n\nfunc decimalToBigInt(num decimal128.Num) *big.Int {\n\thigh := new(big.Int).SetInt64(num.HighBits())\n\tlow := new(big.Int).SetUint64(num.LowBits())\n\treturn new(big.Int).Add(new(big.Int).Mul(high, decimalShift), low)\n}\n\nfunc TestArrowToRecord(t *testing.T) {\n\tpool := memory.NewCheckedAllocator(memory.NewGoAllocator())\n\tdefer pool.AssertSize(t, 0)\n\tvar valids []bool\n\n\tlocalTime := time.Date(2019, 1, 1, 1, 17, 31, 123456789, time.FixedZone(\"-08:00\", -8*3600))\n\tlocalTimeFarIntoFuture := time.Date(9000, 2, 6, 14, 17, 31, 123456789, time.FixedZone(\"-08:00\", -8*3600))\n\n\tepochField := arrow.Field{Name: \"epoch\", Type: &arrow.Int64Type{}}\n\ttimezoneField := arrow.Field{Name: \"timezone\", Type: &arrow.Int32Type{}}\n\tfractionField := arrow.Field{Name: \"fraction\", Type: &arrow.Int32Type{}}\n\ttimestampTzStructWithoutFraction := arrow.StructOf(epochField, timezoneField)\n\ttimestampTzStructWithFraction := arrow.StructOf(epochField, fractionField, timezoneField)\n\ttimestampNtzStruct := arrow.StructOf(epochField, fractionField)\n\ttimestampLtzStruct := arrow.StructOf(epochField, fractionField)\n\n\ttype testObj struct {\n\t\tfield1 int\n\t\tfield2 string\n\t}\n\n\tfor _, tc := range []struct {\n\t\tlogical string\n\t\tphysical string\n\t\tsc *arrow.Schema\n\t\trowType query.ExecResponseRowType\n\t\tvalues any\n\t\texpected any\n\t\terror string\n\t\tarrowBatchesTimestampOption ia.TimestampOption\n\t\tenableArrowBatchesUtf8Validation bool\n\t\twithHigherPrecision bool\n\t\tnrows int\n\t\tbuilder array.Builder\n\t\tappend func(b array.Builder, vs any)\n\t\tcompare func(src any, expected any, rec arrow.Record) int\n\t}{\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"number\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tvalues: []int64{1, 2},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int64Builder).AppendValues(vs.([]int64), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int64\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Decimal128Type{Precision: 38, Scale: 0}}}, nil),\n\t\t\tvalues: []string{\"10000000000000000000000000000000000000\", \"-12345678901234567890123456789012345678\"},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewDecimal128Builder(pool, &arrow.Decimal128Type{Precision: 38, Scale: 0}),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, s := range vs.([]string) {\n\t\t\t\t\tnum, ok := stringIntToDecimal(s)\n\t\t\t\t\tif !ok {\n\t\t\t\t\t\tt.Fatalf(\"failed to convert to Int64\")\n\t\t\t\t\t}\n\t\t\t\t\tb.(*array.Decimal128Builder).Append(num)\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]string)\n\t\t\t\tfor i, dec := range convertedRec.Column(0).(*array.Int64).Int64Values() {\n\t\t\t\t\tnum, ok := stringIntToDecimal(srcvs[i])\n\t\t\t\t\tif !ok {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t\tsrcDec := decimalToBigInt(num).Int64()\n\t\t\t\t\tif srcDec != dec {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"number(38,0)\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Decimal128Type{Precision: 38, Scale: 0}}}, nil),\n\t\t\tvalues: []string{\"10000000000000000000000000000000000000\", \"-12345678901234567890123456789012345678\"},\n\t\t\twithHigherPrecision: true,\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewDecimal128Builder(pool, &arrow.Decimal128Type{Precision: 38, Scale: 0}),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, s := range vs.([]string) {\n\t\t\t\t\tnum, ok := stringIntToDecimal(s)\n\t\t\t\t\tif !ok {\n\t\t\t\t\t\tt.Fatalf(\"failed to convert to Int64\")\n\t\t\t\t\t}\n\t\t\t\t\tb.(*array.Decimal128Builder).Append(num)\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]string)\n\t\t\t\tfor i, dec := range convertedRec.Column(0).(*array.Decimal128).Values() {\n\t\t\t\t\tsrcDec, ok := stringIntToDecimal(srcvs[i])\n\t\t\t\t\tif !ok {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t\tif srcDec != dec {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"float64\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 37},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Decimal128Type{Precision: 38, Scale: 37}}}, nil),\n\t\t\tvalues: []string{\"1.2345678901234567890123456789012345678\", \"-9.999999999999999\"},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewDecimal128Builder(pool, &arrow.Decimal128Type{Precision: 38, Scale: 37}),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, s := range vs.([]string) {\n\t\t\t\t\tnum, err := decimal128.FromString(s, 38, 37)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\tt.Fatalf(\"failed to convert to decimal: %s\", err)\n\t\t\t\t\t}\n\t\t\t\t\tb.(*array.Decimal128Builder).Append(num)\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]string)\n\t\t\t\tfor i, dec := range convertedRec.Column(0).(*array.Float64).Float64Values() {\n\t\t\t\t\tnum, err := decimal128.FromString(srcvs[i], 38, 37)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t\tsrcDec := num.ToFloat64(37)\n\t\t\t\t\tif srcDec != dec {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"number(38,37)\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 37},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Decimal128Type{Precision: 38, Scale: 37}}}, nil),\n\t\t\tvalues: []string{\"1.2345678901234567890123456789012345678\", \"-9.999999999999999\"},\n\t\t\twithHigherPrecision: true,\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewDecimal128Builder(pool, &arrow.Decimal128Type{Precision: 38, Scale: 37}),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, s := range vs.([]string) {\n\t\t\t\t\tnum, err := decimal128.FromString(s, 38, 37)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\tt.Fatalf(\"failed to convert to decimal: %s\", err)\n\t\t\t\t\t}\n\t\t\t\t\tb.(*array.Decimal128Builder).Append(num)\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]string)\n\t\t\t\tfor i, dec := range convertedRec.Column(0).(*array.Decimal128).Values() {\n\t\t\t\t\tsrcDec, err := decimal128.FromString(srcvs[i], 38, 37)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t\tif srcDec != dec {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int8\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int8Type{}}}, nil),\n\t\t\tvalues: []int8{1, 2},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt8Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int8Builder).AppendValues(vs.([]int8), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int16\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int16Type{}}}, nil),\n\t\t\tvalues: []int16{1, 2},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt16Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int16Builder).AppendValues(vs.([]int16), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int32\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int32Type{}}}, nil),\n\t\t\tvalues: []int32{1, 2},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt32Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int32Builder).AppendValues(vs.([]int32), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int64\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tvalues: []int64{1, 2},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int64Builder).AppendValues(vs.([]int64), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"float8\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 1},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int8Type{}}}, nil),\n\t\t\tvalues: []int8{10, 16},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt8Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int8Builder).AppendValues(vs.([]int8), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int8)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Float64).Float64Values() {\n\t\t\t\t\trawFloat, _ := intToBigFloat(int64(srcvs[i]), 1).Float64()\n\t\t\t\t\tif rawFloat != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int8\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 1},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int8Type{}}}, nil),\n\t\t\tvalues: []int8{10, 16},\n\t\t\twithHigherPrecision: true,\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt8Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int8Builder).AppendValues(vs.([]int8), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int8)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Int8).Int8Values() {\n\t\t\t\t\tif srcvs[i] != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"float16\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 1},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int16Type{}}}, nil),\n\t\t\tvalues: []int16{20, 26},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt16Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int16Builder).AppendValues(vs.([]int16), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int16)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Float64).Float64Values() {\n\t\t\t\t\trawFloat, _ := intToBigFloat(int64(srcvs[i]), 1).Float64()\n\t\t\t\t\tif rawFloat != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int16\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 1},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int16Type{}}}, nil),\n\t\t\tvalues: []int16{20, 26},\n\t\t\twithHigherPrecision: true,\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt16Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int16Builder).AppendValues(vs.([]int16), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int16)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Int16).Int16Values() {\n\t\t\t\t\tif srcvs[i] != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"float32\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 2},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int32Type{}}}, nil),\n\t\t\tvalues: []int32{200, 265},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt32Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int32Builder).AppendValues(vs.([]int32), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int32)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Float64).Float64Values() {\n\t\t\t\t\trawFloat, _ := intToBigFloat(int64(srcvs[i]), 2).Float64()\n\t\t\t\t\tif rawFloat != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int32\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 2},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int32Type{}}}, nil),\n\t\t\tvalues: []int32{200, 265},\n\t\t\twithHigherPrecision: true,\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt32Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int32Builder).AppendValues(vs.([]int32), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int32)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Int32).Int32Values() {\n\t\t\t\t\tif srcvs[i] != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"float64\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 5},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tvalues: []int64{12345, 234567},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int64Builder).AppendValues(vs.([]int64), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int64)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Float64).Float64Values() {\n\t\t\t\t\trawFloat, _ := intToBigFloat(srcvs[i], 5).Float64()\n\t\t\t\t\tif rawFloat != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"fixed\",\n\t\t\tphysical: \"int64\",\n\t\t\trowType: query.ExecResponseRowType{Scale: 5},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tvalues: []int64{12345, 234567},\n\t\t\twithHigherPrecision: true,\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Int64Builder).AppendValues(vs.([]int64), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]int64)\n\t\t\t\tfor i, f := range convertedRec.Column(0).(*array.Int64).Int64Values() {\n\t\t\t\t\tif srcvs[i] != f {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"boolean\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.BooleanType{}}}, nil),\n\t\t\tvalues: []bool{true, false},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewBooleanBuilder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.BooleanBuilder).AppendValues(vs.([]bool), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"real\",\n\t\t\tphysical: \"float\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Float64Type{}}}, nil),\n\t\t\tvalues: []float64{1, 2},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewFloat64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.Float64Builder).AppendValues(vs.([]float64), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"text\",\n\t\t\tphysical: \"string\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.StringType{}}}, nil),\n\t\t\tvalues: []string{\"foo\", \"bar\"},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewStringBuilder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.StringBuilder).AppendValues(vs.([]string), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"text\",\n\t\t\tphysical: \"string with invalid utf8\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.StringType{}}}, nil),\n\t\t\trowType: query.ExecResponseRowType{Type: \"TEXT\"},\n\t\t\tvalues: []string{\"\\xFF\", \"bar\", \"baz\\xFF\\xFF\"},\n\t\t\texpected: []string{\"�\", \"bar\", \"baz��\"},\n\t\t\tenableArrowBatchesUtf8Validation: true,\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewStringBuilder(pool),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.StringBuilder).AppendValues(vs.([]string), valids) },\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tarr := convertedRec.Column(0).(*array.String)\n\t\t\t\tfor i := 0; i < arr.Len(); i++ {\n\t\t\t\t\tif expected.([]string)[i] != arr.Value(i) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"binary\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.BinaryType{}}}, nil),\n\t\t\tvalues: [][]byte{[]byte(\"foo\"), []byte(\"bar\")},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewBinaryBuilder(pool, arrow.BinaryTypes.Binary),\n\t\t\tappend: func(b array.Builder, vs any) { b.(*array.BinaryBuilder).AppendValues(vs.([][]byte), valids) },\n\t\t},\n\t\t{\n\t\t\tlogical: \"date\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Date32Type{}}}, nil),\n\t\t\tvalues: []time.Time{time.Now(), localTime},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewDate32Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, d := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Date32Builder).Append(arrow.Date32(d.Unix()))\n\t\t\t\t}\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"time\",\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: arrow.FixedWidthTypes.Time64ns}}, nil),\n\t\t\tvalues: []time.Time{time.Now(), time.Now()},\n\t\t\tnrows: 2,\n\t\t\tbuilder: array.NewTime64Builder(pool, arrow.FixedWidthTypes.Time64ns.(*arrow.Time64Type)),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Time64Builder).Append(arrow.Time64(t.UnixNano()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tarr := convertedRec.Column(0).(*array.Time64)\n\t\t\t\tfor i := 0; i < arr.Len(); i++ {\n\t\t\t\t\tif srcvs[i].UnixNano() != int64(arr.Value(i)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"int64\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Nanosecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now(), localTime},\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Nanosecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Microsecond), localTime.Truncate(time.Microsecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseMicrosecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Microsecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseMillisecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Millisecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Second), localTime.Truncate(time.Second)},\n\t\t\tarrowBatchesTimestampOption: ia.UseSecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Second)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"error\",\n\t\t\tvalues: []time.Time{localTimeFarIntoFuture},\n\t\t\terror: \"Cannot convert timestamp\",\n\t\t\tnrows: 1,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int { return 0 },\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"int64 with original timestamp\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseOriginalTimestamp,\n\t\t\tnrows: 3,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i := 0; i < convertedRec.Column(0).Len(); i++ {\n\t\t\t\t\tts := ArrowSnowflakeTimestampToTime(convertedRec.Column(0), types.GetSnowflakeType(\"timestamp_ntz\"), 3, i, nil)\n\t\t\t\t\tif !srcvs[i].Equal(*ts) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ntz\",\n\t\t\tphysical: \"struct with original timestamp\",\n\t\t\tvalues: []time.Time{time.Now(), localTime, localTimeFarIntoFuture},\n\t\t\tarrowBatchesTimestampOption: ia.UseOriginalTimestamp,\n\t\t\tnrows: 3,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i := 0; i < convertedRec.Column(0).Len(); i++ {\n\t\t\t\t\tts := ArrowSnowflakeTimestampToTime(convertedRec.Column(0), types.GetSnowflakeType(\"timestamp_ntz\"), 9, i, nil)\n\t\t\t\t\tif !srcvs[i].Equal(*ts) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"int64\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Nanosecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now(), localTime},\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Nanosecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Microsecond), localTime.Truncate(time.Microsecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseMicrosecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Microsecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseMillisecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Millisecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"struct\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Second), localTime.Truncate(time.Second)},\n\t\t\tarrowBatchesTimestampOption: ia.UseSecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampNtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Second)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"error\",\n\t\t\tvalues: []time.Time{localTimeFarIntoFuture},\n\t\t\terror: \"Cannot convert timestamp\",\n\t\t\tnrows: 1,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int { return 0 },\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"int64 with original timestamp\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseOriginalTimestamp,\n\t\t\tnrows: 3,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),\n\t\t\tbuilder: array.NewInt64Builder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tb.(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i := 0; i < convertedRec.Column(0).Len(); i++ {\n\t\t\t\t\tts := ArrowSnowflakeTimestampToTime(convertedRec.Column(0), types.GetSnowflakeType(\"timestamp_ltz\"), 3, i, localTime.Location())\n\t\t\t\t\tif !srcvs[i].Equal(*ts) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_ltz\",\n\t\t\tphysical: \"struct with original timestamp\",\n\t\t\tvalues: []time.Time{time.Now(), localTime, localTimeFarIntoFuture},\n\t\t\tarrowBatchesTimestampOption: ia.UseOriginalTimestamp,\n\t\t\tnrows: 3,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampLtzStruct}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampLtzStruct),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i := 0; i < convertedRec.Column(0).Len(); i++ {\n\t\t\t\t\tts := ArrowSnowflakeTimestampToTime(convertedRec.Column(0), types.GetSnowflakeType(\"timestamp_ltz\"), 9, i, localTime.Location())\n\t\t\t\t\tif !srcvs[i].Equal(*ts) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_tz\",\n\t\t\tphysical: \"struct2\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithoutFraction}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampTzStructWithoutFraction),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(0))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Nanosecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_tz\",\n\t\t\tphysical: \"struct3\",\n\t\t\tvalues: []time.Time{time.Now(), localTime},\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampTzStructWithFraction),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t\tsb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Nanosecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_tz\",\n\t\t\tphysical: \"struct3\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Microsecond), localTime.Truncate(time.Microsecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseMicrosecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampTzStructWithFraction),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t\tsb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Microsecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_tz\",\n\t\t\tphysical: \"struct3\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseMillisecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampTzStructWithFraction),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t\tsb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Millisecond)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_tz\",\n\t\t\tphysical: \"struct3\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Second), localTime.Truncate(time.Second)},\n\t\t\tarrowBatchesTimestampOption: ia.UseSecondTimestamp,\n\t\t\tnrows: 2,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampTzStructWithFraction),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t\tsb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {\n\t\t\t\t\tif !srcvs[i].Equal(t.ToTime(arrow.Second)) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_tz\",\n\t\t\tphysical: \"struct2 with original timestamp\",\n\t\t\tvalues: []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},\n\t\t\tarrowBatchesTimestampOption: ia.UseOriginalTimestamp,\n\t\t\tnrows: 3,\n\t\t\trowType: query.ExecResponseRowType{Scale: 3},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithoutFraction}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampTzStructWithoutFraction),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.UnixMilli())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(0))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i := 0; i < convertedRec.Column(0).Len(); i++ {\n\t\t\t\t\tts := ArrowSnowflakeTimestampToTime(convertedRec.Column(0), types.GetSnowflakeType(\"timestamp_tz\"), 3, i, nil)\n\t\t\t\t\tif !srcvs[i].Equal(*ts) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"timestamp_tz\",\n\t\t\tphysical: \"struct3 with original timestamp\",\n\t\t\tvalues: []time.Time{time.Now(), localTime, localTimeFarIntoFuture},\n\t\t\tarrowBatchesTimestampOption: ia.UseOriginalTimestamp,\n\t\t\tnrows: 3,\n\t\t\trowType: query.ExecResponseRowType{Scale: 9},\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),\n\t\t\tbuilder: array.NewStructBuilder(pool, timestampTzStructWithFraction),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tsb := b.(*array.StructBuilder)\n\t\t\t\tvalids = []bool{true, true, true}\n\t\t\t\tsb.AppendValues(valids)\n\t\t\t\tfor _, t := range vs.([]time.Time) {\n\t\t\t\t\tsb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())\n\t\t\t\t\tsb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))\n\t\t\t\t\tsb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0))\n\t\t\t\t}\n\t\t\t},\n\t\t\tcompare: func(src any, expected any, convertedRec arrow.Record) int {\n\t\t\t\tsrcvs := src.([]time.Time)\n\t\t\t\tfor i := 0; i < convertedRec.Column(0).Len(); i++ {\n\t\t\t\t\tts := ArrowSnowflakeTimestampToTime(convertedRec.Column(0), types.GetSnowflakeType(\"timestamp_tz\"), 9, i, nil)\n\t\t\t\t\tif !srcvs[i].Equal(*ts) {\n\t\t\t\t\t\treturn i\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn -1\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"array\",\n\t\t\tvalues: [][]string{{\"foo\", \"bar\"}, {\"baz\", \"quz\", \"quux\"}},\n\t\t\tnrows: 2,\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.StringType{}}}, nil),\n\t\t\tbuilder: array.NewStringBuilder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, a := range vs.([][]string) {\n\t\t\t\t\tb.(*array.StringBuilder).Append(fmt.Sprint(a))\n\t\t\t\t}\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tlogical: \"object\",\n\t\t\tvalues: []testObj{{0, \"foo\"}, {1, \"bar\"}},\n\t\t\tnrows: 2,\n\t\t\tsc: arrow.NewSchema([]arrow.Field{{Type: &arrow.StringType{}}}, nil),\n\t\t\tbuilder: array.NewStringBuilder(pool),\n\t\t\tappend: func(b array.Builder, vs any) {\n\t\t\t\tfor _, o := range vs.([]testObj) {\n\t\t\t\t\tb.(*array.StringBuilder).Append(fmt.Sprint(o))\n\t\t\t\t}\n\t\t\t},\n\t\t},\n\t} {\n\t\ttestName := tc.logical\n\t\tif tc.physical != \"\" {\n\t\t\ttestName += \" \" + tc.physical\n\t\t}\n\t\tt.Run(testName, func(t *testing.T) {\n\t\t\tscope := memory.NewCheckedAllocatorScope(pool)\n\t\t\tdefer scope.CheckSize(t)\n\n\t\t\tb := tc.builder\n\t\t\tdefer b.Release()\n\t\t\ttc.append(b, tc.values)\n\t\t\tarr := b.NewArray()\n\t\t\tdefer arr.Release()\n\t\t\trawRec := array.NewRecord(tc.sc, []arrow.Array{arr}, int64(tc.nrows))\n\t\t\tdefer rawRec.Release()\n\n\t\t\tmeta := tc.rowType\n\t\t\tmeta.Type = tc.logical\n\n\t\t\tctx := context.Background()\n\t\t\tswitch tc.arrowBatchesTimestampOption {\n\t\t\tcase ia.UseOriginalTimestamp:\n\t\t\t\tctx = ia.WithTimestampOption(ctx, ia.UseOriginalTimestamp)\n\t\t\tcase ia.UseSecondTimestamp:\n\t\t\t\tctx = ia.WithTimestampOption(ctx, ia.UseSecondTimestamp)\n\t\t\tcase ia.UseMillisecondTimestamp:\n\t\t\t\tctx = ia.WithTimestampOption(ctx, ia.UseMillisecondTimestamp)\n\t\t\tcase ia.UseMicrosecondTimestamp:\n\t\t\t\tctx = ia.WithTimestampOption(ctx, ia.UseMicrosecondTimestamp)\n\t\t\tdefault:\n\t\t\t\tctx = ia.WithTimestampOption(ctx, ia.UseNanosecondTimestamp)\n\t\t\t}\n\n\t\t\tif tc.enableArrowBatchesUtf8Validation {\n\t\t\t\tctx = ia.EnableUtf8Validation(ctx)\n\t\t\t}\n\n\t\t\tif tc.withHigherPrecision {\n\t\t\t\tctx = ia.WithHigherPrecision(ctx)\n\t\t\t}\n\n\t\t\ttransformedRec, err := arrowToRecord(ctx, rawRec, pool, []query.ExecResponseRowType{meta}, localTime.Location())\n\t\t\tif err != nil {\n\t\t\t\tif tc.error == \"\" || !strings.Contains(err.Error(), tc.error) {\n\t\t\t\t\tt.Fatalf(\"error: %s\", err)\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tdefer transformedRec.Release()\n\t\t\t\tif tc.error != \"\" {\n\t\t\t\t\tt.Fatalf(\"expected error: %s\", tc.error)\n\t\t\t\t}\n\n\t\t\t\tif tc.compare != nil {\n\t\t\t\t\tidx := tc.compare(tc.values, tc.expected, transformedRec)\n\t\t\t\t\tif idx != -1 {\n\t\t\t\t\t\tt.Fatalf(\"error: column array value mismatch at index %v\", idx)\n\t\t\t\t\t}\n\t\t\t\t} else {\n\t\t\t\t\tfor i, c := range transformedRec.Columns() {\n\t\t\t\t\t\trawCol := rawRec.Column(i)\n\t\t\t\t\t\tif rawCol != c {\n\t\t\t\t\t\t\tt.Fatalf(\"error: expected column %s, got column %s\", rawCol, c)\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t}\n}\n"
},
{
"path": "arrowbatches/schema.go",
"content": "package arrowbatches\n\nimport (\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/query\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/types\"\n\t\"time\"\n\n\tia \"github.com/snowflakedb/gosnowflake/v2/internal/arrow\"\n\n\t\"github.com/apache/arrow-go/v18/arrow\"\n)\n\nfunc recordToSchema(sc *arrow.Schema, rowType []query.ExecResponseRowType, loc *time.Location, timestampOption ia.TimestampOption, withHigherPrecision bool) (*arrow.Schema, error) {\n\tfields := recordToSchemaRecursive(sc.Fields(), rowType, loc, timestampOption, withHigherPrecision)\n\tmeta := sc.Metadata()\n\treturn arrow.NewSchema(fields, &meta), nil\n}\n\nfunc recordToSchemaRecursive(inFields []arrow.Field, rowType []query.ExecResponseRowType, loc *time.Location, timestampOption ia.TimestampOption, withHigherPrecision bool) []arrow.Field {\n\tvar outFields []arrow.Field\n\tfor i, f := range inFields {\n\t\tfieldMetadata := rowType[i].ToFieldMetadata()\n\t\tconverted, t := recordToSchemaSingleField(fieldMetadata, f, withHigherPrecision, timestampOption, loc)\n\n\t\tnewField := f\n\t\tif converted {\n\t\t\tnewField = arrow.Field{\n\t\t\t\tName: f.Name,\n\t\t\t\tType: t,\n\t\t\t\tNullable: f.Nullable,\n\t\t\t\tMetadata: f.Metadata,\n\t\t\t}\n\t\t}\n\t\toutFields = append(outFields, newField)\n\t}\n\treturn outFields\n}\n\nfunc recordToSchemaSingleField(fieldMetadata query.FieldMetadata, f arrow.Field, withHigherPrecision bool, timestampOption ia.TimestampOption, loc *time.Location) (bool, arrow.DataType) {\n\tt := f.Type\n\tconverted := true\n\tswitch types.GetSnowflakeType(fieldMetadata.Type) {\n\tcase types.FixedType:\n\t\tswitch f.Type.ID() {\n\t\tcase arrow.DECIMAL:\n\t\t\tif withHigherPrecision {\n\t\t\t\tconverted = false\n\t\t\t} else if fieldMetadata.Scale == 0 {\n\t\t\t\tt = &arrow.Int64Type{}\n\t\t\t} else {\n\t\t\t\tt = &arrow.Float64Type{}\n\t\t\t}\n\t\tdefault:\n\t\t\tif withHigherPrecision {\n\t\t\t\tconverted = false\n\t\t\t} else if fieldMetadata.Scale != 0 {\n\t\t\t\tt = &arrow.Float64Type{}\n\t\t\t} else {\n\t\t\t\tconverted = false\n\t\t\t}\n\t\t}\n\tcase types.TimeType:\n\t\tt = &arrow.Time64Type{Unit: arrow.Nanosecond}\n\tcase types.TimestampNtzType, types.TimestampTzType:\n\t\tswitch timestampOption {\n\t\tcase ia.UseOriginalTimestamp:\n\t\t\tconverted = false\n\t\tcase ia.UseMicrosecondTimestamp:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Microsecond}\n\t\tcase ia.UseMillisecondTimestamp:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Millisecond}\n\t\tcase ia.UseSecondTimestamp:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Second}\n\t\tdefault:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Nanosecond}\n\t\t}\n\tcase types.TimestampLtzType:\n\t\tswitch timestampOption {\n\t\tcase ia.UseOriginalTimestamp:\n\t\t\tconverted = false\n\t\tcase ia.UseMicrosecondTimestamp:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: loc.String()}\n\t\tcase ia.UseMillisecondTimestamp:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Millisecond, TimeZone: loc.String()}\n\t\tcase ia.UseSecondTimestamp:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Second, TimeZone: loc.String()}\n\t\tdefault:\n\t\t\tt = &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: loc.String()}\n\t\t}\n\tcase types.ObjectType:\n\t\tconverted = false\n\t\tif f.Type.ID() == arrow.STRUCT {\n\t\t\tvar internalFields []arrow.Field\n\t\t\tfor idx, internalField := range f.Type.(*arrow.StructType).Fields() {\n\t\t\t\tinternalConverted, convertedDataType := recordToSchemaSingleField(fieldMetadata.Fields[idx], internalField, withHigherPrecision, timestampOption, loc)\n\t\t\t\tconverted = converted || internalConverted\n\t\t\t\tif internalConverted {\n\t\t\t\t\tnewInternalField := arrow.Field{\n\t\t\t\t\t\tName: internalField.Name,\n\t\t\t\t\t\tType: convertedDataType,\n\t\t\t\t\t\tMetadata: internalField.Metadata,\n\t\t\t\t\t\tNullable: internalField.Nullable,\n\t\t\t\t\t}\n\t\t\t\t\tinternalFields = append(internalFields, newInternalField)\n\t\t\t\t} else {\n\t\t\t\t\tinternalFields = append(internalFields, internalField)\n\t\t\t\t}\n\t\t\t}\n\t\t\tt = arrow.StructOf(internalFields...)\n\t\t}\n\tcase types.ArrayType:\n\t\tif _, ok := f.Type.(*arrow.ListType); ok {\n\t\t\tconverted, dataType := recordToSchemaSingleField(fieldMetadata.Fields[0], f.Type.(*arrow.ListType).ElemField(), withHigherPrecision, timestampOption, loc)\n\t\t\tif converted {\n\t\t\t\tt = arrow.ListOf(dataType)\n\t\t\t}\n\t\t} else {\n\t\t\tt = f.Type\n\t\t}\n\tcase types.MapType:\n\t\tconvertedKey, keyDataType := recordToSchemaSingleField(fieldMetadata.Fields[0], f.Type.(*arrow.MapType).KeyField(), withHigherPrecision, timestampOption, loc)\n\t\tconvertedValue, valueDataType := recordToSchemaSingleField(fieldMetadata.Fields[1], f.Type.(*arrow.MapType).ItemField(), withHigherPrecision, timestampOption, loc)\n\t\tconverted = convertedKey || convertedValue\n\t\tif converted {\n\t\t\tt = arrow.MapOf(keyDataType, valueDataType)\n\t\t}\n\tdefault:\n\t\tconverted = false\n\t}\n\treturn converted, t\n}\n"
},
{
"path": "assert_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"bytes\"\n\t\"errors\"\n\t\"fmt\"\n\t\"math\"\n\t\"reflect\"\n\t\"regexp\"\n\t\"slices\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc assertNilE(t *testing.T, actual any, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateNil(actual, descriptions...))\n}\n\nfunc assertNilF(t *testing.T, actual any, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateNil(actual, descriptions...))\n}\n\nfunc assertNotNilE(t *testing.T, actual any, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateNotNil(actual, descriptions...))\n}\n\nfunc assertNotNilF(t *testing.T, actual any, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateNotNil(actual, descriptions...))\n}\n\nfunc assertErrIsF(t *testing.T, actual, expected error, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateErrIs(actual, expected, descriptions...))\n}\n\nfunc assertErrIsE(t *testing.T, actual, expected error, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateErrIs(actual, expected, descriptions...))\n}\n\nfunc assertErrorsAsF(t *testing.T, err error, target any, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateErrorsAs(err, target, descriptions...))\n}\n\nfunc assertEqualE(t *testing.T, actual any, expected any, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateEqual(actual, expected, descriptions...))\n}\n\nfunc assertEqualF(t *testing.T, actual any, expected any, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateEqual(actual, expected, descriptions...))\n}\n\nfunc assertEqualIgnoringWhitespaceE(t *testing.T, actual string, expected string, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateEqualIgnoringWhitespace(actual, expected, descriptions...))\n}\n\nfunc assertEqualEpsilonE(t *testing.T, actual, expected, epsilon float64, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateEqualEpsilon(actual, expected, epsilon, descriptions...))\n}\n\nfunc assertDeepEqualE(t *testing.T, actual any, expected any, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateDeepEqual(actual, expected, descriptions...))\n}\n\nfunc assertNotEqualF(t *testing.T, actual any, expected any, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateNotEqual(actual, expected, descriptions...))\n}\n\nfunc assertNotEqualE(t *testing.T, actual any, expected any, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateNotEqual(actual, expected, descriptions...))\n}\n\nfunc assertBytesEqualE(t *testing.T, actual []byte, expected []byte, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateBytesEqual(actual, expected, descriptions...))\n}\n\nfunc assertTrueF(t *testing.T, actual bool, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateEqual(actual, true, descriptions...))\n}\n\nfunc assertTrueE(t *testing.T, actual bool, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateEqual(actual, true, descriptions...))\n}\n\nfunc assertFalseF(t *testing.T, actual bool, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateEqual(actual, false, descriptions...))\n}\n\nfunc assertFalseE(t *testing.T, actual bool, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateEqual(actual, false, descriptions...))\n}\n\nfunc assertStringContainsE(t *testing.T, actual string, expectedToContain string, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateStringContains(actual, expectedToContain, descriptions...))\n}\n\nfunc assertStringContainsF(t *testing.T, actual string, expectedToContain string, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateStringContains(actual, expectedToContain, descriptions...))\n}\n\nfunc assertEmptyStringE(t *testing.T, actual string, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateEmptyString(actual, descriptions...))\n}\n\nfunc assertHasPrefixF(t *testing.T, actual string, expectedPrefix string, descriptions ...string) {\n\tt.Helper()\n\tfatalOnNonEmpty(t, validateHasPrefix(actual, expectedPrefix, descriptions...))\n}\n\nfunc assertHasPrefixE(t *testing.T, actual string, expectedPrefix string, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateHasPrefix(actual, expectedPrefix, descriptions...))\n}\n\nfunc assertBetweenE(t *testing.T, value float64, min float64, max float64, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateValueBetween(value, min, max, descriptions...))\n}\n\nfunc assertBetweenInclusiveE(t *testing.T, value float64, min float64, max float64, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateValueBetweenInclusive(value, min, max, descriptions...))\n}\n\nfunc assertEmptyE[T any](t *testing.T, actual []T, descriptions ...string) {\n\tt.Helper()\n\terrorOnNonEmpty(t, validateEmpty(actual, descriptions...))\n}\n\nfunc fatalOnNonEmpty(t *testing.T, errMsg string) {\n\tif errMsg != \"\" {\n\t\tt.Helper()\n\t\tt.Fatal(formatErrorMessage(errMsg))\n\t}\n}\n\nfunc errorOnNonEmpty(t *testing.T, errMsg string) {\n\tif errMsg != \"\" {\n\t\tt.Helper()\n\t\tt.Error(formatErrorMessage(errMsg))\n\t}\n}\n\nfunc formatErrorMessage(errMsg string) string {\n\treturn fmt.Sprintf(\"[%s] %s\", time.Now().Format(time.RFC3339Nano), maskSecrets(errMsg))\n}\n\nfunc validateNil(actual any, descriptions ...string) string {\n\tif isNil(actual) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to be nil but was not. %s\", maskSecrets(fmt.Sprintf(\"%v\", actual)), desc)\n}\n\nfunc validateNotNil(actual any, descriptions ...string) string {\n\tif !isNil(actual) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected to be not nil but was not. %s\", desc)\n}\n\nfunc validateErrIs(actual, expected error, descriptions ...string) string {\n\tif errors.Is(actual, expected) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\tactualStr := \"nil\"\n\texpectedStr := \"nil\"\n\tif actual != nil {\n\t\tactualStr = maskSecrets(actual.Error())\n\t}\n\tif expected != nil {\n\t\texpectedStr = maskSecrets(expected.Error())\n\t}\n\treturn fmt.Sprintf(\"expected %v to be %v. %s\", actualStr, expectedStr, desc)\n}\n\nfunc validateErrorsAs(err error, target any, descriptions ...string) string {\n\tif errors.As(err, target) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\terrStr := \"nil\"\n\tif err != nil {\n\t\terrStr = maskSecrets(err.Error())\n\t}\n\ttargetType := reflect.TypeOf(target)\n\treturn fmt.Sprintf(\"expected error %v to be assignable to %v but was not. %s\", errStr, targetType, desc)\n}\n\nfunc validateEqual(actual any, expected any, descriptions ...string) string {\n\tif expected == actual {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to be equal to \\\"%s\\\" but was not. %s\",\n\t\tmaskSecrets(fmt.Sprintf(\"%v\", actual)),\n\t\tmaskSecrets(fmt.Sprintf(\"%v\", expected)),\n\t\tdesc)\n}\n\nfunc removeWhitespaces(s string) string {\n\tpattern, err := regexp.Compile(`\\s+`)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\treturn pattern.ReplaceAllString(s, \"\")\n}\n\nfunc validateEqualIgnoringWhitespace(actual string, expected string, descriptions ...string) string {\n\tif removeWhitespaces(expected) == removeWhitespaces(actual) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to be equal to \\\"%s\\\" but was not. %s\",\n\t\tmaskSecrets(actual),\n\t\tmaskSecrets(expected),\n\t\tdesc)\n}\n\nfunc validateEqualEpsilon(actual, expected, epsilon float64, descriptions ...string) string {\n\tif math.Abs(actual-expected) < epsilon {\n\t\treturn \"\"\n\t}\n\treturn fmt.Sprintf(\"expected \\\"%f\\\" to be equal to \\\"%f\\\" within epsilon \\\"%f\\\" but was not. %s\", actual, expected, epsilon, joinDescriptions(descriptions...))\n}\n\nfunc validateDeepEqual(actual any, expected any, descriptions ...string) string {\n\tif reflect.DeepEqual(actual, expected) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to be equal to \\\"%s\\\" but was not. %s\",\n\t\tmaskSecrets(fmt.Sprintf(\"%v\", actual)),\n\t\tmaskSecrets(fmt.Sprintf(\"%v\", expected)),\n\t\tdesc)\n}\n\nfunc validateNotEqual(actual any, expected any, descriptions ...string) string {\n\tif expected != actual {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" not to be equal to \\\"%s\\\" but they were the same. %s\",\n\t\tmaskSecrets(fmt.Sprintf(\"%v\", actual)),\n\t\tmaskSecrets(fmt.Sprintf(\"%v\", expected)),\n\t\tdesc)\n}\n\nfunc validateBytesEqual(actual []byte, expected []byte, descriptions ...string) string {\n\tif bytes.Equal(actual, expected) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to be equal to \\\"%s\\\" but was not. %s\",\n\t\tmaskSecrets(string(actual)),\n\t\tmaskSecrets(string(expected)),\n\t\tdesc)\n}\n\nfunc validateStringContains(actual string, expectedToContain string, descriptions ...string) string {\n\tif strings.Contains(actual, expectedToContain) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to contain \\\"%s\\\" but did not. %s\",\n\t\tmaskSecrets(actual),\n\t\tmaskSecrets(expectedToContain),\n\t\tdesc)\n}\n\nfunc validateEmptyString(actual string, descriptions ...string) string {\n\tif actual == \"\" {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to be empty, but was not. %s\", maskSecrets(actual), desc)\n}\n\nfunc validateHasPrefix(actual string, expectedPrefix string, descriptions ...string) string {\n\tif strings.HasPrefix(actual, expectedPrefix) {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" to start with \\\"%s\\\" but did not. %s\",\n\t\tmaskSecrets(actual),\n\t\tmaskSecrets(expectedPrefix),\n\t\tdesc)\n}\n\nfunc validateValueBetween(value float64, min float64, max float64, descriptions ...string) string {\n\tif value > min && value < max {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" should be between \\\"%s\\\" and \\\"%s\\\" but did not. %s\",\n\t\tfmt.Sprintf(\"%f\", value),\n\t\tfmt.Sprintf(\"%f\", min),\n\t\tfmt.Sprintf(\"%f\", max),\n\t\tdesc)\n}\n\nfunc validateValueBetweenInclusive(value float64, min float64, max float64, descriptions ...string) string {\n\tif value >= min && value <= max {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%s\\\" should be between \\\"%s\\\" and \\\"%s\\\" inclusively but did not. %s\",\n\t\tfmt.Sprintf(\"%f\", value),\n\t\tfmt.Sprintf(\"%f\", min),\n\t\tfmt.Sprintf(\"%f\", max),\n\t\tdesc)\n}\n\nfunc validateEmpty[T any](value []T, descriptions ...string) string {\n\tif len(value) == 0 {\n\t\treturn \"\"\n\t}\n\tdesc := joinDescriptions(descriptions...)\n\treturn fmt.Sprintf(\"expected \\\"%v\\\" to be empty. %s\", maskSecrets(fmt.Sprintf(\"%v\", value)), desc)\n}\n\nfunc joinDescriptions(descriptions ...string) string {\n\treturn strings.Join(descriptions, \" \")\n}\n\nfunc isNil(value any) bool {\n\tif value == nil {\n\t\treturn true\n\t}\n\tval := reflect.ValueOf(value)\n\treturn slices.Contains([]reflect.Kind{reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface, reflect.Func}, val.Kind()) && val.IsNil()\n}\n"
},
{
"path": "async.go",
"content": "package gosnowflake\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"net/url\"\n\t\"strconv\"\n\t\"time\"\n)\n\nfunc (sr *snowflakeRestful) processAsync(\n\tctx context.Context,\n\trespd *execResponse,\n\theaders map[string]string,\n\ttimeout time.Duration,\n\tcfg *Config) (*execResponse, error) {\n\t// placeholder object to return to user while retrieving results\n\trows := new(snowflakeRows)\n\tres := new(snowflakeResult)\n\tswitch resType := getResultType(ctx); resType {\n\tcase execResultType:\n\t\tres.queryID = respd.Data.QueryID\n\t\tres.status = QueryStatusInProgress\n\t\tres.errChannel = make(chan error)\n\t\trespd.Data.AsyncResult = res\n\tcase queryResultType:\n\t\trows.queryID = respd.Data.QueryID\n\t\trows.status = QueryStatusInProgress\n\t\trows.errChannel = make(chan error)\n\t\trows.ctx = ctx\n\t\trespd.Data.AsyncRows = rows\n\tdefault:\n\t\treturn respd, nil\n\t}\n\n\t// spawn goroutine to retrieve asynchronous results\n\tgo GoroutineWrapper(\n\t\tctx,\n\t\tfunc() {\n\t\t\terr := sr.getAsync(ctx, headers, sr.getFullURL(respd.Data.GetResultURL, nil), timeout, res, rows, cfg)\n\t\t\tif err != nil {\n\t\t\t\tlogger.WithContext(ctx).Errorf(\"error while calling getAsync. %v\", err)\n\t\t\t}\n\t\t},\n\t)\n\treturn respd, nil\n}\n\nfunc (sr *snowflakeRestful) getAsync(\n\tctx context.Context,\n\theaders map[string]string,\n\tURL *url.URL,\n\ttimeout time.Duration,\n\tres *snowflakeResult,\n\trows *snowflakeRows,\n\tcfg *Config) error {\n\tresType := getResultType(ctx)\n\tvar errChannel chan error\n\tsfError := &SnowflakeError{\n\t\tNumber: ErrAsync,\n\t}\n\tif resType == execResultType {\n\t\terrChannel = res.errChannel\n\t\tsfError.QueryID = res.queryID\n\t} else {\n\t\terrChannel = rows.errChannel\n\t\tsfError.QueryID = rows.queryID\n\t}\n\tdefer close(errChannel)\n\ttoken, _, _ := sr.TokenAccessor.GetTokens()\n\theaders[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)\n\n\trespd, err := getQueryResultWithRetriesForAsyncMode(ctx, sr, URL, headers, timeout)\n\tif err != nil {\n\t\tlogger.WithContext(ctx).Errorf(\"error: %v\", err)\n\t\tsfError.Message = err.Error()\n\t\terrChannel <- sfError\n\t\treturn err\n\t}\n\n\tsc := &snowflakeConn{rest: sr, cfg: cfg, currentTimeProvider: defaultTimeProvider}\n\tif respd.Success {\n\t\tif resType == execResultType {\n\t\t\tres.insertID = -1\n\t\t\tif isDml(respd.Data.StatementTypeID) {\n\t\t\t\tres.affectedRows, err = updateRows(respd.Data)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn err\n\t\t\t\t}\n\t\t\t} else if isMultiStmt(&respd.Data) {\n\t\t\t\tr, err := sc.handleMultiExec(ctx, respd.Data)\n\t\t\t\tif err != nil {\n\t\t\t\t\tres.errChannel <- err\n\t\t\t\t\treturn err\n\t\t\t\t}\n\t\t\t\tres.affectedRows, err = r.RowsAffected()\n\t\t\t\tif err != nil {\n\t\t\t\t\tres.errChannel <- err\n\t\t\t\t\treturn err\n\t\t\t\t}\n\t\t\t}\n\t\t\tres.queryID = respd.Data.QueryID\n\t\t\tres.errChannel <- nil // mark exec status complete\n\t\t} else {\n\t\t\trows.sc = sc\n\t\t\trows.queryID = respd.Data.QueryID\n\t\t\tif isMultiStmt(&respd.Data) {\n\t\t\t\tif err = sc.handleMultiQuery(ctx, respd.Data, rows); err != nil {\n\t\t\t\t\trows.errChannel <- err\n\t\t\t\t\treturn err\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\trows.addDownloader(populateChunkDownloader(ctx, sc, respd.Data))\n\t\t\t}\n\t\t\tif err = rows.ChunkDownloader.start(); err != nil {\n\t\t\t\trows.errChannel <- err\n\t\t\t\treturn err\n\t\t\t}\n\t\t\trows.errChannel <- nil // mark query status complete\n\t\t}\n\t} else {\n\t\tvar code int\n\t\tif respd.Code != \"\" {\n\t\t\tcode, err = strconv.Atoi(respd.Code)\n\t\t\tif err != nil {\n\t\t\t\tcode = -1\n\t\t\t}\n\t\t} else {\n\t\t\tcode = -1\n\t\t}\n\t\terrChannel <- &SnowflakeError{\n\t\t\tNumber: code,\n\t\t\tSQLState: respd.Data.SQLState,\n\t\t\tMessage: respd.Message,\n\t\t\tQueryID: respd.Data.QueryID,\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc getQueryResultWithRetriesForAsyncMode(\n\tctx context.Context,\n\tsr *snowflakeRestful,\n\tURL *url.URL,\n\theaders map[string]string,\n\ttimeout time.Duration) (respd *execResponse, err error) {\n\tretry := 0\n\tretryPattern := []int32{1, 1, 2, 3, 4, 8, 10}\n\tretryPatternIndex := 0\n\tretryCountForSessionRenewal := 0\n\n\tfor {\n\t\tlogger.WithContext(ctx).Debugf(\"Retry count for get query result request in async mode: %v\", retry)\n\n\t\trespd, err = getExecResponse(ctx, sr, URL, headers, timeout)\n\t\tif err != nil {\n\t\t\treturn respd, err\n\t\t}\n\t\tif respd.Code == sessionExpiredCode {\n\t\t\t// Update the session token in the header and retry\n\t\t\ttoken, _, _ := sr.TokenAccessor.GetTokens()\n\t\t\tif token != \"\" && headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, token) {\n\t\t\t\theaders[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)\n\t\t\t\tlogger.WithContext(ctx).Debug(\"Session token has been updated.\")\n\t\t\t\tretry++\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\t// Renew the session token\n\t\t\tif err = sr.renewExpiredSessionToken(ctx, timeout, token); err != nil {\n\t\t\t\tlogger.WithContext(ctx).Errorf(\"failed to renew session token. err: %v\", err)\n\t\t\t\treturn respd, err\n\t\t\t}\n\t\t\tretryCountForSessionRenewal++\n\n\t\t\t// If this is the first response, go back to retry the query\n\t\t\t// since it failed due to session expiration\n\t\t\tlogger.WithContext(ctx).Debugf(\"retry count for session renewal: %v\", retryCountForSessionRenewal)\n\t\t\tif retryCountForSessionRenewal < 2 {\n\t\t\t\tretry++\n\t\t\t\tcontinue\n\t\t\t} else {\n\t\t\t\tlogger.WithContext(ctx).Errorf(\"failed to get query result with the renewed session token. err: %v\", err)\n\t\t\t\treturn respd, err\n\t\t\t}\n\t\t} else if respd.Code != queryInProgressAsyncCode {\n\t\t\t// If the query takes longer than 45 seconds to complete the results are not returned.\n\t\t\t// If the query is still in progress after 45 seconds, retry the request to the /results endpoint.\n\t\t\t// For all other scenarios continue processing results response\n\t\t\tbreak\n\t\t} else {\n\t\t\t// Sleep before retrying get result request. Exponential backoff up to 5 seconds.\n\t\t\t// Once 5 second backoff is reached it will keep retrying with this sleeptime.\n\t\t\tsleepTime := time.Millisecond * time.Duration(500*retryPattern[retryPatternIndex])\n\t\t\tlogger.WithContext(ctx).Debugf(\"Query execution still in progress. Response code: %v, message: %v Sleep for %v ms\", respd.Code, respd.Message, sleepTime)\n\t\t\ttime.Sleep(sleepTime)\n\t\t\tretry++\n\n\t\t\tif retryPatternIndex < len(retryPattern)-1 {\n\t\t\t\tretryPatternIndex++\n\t\t\t}\n\t\t}\n\t}\n\tif len(respd.Data.RowType) > 0 {\n\t\tlogger.Infof(\"[Server Response Validation]: RowType: %s, QueryResultFormat: %s\", respd.Data.RowType[0].Name, respd.Data.QueryResultFormat)\n\t}\n\treturn respd, nil\n}\n"
},
{
"path": "async_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"context\"\n\t\"database/sql\"\n\t\"fmt\"\n\t\"testing\"\n)\n\nfunc TestAsyncMode(t *testing.T) {\n\tctx := WithAsyncMode(context.Background())\n\tnumrows := 100000\n\tcnt := 0\n\tvar idx int\n\tvar v string\n\n\trunDBTest(t, func(dbt *DBTest) {\n\t\trows := dbt.mustQueryContext(ctx, fmt.Sprintf(selectRandomGenerator, numrows))\n\t\tdefer rows.Close()\n\n\t\t// Next() will block and wait until results are available\n\t\tfor rows.Next() {\n\t\t\tif err := rows.Scan(&idx, &v); err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\tcnt++\n\t\t}\n\t\tlogger.Infof(\"NextResultSet: %v\", rows.NextResultSet())\n\n\t\tif cnt != numrows {\n\t\t\tt.Errorf(\"number of rows didn't match. expected: %v, got: %v\", numrows, cnt)\n\t\t}\n\n\t\tdbt.mustExec(\"create or replace table test_async_exec (value boolean)\")\n\t\tres := dbt.mustExecContext(ctx, \"insert into test_async_exec values (true)\")\n\t\tcount, err := res.RowsAffected()\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"res.RowsAffected() returned error: %v\", err)\n\t\t}\n\t\tif count != 1 {\n\t\t\tt.Fatalf(\"expected 1 affected row, got %d\", count)\n\t\t}\n\t})\n}\n\nfunc TestAsyncModePing(t *testing.T) {\n\tctx := WithAsyncMode(context.Background())\n\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tdefer func() {\n\t\t\tif r := recover(); r != nil {\n\t\t\t\tt.Fatalf(\"panic during ping: %v\", r)\n\t\t\t}\n\t\t}()\n\t\terr := dbt.conn.PingContext(ctx)\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t})\n}\n\nfunc TestAsyncModeMultiStatement(t *testing.T) {\n\twithMultiStmtCtx := WithMultiStatement(context.Background(), 6)\n\tctx := WithAsyncMode(withMultiStmtCtx)\n\tmultiStmtQuery := \"begin;\\n\" +\n\t\t\"delete from test_multi_statement_async;\\n\" +\n\t\t\"insert into test_multi_statement_async values (1, 'a'), (2, 'b');\\n\" +\n\t\t\"select 1;\\n\" +\n\t\t\"select 2;\\n\" +\n\t\t\"rollback;\"\n\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tdbt.mustExec(\"drop table if exists test_multi_statement_async\")\n\t\tdbt.mustExec(`create or replace table test_multi_statement_async(\n\t\t\tc1 number, c2 string) as select 10, 'z'`)\n\t\tdefer dbt.mustExec(\"drop table if exists test_multi_statement_async\")\n\n\t\tres := dbt.mustExecContext(ctx, multiStmtQuery)\n\t\tcount, err := res.RowsAffected()\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"res.RowsAffected() returned error: %v\", err)\n\t\t}\n\t\tif count != 3 {\n\t\t\tt.Fatalf(\"expected 3 affected rows, got %d\", count)\n\t\t}\n\t})\n}\n\nfunc TestAsyncModeCancel(t *testing.T) {\n\twithCancelCtx, cancel := context.WithCancel(context.Background())\n\tctx := WithAsyncMode(withCancelCtx)\n\tnumrows := 100000\n\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tdbt.mustQueryContext(ctx, fmt.Sprintf(selectRandomGenerator, numrows))\n\t\tcancel()\n\t})\n}\n\nfunc TestAsyncQueryFail(t *testing.T) {\n\tctx := WithAsyncMode(context.Background())\n\trunDBTest(t, func(dbt *DBTest) {\n\t\trows := dbt.mustQueryContext(ctx, \"selectt 1\")\n\t\tdefer rows.Close()\n\n\t\tif rows.Next() {\n\t\t\tt.Fatal(\"should have no rows available\")\n\t\t} else {\n\t\t\tif err := rows.Err(); err == nil {\n\t\t\t\tt.Fatal(\"should return a syntax error\")\n\t\t\t}\n\t\t}\n\t})\n}\n\n// TestMultipleAsyncQueries validates that shorter async queries return before\n// longer ones. The TIMELIMIT values (30 and 10) must have sufficient separation\n// to avoid flaky ordering. Do not reduce these values significantly.\nfunc TestMultipleAsyncQueries(t *testing.T) {\n\tctx := WithAsyncMode(context.Background())\n\ts1 := \"foo\"\n\ts2 := \"bar\"\n\tch1 := make(chan string)\n\tch2 := make(chan string)\n\n\tdb := openDB(t)\n\n\trunDBTest(t, func(dbt *DBTest) {\n\t\trows1, err := db.QueryContext(ctx, fmt.Sprintf(\"select distinct '%v' from table (generator(timelimit=>%v))\", s1, 30))\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"can't read rows1: %v\", err)\n\t\t}\n\t\tdefer rows1.Close()\n\t\trows2, err := db.QueryContext(ctx, fmt.Sprintf(\"select distinct '%v' from table (generator(timelimit=>%v))\", s2, 10))\n\t\tif err != nil {\n\t\t\tt.Fatalf(\"can't read rows2: %v\", err)\n\t\t}\n\t\tdefer rows2.Close()\n\n\t\tgo retrieveRows(rows1, ch1)\n\t\tgo retrieveRows(rows2, ch2)\n\t\tselect {\n\t\tcase res := <-ch1:\n\t\t\tt.Fatalf(\"value %v should not have been called earlier.\", res)\n\t\tcase res := <-ch2:\n\t\t\tif res != s2 {\n\t\t\t\tt.Fatalf(\"query failed. expected: %v, got: %v\", s2, res)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc retrieveRows(rows *sql.Rows, ch chan string) {\n\tvar s string\n\tfor rows.Next() {\n\t\tif err := rows.Scan(&s); err != nil {\n\t\t\tch <- err.Error()\n\t\t\tclose(ch)\n\t\t\treturn\n\t\t}\n\t}\n\tch <- s\n\tclose(ch)\n}\n\n// TestLongRunningAsyncQuery validates the retry logic for async queries that\n// exceed Snowflake's 45-second threshold. After 45 seconds, the /results\n// endpoint returns \"query in progress\" (code 333334) and the driver must retry.\n// The 50-second wait MUST exceed 45 seconds to exercise this code path.\nfunc TestLongRunningAsyncQuery(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tctx := WithMultiStatement(context.Background(), 0)\n\t\tquery := \"CALL SYSTEM$WAIT(50, 'SECONDS');use snowflake_sample_data\"\n\n\t\trows := dbt.mustQueryContext(WithAsyncMode(ctx), query)\n\t\tdefer rows.Close()\n\t\tvar v string\n\t\ti := 0\n\t\tfor {\n\t\t\tfor rows.Next() {\n\t\t\t\terr := rows.Scan(&v)\n\t\t\t\tif err != nil {\n\t\t\t\t\tt.Fatalf(\"failed to get result. err: %v\", err)\n\t\t\t\t}\n\t\t\t\tif v == \"\" {\n\t\t\t\t\tt.Fatal(\"should have returned a result\")\n\t\t\t\t}\n\t\t\t\tresults := []string{\"waited 50 seconds\", \"Statement executed successfully.\"}\n\t\t\t\tif v != results[i] {\n\t\t\t\t\tt.Fatalf(\"unexpected result returned. expected: %v, but got: %v\", results[i], v)\n\t\t\t\t}\n\t\t\t\ti++\n\t\t\t}\n\t\t\tif !rows.NextResultSet() {\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc TestLongRunningAsyncQueryFetchResultByID(t *testing.T) {\n\trunDBTest(t, func(dbt *DBTest) {\n\t\tqueryIDChan := make(chan string, 1)\n\t\tctx := WithAsyncMode(context.Background())\n\t\tctx = WithQueryIDChan(ctx, queryIDChan)\n\n\t\t// Run a long running query asynchronously\n\t\tgo dbt.mustExecContext(ctx, \"CALL SYSTEM$WAIT(50, 'SECONDS')\")\n\n\t\t// Get the query ID without waiting for the query to finish\n\t\tqueryID := <-queryIDChan\n\t\tassertNotNilF(t, queryID, \"expected a nonempty query ID\")\n\n\t\tctx = WithFetchResultByID(ctx, queryID)\n\t\trows := dbt.mustQueryContext(ctx, \"\")\n\t\tdefer rows.Close()\n\n\t\tvar v string\n\t\tassertTrueF(t, rows.Next())\n\t\terr := rows.Scan(&v)\n\t\tassertNilF(t, err, fmt.Sprintf(\"failed to get result. err: %v\", err))\n\t\tassertNotNilF(t, v, \"should have returned a result\")\n\n\t\texpected := \"waited 50 seconds\"\n\t\tif v != expected {\n\t\t\tt.Fatalf(\"unexpected result returned. expected: %v, but got: %v\", expected, v)\n\t\t}\n\t\tassertFalseF(t, rows.NextResultSet())\n\t})\n}\n"
},
{
"path": "auth.go",
"content": "package gosnowflake\n\nimport (\n\t\"context\"\n\t\"crypto/sha256\"\n\t\"crypto/x509\"\n\t\"encoding/base64\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"os\"\n\t\"runtime\"\n\t\"slices\"\n\t\"strconv\"\n\t\"strings\"\n\t\"time\"\n\n\tsferrors \"github.com/snowflakedb/gosnowflake/v2/internal/errors\"\n\n\t\"github.com/golang-jwt/jwt/v5\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/compilation\"\n\tsfconfig \"github.com/snowflakedb/gosnowflake/v2/internal/config\"\n\tinternalos \"github.com/snowflakedb/gosnowflake/v2/internal/os\"\n)\n\nconst (\n\tclientType = \"Go\"\n)\n\nconst (\n\tclientStoreTemporaryCredential = \"CLIENT_STORE_TEMPORARY_CREDENTIAL\"\n\tclientRequestMfaToken = \"CLIENT_REQUEST_MFA_TOKEN\"\n\tidTokenAuthenticator = \"ID_TOKEN\"\n)\n\n// AuthType indicates the type of authentication in Snowflake\ntype AuthType = sfconfig.AuthType\n\nconst (\n\t// AuthTypeSnowflake is the general username password authentication\n\tAuthTypeSnowflake = sfconfig.AuthTypeSnowflake\n\t// AuthTypeOAuth is the OAuth authentication\n\tAuthTypeOAuth = sfconfig.AuthTypeOAuth\n\t// AuthTypeExternalBrowser is to use a browser to access an Fed and perform SSO authentication\n\tAuthTypeExternalBrowser = sfconfig.AuthTypeExternalBrowser\n\t// AuthTypeOkta is to use a native okta URL to perform SSO authentication on Okta\n\tAuthTypeOkta = sfconfig.AuthTypeOkta\n\t// AuthTypeJwt is to use Jwt to perform authentication\n\tAuthTypeJwt = sfconfig.AuthTypeJwt\n\t// AuthTypeTokenAccessor is to use the provided token accessor and bypass authentication\n\tAuthTypeTokenAccessor = sfconfig.AuthTypeTokenAccessor\n\t// AuthTypeUsernamePasswordMFA is to use username and password with mfa\n\tAuthTypeUsernamePasswordMFA = sfconfig.AuthTypeUsernamePasswordMFA\n\t// AuthTypePat is to use programmatic access token\n\tAuthTypePat = sfconfig.AuthTypePat\n\t// AuthTypeOAuthAuthorizationCode is to use browser-based OAuth2 flow\n\tAuthTypeOAuthAuthorizationCode = sfconfig.AuthTypeOAuthAuthorizationCode\n\t// AuthTypeOAuthClientCredentials is to use non-interactive OAuth2 flow\n\tAuthTypeOAuthClientCredentials = sfconfig.AuthTypeOAuthClientCredentials\n\t// AuthTypeWorkloadIdentityFederation is to use CSP identity for authentication\n\tAuthTypeWorkloadIdentityFederation = sfconfig.AuthTypeWorkloadIdentityFederation\n)\n\nfunc isOauthNativeFlow(authType AuthType) bool {\n\treturn authType == AuthTypeOAuthAuthorizationCode || authType == AuthTypeOAuthClientCredentials\n}\n\nvar refreshOAuthTokenErrorCodes = []string{\n\tstrconv.Itoa(ErrMissingAccessATokenButRefreshTokenPresent),\n\tinvalidOAuthAccessTokenCode,\n\texpiredOAuthAccessTokenCode,\n}\n\n// userAgent shows up in User-Agent HTTP header\nvar userAgent = fmt.Sprintf(\"%v/%v (%v-%v) %v/%v\",\n\tclientType,\n\tSnowflakeGoDriverVersion,\n\truntime.GOOS,\n\truntime.GOARCH,\n\truntime.Compiler,\n\truntime.Version())\n\ntype authRequestClientEnvironment struct {\n\tApplication string `json:\"APPLICATION\"`\n\tApplicationPath string `json:\"APPLICATION_PATH\"`\n\tOs string `json:\"OS\"`\n\tOsVersion string `json:\"OS_VERSION\"`\n\tOsDetails map[string]string `json:\"OS_DETAILS,omitempty\"`\n\tIsa string `json:\"ISA,omitempty\"`\n\tOCSPMode string `json:\"OCSP_MODE\"`\n\tGoVersion string `json:\"GO_VERSION\"`\n\tOAuthType string `json:\"OAUTH_TYPE,omitempty\"`\n\tCertRevocationCheckMode string `json:\"CERT_REVOCATION_CHECK_MODE,omitempty\"`\n\tPlatform []string `json:\"PLATFORM,omitempty\"`\n\tCoreVersion string `json:\"CORE_VERSION,omitempty\"`\n\tCoreLoadError string `json:\"CORE_LOAD_ERROR,omitempty\"`\n\tCoreFileName string `json:\"CORE_FILE_NAME,omitempty\"`\n\tCgoEnabled bool `json:\"CGO_ENABLED,omitempty\"`\n\tLinkingMode string `json:\"LINKING_MODE,omitempty\"`\n\tLibcFamily string `json:\"LIBC_FAMILY,omitempty\"`\n\tLibcVersion string `json:\"LIBC_VERSION,omitempty\"`\n}\n\ntype authRequestData struct {\n\tClientAppID string `json:\"CLIENT_APP_ID\"`\n\tClientAppVersion string `json:\"CLIENT_APP_VERSION\"`\n\tSvnRevision string `json:\"SVN_REVISION\"`\n\tAccountName string `json:\"ACCOUNT_NAME\"`\n\tLoginName string `json:\"LOGIN_NAME,omitempty\"`\n\tPassword string `json:\"PASSWORD,omitempty\"`\n\tRawSAMLResponse string `json:\"RAW_SAML_RESPONSE,omitempty\"`\n\tExtAuthnDuoMethod string `json:\"EXT_AUTHN_DUO_METHOD,omitempty\"`\n\tPasscode string `json:\"PASSCODE,omitempty\"`\n\tAuthenticator string `json:\"AUTHENTICATOR,omitempty\"`\n\tSessionParameters map[string]any `json:\"SESSION_PARAMETERS,omitempty\"`\n\tClientEnvironment authRequestClientEnvironment `json:\"CLIENT_ENVIRONMENT\"`\n\tBrowserModeRedirectPort string `json:\"BROWSER_MODE_REDIRECT_PORT,omitempty\"`\n\tProofKey string `json:\"PROOF_KEY,omitempty\"`\n\tToken string `json:\"TOKEN,omitempty\"`\n\tProvider string `json:\"PROVIDER,omitempty\"`\n}\ntype authRequest struct {\n\tData authRequestData `json:\"data\"`\n}\n\ntype nameValueParameter struct {\n\tName string `json:\"name\"`\n\tValue any `json:\"value\"`\n}\n\ntype authResponseSessionInfo struct {\n\tDatabaseName string `json:\"databaseName\"`\n\tSchemaName string `json:\"schemaName\"`\n\tWarehouseName string `json:\"warehouseName\"`\n\tRoleName string `json:\"roleName\"`\n}\n\ntype authResponseMain struct {\n\tToken string `json:\"token,omitempty\"`\n\tValidity time.Duration `json:\"validityInSeconds,omitempty\"`\n\tMasterToken string `json:\"masterToken,omitempty\"`\n\tMasterValidity time.Duration `json:\"masterValidityInSeconds\"`\n\tMfaToken string `json:\"mfaToken,omitempty\"`\n\tMfaTokenValidity time.Duration `json:\"mfaTokenValidityInSeconds\"`\n\tIDToken string `json:\"idToken,omitempty\"`\n\tIDTokenValidity time.Duration `json:\"idTokenValidityInSeconds\"`\n\tDisplayUserName string `json:\"displayUserName\"`\n\tServerVersion string `json:\"serverVersion\"`\n\tFirstLogin bool `json:\"firstLogin\"`\n\tRemMeToken string `json:\"remMeToken\"`\n\tRemMeValidity time.Duration `json:\"remMeValidityInSeconds\"`\n\tHealthCheckInterval time.Duration `json:\"healthCheckInterval\"`\n\tNewClientForUpgrade string `json:\"newClientForUpgrade\"`\n\tSessionID int64 `json:\"sessionId\"`\n\tParameters []nameValueParameter `json:\"parameters\"`\n\tSessionInfo authResponseSessionInfo `json:\"sessionInfo\"`\n\tTokenURL string `json:\"tokenUrl,omitempty\"`\n\tSSOURL string `json:\"ssoUrl,omitempty\"`\n\tProofKey string `json:\"proofKey,omitempty\"`\n}\n\ntype authResponse struct {\n\tData authResponseMain `json:\"data\"`\n\tMessage string `json:\"message\"`\n\tCode string `json:\"code\"`\n\tSuccess bool `json:\"success\"`\n}\n\nfunc postAuth(\n\tctx context.Context,\n\tsr *snowflakeRestful,\n\tclient *http.Client,\n\tparams *url.Values,\n\theaders map[string]string,\n\tbodyCreator bodyCreatorType,\n\ttimeout time.Duration) (\n\tdata *authResponse, err error) {\n\tparams.Set(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())\n\tparams.Set(requestGUIDKey, NewUUID().String())\n\n\tfullURL := sr.getFullURL(loginRequestPath, params)\n\tlogger.WithContext(ctx).Infof(\"full URL: %v\", fullURL)\n\tresp, err := sr.FuncAuthPost(ctx, client, fullURL, headers, bodyCreator, timeout, sr.MaxRetryCount)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdefer func() {\n\t\tif closeErr := resp.Body.Close(); closeErr != nil {\n\t\t\tlogger.WithContext(ctx).Errorf(\"failed to close HTTP response body for %v. err: %v\", fullURL, closeErr)\n\t\t}\n\t}()\n\tif resp.StatusCode == http.StatusOK {\n\t\tvar respd authResponse\n\t\terr = json.NewDecoder(resp.Body).Decode(&respd)\n\t\tif err != nil {\n\t\t\tlogger.WithContext(ctx).Errorf(\"failed to decode JSON. err: %v\", err)\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &respd, nil\n\t}\n\tswitch resp.StatusCode {\n\tcase http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:\n\t\t// service availability or connectivity issue. Most likely server side issue.\n\t\treturn nil, &SnowflakeError{\n\t\t\tNumber: ErrCodeServiceUnavailable,\n\t\t\tSQLState: SQLStateConnectionWasNotEstablished,\n\t\t\tMessage: sferrors.ErrMsgServiceUnavailable,\n\t\t\tMessageArgs: []any{resp.StatusCode, fullURL},\n\t\t}\n\tcase http.StatusUnauthorized, http.StatusForbidden:\n\t\t// failed to connect to db. account name may be wrong\n\t\treturn nil, &SnowflakeError{\n\t\t\tNumber: ErrCodeFailedToConnect,\n\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\tMessage: sferrors.ErrMsgFailedToConnect,\n\t\t\tMessageArgs: []any{resp.StatusCode, fullURL},\n\t\t}\n\t}\n\tb, err := io.ReadAll(resp.Body)\n\tif err != nil {\n\t\tlogger.WithContext(ctx).Errorf(\"failed to extract HTTP response body. err: %v\", err)\n\t\treturn nil, err\n\t}\n\tlogger.WithContext(ctx).Infof(\"HTTP: %v, URL: %v, Body: %v\", resp.StatusCode, fullURL, b)\n\tlogger.WithContext(ctx).Infof(\"Header: %v\", resp.Header)\n\treturn nil, &SnowflakeError{\n\t\tNumber: ErrFailedToAuth,\n\t\tSQLState: SQLStateConnectionRejected,\n\t\tMessage: sferrors.ErrMsgFailedToAuth,\n\t\tMessageArgs: []any{resp.StatusCode, fullURL},\n\t}\n}\n\n// Generates a map of headers needed to authenticate\n// with Snowflake.\nfunc getHeaders() map[string]string {\n\theaders := make(map[string]string)\n\theaders[httpHeaderContentType] = headerContentTypeApplicationJSON\n\theaders[httpHeaderAccept] = headerAcceptTypeApplicationSnowflake\n\theaders[httpClientAppID] = clientType\n\theaders[httpClientAppVersion] = SnowflakeGoDriverVersion\n\theaders[httpHeaderUserAgent] = userAgent\n\treturn headers\n}\n\n// Used to authenticate the user with Snowflake.\nfunc authenticate(\n\tctx context.Context,\n\tsc *snowflakeConn,\n\tsamlResponse []byte,\n\tproofKey []byte,\n) (resp *authResponseMain, err error) {\n\tif sc.cfg.Authenticator == AuthTypeTokenAccessor {\n\t\tlogger.WithContext(ctx).Info(\"Bypass authentication using existing token from token accessor\")\n\t\tsessionInfo := authResponseSessionInfo{\n\t\t\tDatabaseName: sc.cfg.Database,\n\t\t\tSchemaName: sc.cfg.Schema,\n\t\t\tWarehouseName: sc.cfg.Warehouse,\n\t\t\tRoleName: sc.cfg.Role,\n\t\t}\n\t\ttoken, masterToken, sessionID := sc.cfg.TokenAccessor.GetTokens()\n\t\treturn &authResponseMain{\n\t\t\tToken: token,\n\t\t\tMasterToken: masterToken,\n\t\t\tSessionID: sessionID,\n\t\t\tSessionInfo: sessionInfo,\n\t\t}, nil\n\t}\n\n\theaders := getHeaders()\n\t// Get the current application path\n\tapplicationPath, err := os.Executable()\n\tif err != nil {\n\t\tlogger.WithContext(ctx).Warnf(\"Failed to get executable path: %v\", err)\n\t\tapplicationPath = \"unknown\"\n\t}\n\n\toauthType := \"\"\n\tswitch sc.cfg.Authenticator {\n\tcase AuthTypeOAuthAuthorizationCode:\n\t\toauthType = \"OAUTH_AUTHORIZATION_CODE\"\n\tcase AuthTypeOAuthClientCredentials:\n\t\toauthType = \"OAUTH_CLIENT_CREDENTIALS\"\n\t}\n\n\tclientEnvironment := newAuthRequestClientEnvironment()\n\tclientEnvironment.Application = sc.cfg.Application\n\tclientEnvironment.ApplicationPath = applicationPath\n\tclientEnvironment.OAuthType = oauthType\n\tclientEnvironment.CertRevocationCheckMode = sc.cfg.CertRevocationCheckMode.String()\n\tclientEnvironment.Platform = getDetectedPlatforms()\n\n\tsessionParameters := make(map[string]any)\n\tfor k, v := range sc.syncParams.All() {\n\t\t// upper casing to normalize keys\n\t\tsessionParameters[strings.ToUpper(k)] = v\n\t}\n\n\tsessionParameters[sessionClientValidateDefaultParameters] = sc.cfg.ValidateDefaultParameters != ConfigBoolFalse\n\tif sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {\n\t\tsessionParameters[clientRequestMfaToken] = true\n\t}\n\tif sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {\n\t\tsessionParameters[clientStoreTemporaryCredential] = true\n\t}\n\tbodyCreator := func() ([]byte, error) {\n\t\treturn createRequestBody(sc, sessionParameters, clientEnvironment, proofKey, samlResponse)\n\t}\n\n\tparams := &url.Values{}\n\tif sc.cfg.Database != \"\" {\n\t\tparams.Add(\"databaseName\", sc.cfg.Database)\n\t}\n\tif sc.cfg.Schema != \"\" {\n\t\tparams.Add(\"schemaName\", sc.cfg.Schema)\n\t}\n\tif sc.cfg.Warehouse != \"\" {\n\t\tparams.Add(\"warehouse\", sc.cfg.Warehouse)\n\t}\n\tif sc.cfg.Role != \"\" {\n\t\tparams.Add(\"roleName\", sc.cfg.Role)\n\t}\n\n\tlogger.WithContext(ctx).Infof(\"Information for Auth: Host: %v, User: %v, Authenticator: %v, Params: %v, Protocol: %v, Port: %v, LoginTimeout: %v\",\n\t\tsc.rest.Host, sc.cfg.User, sc.cfg.Authenticator.String(), params, sc.rest.Protocol, sc.rest.Port, sc.rest.LoginTimeout)\n\n\trespd, err := sc.rest.FuncPostAuth(ctx, sc.rest, sc.rest.getClientFor(sc.cfg.Authenticator), params, headers, bodyCreator, sc.rest.LoginTimeout)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif !respd.Success {\n\t\tlogger.WithContext(ctx).Error(\"Authentication FAILED\")\n\t\tsc.rest.TokenAccessor.SetTokens(\"\", \"\", -1)\n\t\tif sessionParameters[clientRequestMfaToken] == true {\n\t\t\tcredentialsStorage.deleteCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))\n\t\t}\n\t\tif sessionParameters[clientStoreTemporaryCredential] == true && sc.cfg.Authenticator == AuthTypeExternalBrowser {\n\t\t\tcredentialsStorage.deleteCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))\n\t\t}\n\t\tif sessionParameters[clientStoreTemporaryCredential] == true && isOauthNativeFlow(sc.cfg.Authenticator) {\n\t\t\tcredentialsStorage.deleteCredential(newOAuthAccessTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))\n\t\t}\n\t\tcode, err := strconv.Atoi(respd.Code)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn nil, exceptionTelemetry(&SnowflakeError{\n\t\t\tNumber: code,\n\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\tMessage: respd.Message,\n\t\t}, sc)\n\t}\n\tlogger.WithContext(ctx).Info(\"Authentication SUCCESS\")\n\tsc.rest.TokenAccessor.SetTokens(respd.Data.Token, respd.Data.MasterToken, respd.Data.SessionID)\n\tif sessionParameters[clientRequestMfaToken] == true {\n\t\ttoken := respd.Data.MfaToken\n\t\tcredentialsStorage.setCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User), token)\n\t}\n\tif sessionParameters[clientStoreTemporaryCredential] == true {\n\t\ttoken := respd.Data.IDToken\n\t\tcredentialsStorage.setCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User), token)\n\t}\n\treturn &respd.Data, nil\n}\n\nfunc newAuthRequestClientEnvironment() authRequestClientEnvironment {\n\tvar coreVersion string\n\tvar coreLoadError string\n\n\t// Try to get minicore version, but don't block if it's not loaded yet\n\tif !compilation.MinicoreEnabled {\n\t\tlogger.Trace(\"minicore disabled at compile time\")\n\t\tcoreLoadError = \"Minicore is disabled at compile time (built with -tags minicore_disabled)\"\n\t} else if strings.EqualFold(os.Getenv(disableMinicoreEnv), \"true\") {\n\t\tlogger.Trace(\"minicore loading disabled\")\n\t\tcoreLoadError = \"Minicore is disabled with SF_DISABLE_MINICORE env variable\"\n\t} else if mc := getMiniCore(); mc != nil {\n\t\tvar err error\n\t\tcoreVersion, err = mc.FullVersion()\n\t\tif err != nil {\n\t\t\tlogger.Debugf(\"Minicore loading failed. %v\", err)\n\t\t\tvar mcErr *miniCoreError\n\t\t\tif errors.As(err, &mcErr) {\n\t\t\t\tcoreLoadError = fmt.Sprintf(\"Failed to load binary: %v\", mcErr.errorType)\n\t\t\t} else {\n\t\t\t\tcoreLoadError = \"Failed to load binary: unknown\"\n\t\t\t}\n\t\t}\n\t} else {\n\t\t// Minicore not loaded yet - this is expected during startup\n\t\tcoreVersion = \"\"\n\t\tcoreLoadError = \"Minicore is still loading\"\n\t\tlogger.Debugf(\"Minicore not yet loaded for client environment telemetry\")\n\t}\n\tlibcInfo := internalos.GetLibcInfo()\n\tlinkingMode, err := compilation.CheckDynamicLinking()\n\tif err != nil {\n\t\tlogger.Debugf(\"cannot determine if app is dynamically linked: %v\", err)\n\t}\n\treturn authRequestClientEnvironment{\n\t\tOs: runtime.GOOS,\n\t\tOsVersion: osVersion,\n\t\tOsDetails: internalos.GetOsDetails(),\n\t\tIsa: runtime.GOARCH,\n\t\tGoVersion: runtime.Version(),\n\t\tCoreVersion: coreVersion,\n\t\tCoreFileName: getMiniCoreFileName(),\n\t\tCoreLoadError: coreLoadError,\n\t\tCgoEnabled: compilation.CgoEnabled,\n\t\tLinkingMode: linkingMode.String(),\n\t\tLibcFamily: libcInfo.Family,\n\t\tLibcVersion: libcInfo.Version,\n\t}\n}\n\nfunc createRequestBody(sc *snowflakeConn, sessionParameters map[string]any,\n\tclientEnvironment authRequestClientEnvironment, proofKey []byte, samlResponse []byte,\n) ([]byte, error) {\n\trequestMain := authRequestData{\n\t\tClientAppID: clientType,\n\t\tClientAppVersion: SnowflakeGoDriverVersion,\n\t\tAccountName: sc.cfg.Account,\n\t\tSessionParameters: sessionParameters,\n\t\tClientEnvironment: clientEnvironment,\n\t}\n\n\tswitch sc.cfg.Authenticator {\n\tcase AuthTypeExternalBrowser:\n\t\tif sc.idToken != \"\" {\n\t\t\trequestMain.Authenticator = idTokenAuthenticator\n\t\t\trequestMain.Token = sc.idToken\n\t\t\trequestMain.LoginName = sc.cfg.User\n\t\t} else {\n\t\t\trequestMain.ProofKey = string(proofKey)\n\t\t\trequestMain.Token = string(samlResponse)\n\t\t\trequestMain.LoginName = sc.cfg.User\n\t\t\trequestMain.Authenticator = AuthTypeExternalBrowser.String()\n\t\t}\n\tcase AuthTypeOAuth:\n\t\trequestMain.LoginName = sc.cfg.User\n\t\trequestMain.Authenticator = AuthTypeOAuth.String()\n\t\tvar err error\n\t\tif requestMain.Token, err = sfconfig.GetToken(sc.cfg); err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to get OAuth token: %w\", err)\n\t\t}\n\tcase AuthTypeOkta:\n\t\tsamlResponse, err := authenticateBySAML(\n\t\t\tsc.ctx,\n\t\t\tsc.rest,\n\t\t\tsc.cfg.OktaURL,\n\t\t\tsc.cfg.Application,\n\t\t\tsc.cfg.Account,\n\t\t\tsc.cfg.User,\n\t\t\tsc.cfg.Password,\n\t\t\tsc.cfg.DisableSamlURLCheck)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\trequestMain.RawSAMLResponse = string(samlResponse)\n\tcase AuthTypeJwt:\n\t\trequestMain.Authenticator = AuthTypeJwt.String()\n\n\t\tjwtTokenString, err := prepareJWTToken(sc.cfg)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\trequestMain.Token = jwtTokenString\n\tcase AuthTypePat:\n\t\tlogger.WithContext(sc.ctx).Info(\"Programmatic access token\")\n\t\trequestMain.Authenticator = AuthTypePat.String()\n\t\trequestMain.LoginName = sc.cfg.User\n\t\tvar err error\n\t\tif requestMain.Token, err = sfconfig.GetToken(sc.cfg); err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to get PAT token: %w\", err)\n\t\t}\n\tcase AuthTypeSnowflake:\n\t\tlogger.WithContext(sc.ctx).Debug(\"Username and password\")\n\t\trequestMain.LoginName = sc.cfg.User\n\t\trequestMain.Password = sc.cfg.Password\n\t\tswitch {\n\t\tcase sc.cfg.PasscodeInPassword:\n\t\t\trequestMain.ExtAuthnDuoMethod = \"passcode\"\n\t\tcase sc.cfg.Passcode != \"\":\n\t\t\trequestMain.Passcode = sc.cfg.Passcode\n\t\t\trequestMain.ExtAuthnDuoMethod = \"passcode\"\n\t\t}\n\tcase AuthTypeUsernamePasswordMFA:\n\t\tlogger.WithContext(sc.ctx).Debug(\"Username and password MFA\")\n\t\trequestMain.LoginName = sc.cfg.User\n\t\trequestMain.Password = sc.cfg.Password\n\t\tswitch {\n\t\tcase sc.mfaToken != \"\":\n\t\t\trequestMain.Token = sc.mfaToken\n\t\tcase sc.cfg.PasscodeInPassword:\n\t\t\trequestMain.ExtAuthnDuoMethod = \"passcode\"\n\t\tcase sc.cfg.Passcode != \"\":\n\t\t\trequestMain.Passcode = sc.cfg.Passcode\n\t\t\trequestMain.ExtAuthnDuoMethod = \"passcode\"\n\t\t}\n\tcase AuthTypeOAuthAuthorizationCode:\n\t\tlogger.WithContext(sc.ctx).Debug(\"OAuth authorization code\")\n\t\ttoken, err := authenticateByAuthorizationCode(sc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\trequestMain.LoginName = sc.cfg.User\n\t\trequestMain.Token = token\n\tcase AuthTypeOAuthClientCredentials:\n\t\tlogger.WithContext(sc.ctx).Debug(\"OAuth client credentials\")\n\t\toauthClient, err := newOauthClient(sc.ctx, sc.cfg, sc)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\ttoken, err := oauthClient.authenticateByOAuthClientCredentials()\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\trequestMain.LoginName = sc.cfg.User\n\t\trequestMain.Token = token\n\tcase AuthTypeWorkloadIdentityFederation:\n\t\tlogger.WithContext(sc.ctx).Debug(\"Workload Identity Federation\")\n\t\twifAttestationProvider := createWifAttestationProvider(sc.ctx, sc.cfg, sc.telemetry)\n\t\twifAttestation, err := wifAttestationProvider.getAttestation(sc.cfg.WorkloadIdentityProvider)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tif wifAttestation == nil {\n\t\t\treturn nil, errors.New(\"workload identity federation attestation is not available, please check your configuration\")\n\t\t}\n\t\trequestMain.Authenticator = AuthTypeWorkloadIdentityFederation.String()\n\t\trequestMain.Token = wifAttestation.Credential\n\t\trequestMain.Provider = wifAttestation.ProviderType\n\t}\n\n\tlogger.WithContext(sc.ctx).Debugf(\"Request body is created for the authentication. Authenticator: %s, User: %s, Account: %s\", sc.cfg.Authenticator.String(), sc.cfg.User, sc.cfg.Account)\n\n\tauthRequest := authRequest{\n\t\tData: requestMain,\n\t}\n\tjsonBody, err := json.Marshal(authRequest)\n\tif err != nil {\n\t\tlogger.WithContext(sc.ctx).Errorf(\"Failed to marshal JSON. err: %v\", err)\n\t\treturn nil, err\n\t}\n\treturn jsonBody, nil\n}\n\ntype oauthLockKey struct {\n\ttokenRequestURL string\n\tuser string\n\tflowType string\n}\n\nfunc newOAuthAuthorizationCodeLockKey(tokenRequestURL, user string) *oauthLockKey {\n\treturn &oauthLockKey{\n\t\ttokenRequestURL: tokenRequestURL,\n\t\tuser: user,\n\t\tflowType: \"authorization_code\",\n\t}\n}\n\nfunc newRefreshTokenLockKey(tokenRequestURL, user string) *oauthLockKey {\n\treturn &oauthLockKey{\n\t\ttokenRequestURL: tokenRequestURL,\n\t\tuser: user,\n\t\tflowType: \"refresh_token\",\n\t}\n}\n\nfunc (o *oauthLockKey) lockID() string {\n\treturn o.tokenRequestURL + \"|\" + o.user + \"|\" + o.flowType\n}\n\nfunc authenticateByAuthorizationCode(sc *snowflakeConn) (string, error) {\n\toauthClient, err := newOauthClient(sc.ctx, sc.cfg, sc)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tif !isEligibleForParallelLogin(sc.cfg, sc.cfg.ClientStoreTemporaryCredential) {\n\t\treturn oauthClient.authenticateByOAuthAuthorizationCode()\n\t}\n\n\tlockKey := newOAuthAuthorizationCodeLockKey(oauthClient.tokenURL(), sc.cfg.User)\n\tvalueAwaiter := valueAwaitHolder.get(lockKey)\n\tdefer valueAwaiter.resumeOne()\n\ttoken, err := awaitValue(valueAwaiter, func() (string, error) {\n\t\treturn credentialsStorage.getCredential(newOAuthAccessTokenSpec(oauthClient.tokenURL(), sc.cfg.User)), nil\n\t}, func(s string, err error) bool {\n\t\treturn s != \"\"\n\t}, func() string {\n\t\treturn \"\"\n\t})\n\tif err != nil || token != \"\" {\n\t\treturn token, err\n\t}\n\ttoken, err = oauthClient.authenticateByOAuthAuthorizationCode()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tvalueAwaiter.done()\n\treturn token, err\n}\n\n// Generate a JWT token in string given the configuration\nfunc prepareJWTToken(config *Config) (string, error) {\n\tif config.PrivateKey == nil {\n\t\treturn \"\", errors.New(\"trying to use keypair authentication, but PrivateKey was not provided in the driver config\")\n\t}\n\tlogger.Debug(\"preparing JWT for keypair authentication\")\n\tpubBytes, err := x509.MarshalPKIXPublicKey(config.PrivateKey.Public())\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\thash := sha256.Sum256(pubBytes)\n\n\taccountName := sfconfig.ExtractAccountName(config.Account)\n\tuserName := strings.ToUpper(config.User)\n\n\tissueAtTime := time.Now().UTC()\n\tjwtClaims := jwt.MapClaims{\n\t\t\"iss\": fmt.Sprintf(\"%s.%s.%s\", accountName, userName, \"SHA256:\"+base64.StdEncoding.EncodeToString(hash[:])),\n\t\t\"sub\": fmt.Sprintf(\"%s.%s\", accountName, userName),\n\t\t\"iat\": issueAtTime.Unix(),\n\t\t\"nbf\": time.Date(2015, 10, 10, 12, 0, 0, 0, time.UTC).Unix(),\n\t\t\"exp\": issueAtTime.Add(config.JWTExpireTimeout).Unix(),\n\t}\n\ttoken := jwt.NewWithClaims(jwt.SigningMethodRS256, jwtClaims)\n\n\ttokenString, err := token.SignedString(config.PrivateKey)\n\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tlogger.Debugf(\"successfully generated JWT with following claims: %v\", jwtClaims)\n\treturn tokenString, err\n}\n\ntype tokenLockKey struct {\n\tsnowflakeHost string\n\tuser string\n\ttokenType string\n}\n\nfunc newMfaTokenLockKey(snowflakeHost, user string) *tokenLockKey {\n\treturn &tokenLockKey{\n\t\tsnowflakeHost: snowflakeHost,\n\t\tuser: user,\n\t\ttokenType: \"MFA\",\n\t}\n}\n\nfunc newIDTokenLockKey(snowflakeHost, user string) *tokenLockKey {\n\treturn &tokenLockKey{\n\t\tsnowflakeHost: snowflakeHost,\n\t\tuser: user,\n\t\ttokenType: \"ID\",\n\t}\n}\n\nfunc (m *tokenLockKey) lockID() string {\n\treturn m.snowflakeHost + \"|\" + m.user + \"|\" + m.tokenType\n}\n\nfunc authenticateWithConfig(sc *snowflakeConn) error {\n\tvar authData *authResponseMain\n\tvar samlResponse []byte\n\tvar proofKey []byte\n\tvar err error\n\n\tmfaTokenLockKey := newMfaTokenLockKey(sc.cfg.Host, sc.cfg.User)\n\tidTokenLockKey := newIDTokenLockKey(sc.cfg.Host, sc.cfg.User)\n\n\tif sc.cfg.Authenticator == AuthTypeExternalBrowser || sc.cfg.Authenticator == AuthTypeOAuthAuthorizationCode || sc.cfg.Authenticator == AuthTypeOAuthClientCredentials {\n\t\tif (runtime.GOOS == \"windows\" || runtime.GOOS == \"darwin\") && sc.cfg.ClientStoreTemporaryCredential == sfconfig.BoolNotSet {\n\t\t\tsc.cfg.ClientStoreTemporaryCredential = ConfigBoolTrue\n\t\t}\n\t\tif sc.cfg.Authenticator == AuthTypeExternalBrowser {\n\t\t\tif isEligibleForParallelLogin(sc.cfg, sc.cfg.ClientStoreTemporaryCredential) {\n\t\t\t\tvalueAwaiter := valueAwaitHolder.get(idTokenLockKey)\n\t\t\t\tdefer valueAwaiter.resumeOne()\n\t\t\t\tsc.idToken, _ = awaitValue(valueAwaiter, func() (string, error) {\n\t\t\t\t\tcredential := credentialsStorage.getCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))\n\t\t\t\t\treturn credential, nil\n\t\t\t\t}, func(s string, err error) bool {\n\t\t\t\t\treturn s != \"\"\n\t\t\t\t}, func() string {\n\t\t\t\t\treturn \"\"\n\t\t\t\t})\n\t\t\t} else if sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {\n\t\t\t\tsc.idToken = credentialsStorage.getCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))\n\t\t\t}\n\t\t}\n\t\t// Disable console login by default\n\t\tif sc.cfg.DisableConsoleLogin == sfconfig.BoolNotSet {\n\t\t\tsc.cfg.DisableConsoleLogin = ConfigBoolTrue\n\t\t}\n\t}\n\n\tif sc.cfg.Authenticator == AuthTypeUsernamePasswordMFA {\n\t\tif (runtime.GOOS == \"windows\" || runtime.GOOS == \"darwin\") && sc.cfg.ClientRequestMfaToken == sfconfig.BoolNotSet {\n\t\t\tsc.cfg.ClientRequestMfaToken = ConfigBoolTrue\n\t\t}\n\t\tif isEligibleForParallelLogin(sc.cfg, sc.cfg.ClientRequestMfaToken) {\n\t\t\tvalueAwaiter := valueAwaitHolder.get(mfaTokenLockKey)\n\t\t\tdefer valueAwaiter.resumeOne()\n\t\t\tsc.mfaToken, _ = awaitValue(valueAwaiter, func() (string, error) {\n\t\t\t\tcredential := credentialsStorage.getCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))\n\t\t\t\treturn credential, nil\n\t\t\t}, func(s string, err error) bool {\n\t\t\t\treturn s != \"\"\n\t\t\t}, func() string {\n\t\t\t\treturn \"\"\n\t\t\t})\n\t\t} else if sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {\n\t\t\tsc.mfaToken = credentialsStorage.getCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))\n\t\t}\n\t}\n\n\tlogger.WithContext(sc.ctx).Infof(\"Authenticating via %v\", sc.cfg.Authenticator.String())\n\tswitch sc.cfg.Authenticator {\n\tcase AuthTypeExternalBrowser:\n\t\tif sc.idToken == \"\" {\n\t\t\tsamlResponse, proofKey, err = authenticateByExternalBrowser(\n\t\t\t\tsc.ctx,\n\t\t\t\tsc.rest,\n\t\t\t\tsc.cfg.Authenticator.String(),\n\t\t\t\tsc.cfg.Application,\n\t\t\t\tsc.cfg.Account,\n\t\t\t\tsc.cfg.User,\n\t\t\t\tsc.cfg.ExternalBrowserTimeout,\n\t\t\t\tsc.cfg.DisableConsoleLogin)\n\t\t\tif err != nil {\n\t\t\t\tsc.cleanup()\n\t\t\t\treturn err\n\t\t\t}\n\t\t}\n\t}\n\tauthData, err = authenticate(\n\t\tsc.ctx,\n\t\tsc,\n\t\tsamlResponse,\n\t\tproofKey)\n\tif err != nil {\n\t\tvar se *SnowflakeError\n\t\tif errors.As(err, &se) && slices.Contains(refreshOAuthTokenErrorCodes, strconv.Itoa(se.Number)) {\n\t\t\tcredentialsStorage.deleteCredential(newOAuthAccessTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))\n\n\t\t\tif sc.cfg.Authenticator == AuthTypeOAuthAuthorizationCode {\n\t\t\t\tdoRefreshTokenWithLock(sc)\n\t\t\t}\n\n\t\t\t// if refreshing succeeds for authorization code, we will take a token from cache\n\t\t\t// if it fails, we will just run the full flow\n\t\t\tauthData, err = authenticate(sc.ctx, sc, nil, nil)\n\t\t}\n\t\tif err != nil {\n\t\t\tsc.cleanup()\n\t\t\treturn err\n\t\t}\n\t}\n\tif sc.cfg.Authenticator == AuthTypeUsernamePasswordMFA && isEligibleForParallelLogin(sc.cfg, sc.cfg.ClientRequestMfaToken) {\n\t\tvalueAwaiter := valueAwaitHolder.get(mfaTokenLockKey)\n\t\tvalueAwaiter.done()\n\t}\n\tif sc.cfg.Authenticator == AuthTypeExternalBrowser && isEligibleForParallelLogin(sc.cfg, sc.cfg.ClientStoreTemporaryCredential) {\n\t\tvalueAwaiter := valueAwaitHolder.get(idTokenLockKey)\n\t\tvalueAwaiter.done()\n\t}\n\tsc.populateSessionParameters(authData.Parameters)\n\tsc.configureTelemetry()\n\tsc.ctx = context.WithValue(sc.ctx, SFSessionIDKey, authData.SessionID)\n\treturn nil\n}\n\nfunc doRefreshTokenWithLock(sc *snowflakeConn) {\n\tif oauthClient, err := newOauthClient(sc.ctx, sc.cfg, sc); err != nil {\n\t\tlogger.Warnf(\"failed to create oauth client. %v\", err)\n\t} else {\n\t\tlockKey := newRefreshTokenLockKey(oauthClient.tokenURL(), sc.cfg.User)\n\t\tif _, err = getValueWithLock(chooseLockerForAuth(sc.cfg), lockKey, func() (string, error) {\n\t\t\tif err = oauthClient.refreshToken(); err != nil {\n\t\t\t\tlogger.Warnf(\"cannot refresh token. %v\", err)\n\t\t\t\tcredentialsStorage.deleteCredential(newOAuthRefreshTokenSpec(sc.cfg.OauthTokenRequestURL, sc.cfg.User))\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t\treturn \"\", nil\n\t\t}); err != nil {\n\t\t\tlogger.Warnf(\"failed to refresh token with lock. %v\", err)\n\t\t}\n\t}\n}\n\nfunc chooseLockerForAuth(cfg *Config) locker {\n\tif cfg.SingleAuthenticationPrompt == ConfigBoolFalse {\n\t\treturn noopLocker\n\t}\n\tif cfg.User == \"\" {\n\t\treturn noopLocker\n\t}\n\treturn exclusiveLocker\n}\n\nfunc isEligibleForParallelLogin(cfg *Config, cacheEnabled ConfigBool) bool {\n\treturn cfg.SingleAuthenticationPrompt != ConfigBoolFalse && cfg.User != \"\" && cacheEnabled == ConfigBoolTrue\n}\n"
},
{
"path": "auth_generic_test_methods_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"testing\"\n)\n\nfunc getAuthTestConfigFromEnv() (*Config, error) {\n\treturn GetConfigFromEnv([]*ConfigParam{\n\t\t{Name: \"Account\", EnvName: \"SNOWFLAKE_TEST_ACCOUNT\", FailOnMissing: true},\n\t\t{Name: \"User\", EnvName: \"SNOWFLAKE_AUTH_TEST_OKTA_USER\", FailOnMissing: true},\n\t\t{Name: \"Password\", EnvName: \"SNOWFLAKE_AUTH_TEST_OKTA_PASS\", FailOnMissing: true},\n\t\t{Name: \"Host\", EnvName: \"SNOWFLAKE_TEST_HOST\", FailOnMissing: false},\n\t\t{Name: \"Port\", EnvName: \"SNOWFLAKE_TEST_PORT\", FailOnMissing: false},\n\t\t{Name: \"Protocol\", EnvName: \"SNOWFLAKE_AUTH_TEST_PROTOCOL\", FailOnMissing: false},\n\t\t{Name: \"Role\", EnvName: \"SNOWFLAKE_TEST_ROLE\", FailOnMissing: false},\n\t\t{Name: \"Warehouse\", EnvName: \"SNOWFLAKE_TEST_WAREHOUSE\", FailOnMissing: false},\n\t})\n}\n\nfunc getAuthTestsConfig(t *testing.T, authMethod AuthType) (*Config, error) {\n\tcfg, err := getAuthTestConfigFromEnv()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get config: %v\", err))\n\n\tcfg.Authenticator = authMethod\n\n\treturn cfg, nil\n}\n\nfunc isTestRunningInDockerContainer() bool {\n\treturn os.Getenv(\"AUTHENTICATION_TESTS_ENV\") == \"docker\"\n}\n"
},
{
"path": "auth_oauth.go",
"content": "package gosnowflake\n\nimport (\n\t\"bufio\"\n\t\"bytes\"\n\t\"cmp\"\n\t\"context\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"html\"\n\t\"io\"\n\t\"net\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"strconv\"\n\t\"strings\"\n\t\"time\"\n\n\t\"golang.org/x/oauth2\"\n\t\"golang.org/x/oauth2/clientcredentials\"\n)\n\nconst (\n\toauthSuccessHTML = `\nOAuth for Snowflake\n\nOAuth authentication completed successfully.\n`\n\tlocalApplicationClientCredentials = \"LOCAL_APPLICATION\"\n)\n\nvar defaultAuthorizationCodeProviderFactory = func() authorizationCodeProvider {\n\treturn &browserBasedAuthorizationCodeProvider{}\n}\n\ntype oauthClient struct {\n\tctx context.Context\n\tcfg *Config\n\tclient *http.Client\n\n\tport int\n\tredirectURITemplate string\n\n\tauthorizationCodeProviderFactory func() authorizationCodeProvider\n}\n\nfunc newOauthClient(ctx context.Context, cfg *Config, sc *snowflakeConn) (*oauthClient, error) {\n\tport := 0\n\tif cfg.OauthRedirectURI != \"\" {\n\t\tlogger.Debugf(\"Using oauthRedirectUri from config: %v\", cfg.OauthRedirectURI)\n\t\turi, err := url.Parse(cfg.OauthRedirectURI)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tportStr := uri.Port()\n\t\tif portStr != \"\" {\n\t\t\tif port, err = strconv.Atoi(portStr); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\n\tredirectURITemplate := \"\"\n\tif cfg.OauthRedirectURI == \"\" {\n\t\tredirectURITemplate = \"http://127.0.0.1:%v\"\n\t}\n\tlogger.Debugf(\"Redirect URI template: %v, port: %v\", redirectURITemplate, port)\n\n\ttransport, err := newTransportFactory(cfg, sc.telemetry).createTransport(transportConfigFor(transportTypeOAuth))\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tclient := &http.Client{\n\t\tTransport: transport,\n\t}\n\treturn &oauthClient{\n\t\tctx: context.WithValue(ctx, oauth2.HTTPClient, client),\n\t\tcfg: cfg,\n\t\tclient: client,\n\t\tport: port,\n\t\tredirectURITemplate: redirectURITemplate,\n\t\tauthorizationCodeProviderFactory: defaultAuthorizationCodeProviderFactory,\n\t}, nil\n}\n\ntype oauthBrowserResult struct {\n\taccessToken string\n\trefreshToken string\n\terr error\n}\n\nfunc (oauthClient *oauthClient) authenticateByOAuthAuthorizationCode() (string, error) {\n\taccessTokenSpec := oauthClient.accessTokenSpec()\n\tif oauthClient.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {\n\t\tif accessToken := credentialsStorage.getCredential(accessTokenSpec); accessToken != \"\" {\n\t\t\tlogger.Debugf(\"Access token retrieved from cache\")\n\t\t\treturn accessToken, nil\n\t\t}\n\t\tif refreshToken := credentialsStorage.getCredential(oauthClient.refreshTokenSpec()); refreshToken != \"\" {\n\t\t\treturn \"\", &SnowflakeError{Number: ErrMissingAccessATokenButRefreshTokenPresent}\n\t\t}\n\t}\n\tlogger.Debugf(\"Access token not present in cache, running full auth code flow\")\n\n\tresultChan := make(chan oauthBrowserResult, 1)\n\ttcpListener, callbackPort, err := oauthClient.setupListener()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tdefer func() {\n\t\tlogger.Debug(\"Closing tcp listener\")\n\t\tif err := tcpListener.Close(); err != nil {\n\t\t\tlogger.Warnf(\"error while closing TCP listener. %v\", err)\n\t\t}\n\t}()\n\tgo GoroutineWrapper(oauthClient.ctx, func() {\n\t\tresultChan <- oauthClient.doAuthenticateByOAuthAuthorizationCode(tcpListener, callbackPort)\n\t})\n\tselect {\n\tcase <-time.After(oauthClient.cfg.ExternalBrowserTimeout):\n\t\treturn \"\", errors.New(\"authentication via browser timed out\")\n\tcase result := <-resultChan:\n\t\tif oauthClient.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {\n\t\t\tlogger.Debug(\"saving oauth access token in cache\")\n\t\t\tcredentialsStorage.setCredential(oauthClient.accessTokenSpec(), result.accessToken)\n\t\t\tcredentialsStorage.setCredential(oauthClient.refreshTokenSpec(), result.refreshToken)\n\t\t}\n\t\treturn result.accessToken, result.err\n\t}\n}\n\nfunc (oauthClient *oauthClient) doAuthenticateByOAuthAuthorizationCode(tcpListener *net.TCPListener, callbackPort int) oauthBrowserResult {\n\tauthCodeProvider := oauthClient.authorizationCodeProviderFactory()\n\n\tsuccessChan := make(chan []byte)\n\terrChan := make(chan error)\n\tresponseBodyChan := make(chan string, 2)\n\tcloseListenerChan := make(chan bool, 2)\n\n\tdefer func() {\n\t\tcloseListenerChan <- true\n\t\tclose(successChan)\n\t\tclose(errChan)\n\t\tclose(responseBodyChan)\n\t\tclose(closeListenerChan)\n\t}()\n\n\tlogger.Debugf(\"opening socket on port %v\", callbackPort)\n\tdefer func(tcpListener *net.TCPListener) {\n\t\t<-closeListenerChan\n\t}(tcpListener)\n\n\tgo handleOAuthSocket(tcpListener, successChan, errChan, responseBodyChan, closeListenerChan)\n\n\toauth2cfg := oauthClient.buildAuthorizationCodeConfig(callbackPort)\n\tcodeVerifier := authCodeProvider.createCodeVerifier()\n\tstate := authCodeProvider.createState()\n\tauthorizationURL := oauth2cfg.AuthCodeURL(state, oauth2.S256ChallengeOption(codeVerifier))\n\tif err := authCodeProvider.run(authorizationURL); err != nil {\n\t\tresponseBodyChan <- err.Error()\n\t\tcloseListenerChan <- true\n\t\treturn oauthBrowserResult{\"\", \"\", err}\n\t}\n\n\terr := <-errChan\n\tif err != nil {\n\t\tresponseBodyChan <- err.Error()\n\t\treturn oauthBrowserResult{\"\", \"\", err}\n\t}\n\tcodeReqBytes := <-successChan\n\n\tcodeReq, err := http.ReadRequest(bufio.NewReader(bytes.NewReader(codeReqBytes)))\n\tif err != nil {\n\t\tresponseBodyChan <- err.Error()\n\t\treturn oauthBrowserResult{\"\", \"\", err}\n\t}\n\tlogger.Debugf(\"Received authorization code from %v\", oauthClient.authorizationURL())\n\ttokenResponse, err := oauthClient.exchangeAccessToken(codeReq, state, oauth2cfg, codeVerifier, responseBodyChan)\n\tif err != nil {\n\t\treturn oauthBrowserResult{\"\", \"\", err}\n\t}\n\tlogger.Debugf(\"Received token from %v\", oauthClient.tokenURL())\n\treturn oauthBrowserResult{tokenResponse.AccessToken, tokenResponse.RefreshToken, err}\n}\n\nfunc (oauthClient *oauthClient) setupListener() (*net.TCPListener, int, error) {\n\ttcpListener, err := createLocalTCPListener(oauthClient.port)\n\tif err != nil {\n\t\treturn nil, 0, err\n\t}\n\tcallbackPort := tcpListener.Addr().(*net.TCPAddr).Port\n\tlogger.Debugf(\"oauthClient.port: %v, callbackPort: %v\", oauthClient.port, callbackPort)\n\treturn tcpListener, callbackPort, nil\n}\n\nfunc (oauthClient *oauthClient) exchangeAccessToken(codeReq *http.Request, state string, oauth2cfg *oauth2.Config, codeVerifier string, responseBodyChan chan string) (*oauth2.Token, error) {\n\tqueryParams := codeReq.URL.Query()\n\terrorMsg := queryParams.Get(\"error\")\n\tif errorMsg != \"\" {\n\t\terrorDesc := queryParams.Get(\"error_description\")\n\t\terrMsg := fmt.Sprintf(\"error while getting authentication from oauth: %v. Details: %v\", errorMsg, errorDesc)\n\t\tresponseBodyChan <- html.EscapeString(errMsg)\n\t\treturn nil, errors.New(errMsg)\n\t}\n\n\treceivedState := queryParams.Get(\"state\")\n\tif state != receivedState {\n\t\terrMsg := \"invalid oauth state received\"\n\t\tresponseBodyChan <- errMsg\n\t\treturn nil, errors.New(errMsg)\n\t}\n\n\tcode := queryParams.Get(\"code\")\n\topts := []oauth2.AuthCodeOption{oauth2.VerifierOption(codeVerifier)}\n\tif oauthClient.cfg.EnableSingleUseRefreshTokens {\n\t\topts = append(opts, oauth2.SetAuthURLParam(\"enable_single_use_refresh_tokens\", \"true\"))\n\t}\n\ttoken, err := oauth2cfg.Exchange(oauthClient.ctx, code, opts...)\n\tif err != nil {\n\t\tresponseBodyChan <- err.Error()\n\t\treturn nil, err\n\t}\n\tresponseBodyChan <- oauthSuccessHTML\n\treturn token, nil\n}\n\nfunc (oauthClient *oauthClient) buildAuthorizationCodeConfig(callbackPort int) *oauth2.Config {\n\tclientID, clientSecret := oauthClient.cfg.OauthClientID, oauthClient.cfg.OauthClientSecret\n\tif oauthClient.eligibleForDefaultClientCredentials() {\n\t\tclientID, clientSecret = localApplicationClientCredentials, localApplicationClientCredentials\n\t}\n\toauthClient.logIfHTTPInUse(oauthClient.authorizationURL())\n\toauthClient.logIfHTTPInUse(oauthClient.tokenURL())\n\treturn &oauth2.Config{\n\t\tClientID: clientID,\n\t\tClientSecret: clientSecret,\n\t\tRedirectURL: oauthClient.buildRedirectURI(callbackPort),\n\t\tScopes: oauthClient.buildScopes(),\n\t\tEndpoint: oauth2.Endpoint{\n\t\t\tAuthURL: oauthClient.authorizationURL(),\n\t\t\tTokenURL: oauthClient.tokenURL(),\n\t\t\tAuthStyle: oauth2.AuthStyleInHeader,\n\t\t},\n\t}\n}\nfunc (oauthClient *oauthClient) eligibleForDefaultClientCredentials() bool {\n\treturn oauthClient.cfg.OauthClientID == \"\" && oauthClient.cfg.OauthClientSecret == \"\" && oauthClient.isSnowflakeAsIDP()\n}\n\nfunc (oauthClient *oauthClient) isSnowflakeAsIDP() bool {\n\treturn (oauthClient.cfg.OauthAuthorizationURL == \"\" || strings.Contains(oauthClient.cfg.OauthAuthorizationURL, oauthClient.cfg.Host)) &&\n\t\t(oauthClient.cfg.OauthTokenRequestURL == \"\" || strings.Contains(oauthClient.cfg.OauthTokenRequestURL, oauthClient.cfg.Host))\n}\n\nfunc (oauthClient *oauthClient) authorizationURL() string {\n\treturn cmp.Or(oauthClient.cfg.OauthAuthorizationURL, oauthClient.defaultAuthorizationURL())\n}\n\nfunc (oauthClient *oauthClient) defaultAuthorizationURL() string {\n\treturn fmt.Sprintf(\"%v://%v:%v/oauth/authorize\", oauthClient.cfg.Protocol, oauthClient.cfg.Host, oauthClient.cfg.Port)\n}\n\nfunc (oauthClient *oauthClient) tokenURL() string {\n\treturn cmp.Or(oauthClient.cfg.OauthTokenRequestURL, oauthClient.defaultTokenURL())\n}\n\nfunc (oauthClient *oauthClient) defaultTokenURL() string {\n\treturn fmt.Sprintf(\"%v://%v:%v/oauth/token-request\", oauthClient.cfg.Protocol, oauthClient.cfg.Host, oauthClient.cfg.Port)\n}\n\nfunc (oauthClient *oauthClient) buildRedirectURI(port int) string {\n\tif oauthClient.cfg.OauthRedirectURI != \"\" {\n\t\treturn oauthClient.cfg.OauthRedirectURI\n\t}\n\treturn fmt.Sprintf(oauthClient.redirectURITemplate, port)\n}\n\nfunc (oauthClient *oauthClient) buildScopes() []string {\n\tif oauthClient.cfg.OauthScope == \"\" {\n\t\treturn []string{\"session:role:\" + oauthClient.cfg.Role}\n\t}\n\tscopes := strings.Split(oauthClient.cfg.OauthScope, \" \")\n\tfor i, scope := range scopes {\n\t\tscopes[i] = strings.TrimSpace(scope)\n\t}\n\treturn scopes\n}\n\nfunc handleOAuthSocket(tcpListener *net.TCPListener, successChan chan []byte, errChan chan error, responseBodyChan chan string, closeListenerChan chan bool) {\n\tconn, err := tcpListener.AcceptTCP()\n\tif err != nil {\n\t\tlogger.Warnf(\"error creating socket. %v\", err)\n\t\treturn\n\t}\n\tdefer func() {\n\t\tif err := conn.Close(); err != nil {\n\t\t\tlogger.Warnf(\"error while closing connection (%v -> %v). %v\", conn.LocalAddr(), conn.RemoteAddr(), err)\n\t\t}\n\t}()\n\tvar buf [bufSize]byte\n\tcodeResp := bytes.NewBuffer(nil)\n\tfor {\n\t\treadBytes, err := conn.Read(buf[:])\n\t\tif err == io.EOF {\n\t\t\tbreak\n\t\t}\n\t\tif err != nil {\n\t\t\terrChan <- err\n\t\t\treturn\n\t\t}\n\t\tcodeResp.Write(buf[0:readBytes])\n\t\tif readBytes < bufSize {\n\t\t\tbreak\n\t\t}\n\t}\n\n\terrChan <- nil\n\tsuccessChan <- codeResp.Bytes()\n\n\tresponseBody := <-responseBodyChan\n\trespToBrowser, err := buildResponse(responseBody)\n\tif err != nil {\n\t\tlogger.Warnf(\"cannot create response to browser. %v\", err)\n\t}\n\t_, err = conn.Write(respToBrowser.Bytes())\n\tif err != nil {\n\t\tlogger.Warnf(\"cannot write response to browser. %v\", err)\n\t}\n\tcloseListenerChan <- true\n}\n\ntype authorizationCodeProvider interface {\n\trun(authorizationURL string) error\n\tcreateState() string\n\tcreateCodeVerifier() string\n}\n\ntype browserBasedAuthorizationCodeProvider struct {\n}\n\nfunc (provider *browserBasedAuthorizationCodeProvider) run(authorizationURL string) error {\n\treturn openBrowser(authorizationURL)\n}\n\nfunc (provider *browserBasedAuthorizationCodeProvider) createState() string {\n\treturn NewUUID().String()\n}\n\nfunc (provider *browserBasedAuthorizationCodeProvider) createCodeVerifier() string {\n\treturn oauth2.GenerateVerifier()\n}\n\nfunc (oauthClient *oauthClient) authenticateByOAuthClientCredentials() (string, error) {\n\taccessTokenSpec := oauthClient.accessTokenSpec()\n\tif oauthClient.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {\n\t\tif accessToken := credentialsStorage.getCredential(accessTokenSpec); accessToken != \"\" {\n\t\t\treturn accessToken, nil\n\t\t}\n\t}\n\toauth2Cfg, err := oauthClient.buildClientCredentialsConfig()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\ttoken, err := oauth2Cfg.Token(oauthClient.ctx)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tif oauthClient.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {\n\t\tcredentialsStorage.setCredential(accessTokenSpec, token.AccessToken)\n\t}\n\treturn token.AccessToken, nil\n}\n\nfunc (oauthClient *oauthClient) buildClientCredentialsConfig() (*clientcredentials.Config, error) {\n\tif oauthClient.cfg.OauthTokenRequestURL == \"\" {\n\t\treturn nil, errors.New(\"client credentials flow requires tokenRequestURL\")\n\t}\n\treturn &clientcredentials.Config{\n\t\tClientID: oauthClient.cfg.OauthClientID,\n\t\tClientSecret: oauthClient.cfg.OauthClientSecret,\n\t\tTokenURL: oauthClient.cfg.OauthTokenRequestURL,\n\t\tScopes: oauthClient.buildScopes(),\n\t}, nil\n}\n\nfunc (oauthClient *oauthClient) refreshToken() error {\n\tif oauthClient.cfg.ClientStoreTemporaryCredential != ConfigBoolTrue {\n\t\tlogger.Debug(\"credentials storage is disabled, cannot use refresh tokens\")\n\t\treturn nil\n\t}\n\trefreshTokenSpec := newOAuthRefreshTokenSpec(oauthClient.cfg.OauthTokenRequestURL, oauthClient.cfg.User)\n\trefreshToken := credentialsStorage.getCredential(refreshTokenSpec)\n\tif refreshToken == \"\" {\n\t\tlogger.Debug(\"no refresh token in cache, full flow must be run\")\n\t\treturn nil\n\t}\n\tbody := url.Values{}\n\tbody.Add(\"grant_type\", \"refresh_token\")\n\tbody.Add(\"refresh_token\", refreshToken)\n\tbody.Add(\"scope\", strings.Join(oauthClient.buildScopes(), \" \"))\n\treq, err := http.NewRequest(\"POST\", oauthClient.tokenURL(), strings.NewReader(body.Encode()))\n\tif err != nil {\n\t\treturn err\n\t}\n\treq.SetBasicAuth(oauthClient.cfg.OauthClientID, oauthClient.cfg.OauthClientSecret)\n\treq.Header.Add(\"Content-Type\", \"application/x-www-form-urlencoded\")\n\tresp, err := oauthClient.client.Do(req)\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer func() {\n\t\tif err := resp.Body.Close(); err != nil {\n\t\t\tlogger.Warnf(\"error while closing response body for %v. %v\", req.URL, err)\n\t\t}\n\t}()\n\tif resp.StatusCode != 200 {\n\t\trespBody, err := io.ReadAll(resp.Body)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tcredentialsStorage.deleteCredential(refreshTokenSpec)\n\t\treturn errors.New(string(respBody))\n\t}\n\tvar tokenResponse tokenExchangeResponseBody\n\tif err = json.NewDecoder(resp.Body).Decode(&tokenResponse); err != nil {\n\t\treturn err\n\t}\n\taccessTokenSpec := oauthClient.accessTokenSpec()\n\tcredentialsStorage.setCredential(accessTokenSpec, tokenResponse.AccessToken)\n\tif tokenResponse.RefreshToken != \"\" {\n\t\tcredentialsStorage.setCredential(refreshTokenSpec, tokenResponse.RefreshToken)\n\t}\n\treturn nil\n}\n\ntype tokenExchangeResponseBody struct {\n\tAccessToken string `json:\"access_token,omitempty\"`\n\tRefreshToken string `json:\"refresh_token\"`\n}\n\nfunc (oauthClient *oauthClient) accessTokenSpec() *secureTokenSpec {\n\treturn newOAuthAccessTokenSpec(oauthClient.tokenURL(), oauthClient.cfg.User)\n}\n\nfunc (oauthClient *oauthClient) refreshTokenSpec() *secureTokenSpec {\n\treturn newOAuthRefreshTokenSpec(oauthClient.tokenURL(), oauthClient.cfg.User)\n}\n\nfunc (oauthClient *oauthClient) logIfHTTPInUse(u string) {\n\tparsed, err := url.Parse(u)\n\tif err != nil {\n\t\tlogger.Warnf(\"Cannot parse URL: %v. %v\", u, err)\n\t\treturn\n\t}\n\tif parsed.Scheme == \"http\" {\n\t\tlogger.Warnf(\"OAuth URL uses insecure HTTP protocol: %v\", u)\n\t}\n}\n"
},
{
"path": "auth_oauth_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"context\"\n\t\"database/sql\"\n\t\"errors\"\n\tsfconfig \"github.com/snowflakedb/gosnowflake/v2/internal/config\"\n\t\"io\"\n\t\"net/http\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\n\t\"golang.org/x/oauth2\"\n)\n\nfunc TestUnitOAuthAuthorizationCode(t *testing.T) {\n\tskipOnMac(t, \"keychain requires password\")\n\troundTripper := newCountingRoundTripper(createTestNoRevocationTransport())\n\thttpClient := &http.Client{\n\t\tTransport: roundTripper,\n\t}\n\tcfg := &Config{\n\t\tUser: \"testUser\",\n\t\tRole: \"ANALYST\",\n\t\tOauthClientID: \"testClientId\",\n\t\tOauthClientSecret: \"testClientSecret\",\n\t\tOauthAuthorizationURL: wiremock.baseURL() + \"/oauth/authorize\",\n\t\tOauthTokenRequestURL: wiremock.baseURL() + \"/oauth/token\",\n\t\tOauthRedirectURI: \"http://localhost:1234/snowflake/oauth-redirect\",\n\t\tTransporter: roundTripper,\n\t\tClientStoreTemporaryCredential: ConfigBoolTrue,\n\t\tExternalBrowserTimeout: time.Duration(sfconfig.DefaultExternalBrowserTimeout),\n\t}\n\tclient, err := newOauthClient(context.WithValue(context.Background(), oauth2.HTTPClient, httpClient), cfg, &snowflakeConn{})\n\tassertNilF(t, err)\n\taccessTokenSpec := newOAuthAccessTokenSpec(wiremock.connectionConfig().OauthTokenRequestURL, wiremock.connectionConfig().User)\n\trefreshTokenSpec := newOAuthRefreshTokenSpec(wiremock.connectionConfig().OauthTokenRequestURL, wiremock.connectionConfig().User)\n\n\tt.Run(\"Success\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(accessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(refreshTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"))\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{t: t}\n\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\treturn authCodeProvider\n\t\t}\n\t\ttoken, err := client.authenticateByOAuthAuthorizationCode()\n\t\tassertNilF(t, err)\n\t\tassertEqualE(t, token, \"access-token-123\")\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tauthCodeProvider.assertResponseBodyContains(\"OAuth authentication completed successfully.\")\n\t})\n\n\tt.Run(\"Store access token in cache\", func(t *testing.T) {\n\t\tskipOnMissingHome(t)\n\t\troundTripper.reset()\n\t\tcredentialsStorage.deleteCredential(accessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(refreshTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"))\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{t: t}\n\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\treturn authCodeProvider\n\t\t}\n\t\t_, err = client.authenticateByOAuthAuthorizationCode()\n\t\tassertNilF(t, err)\n\t\tassertEqualE(t, credentialsStorage.getCredential(accessTokenSpec), \"access-token-123\")\n\t})\n\n\tt.Run(\"Use cache for consecutive calls\", func(t *testing.T) {\n\t\tskipOnMissingHome(t)\n\t\troundTripper.reset()\n\t\tcredentialsStorage.setCredential(accessTokenSpec, \"access-token-123\")\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"))\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{t: t}\n\t\tfor range 3 {\n\t\t\tclient, err := newOauthClient(context.WithValue(context.Background(), oauth2.HTTPClient, httpClient), cfg, &snowflakeConn{})\n\t\t\tassertNilF(t, err)\n\t\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\t\treturn authCodeProvider\n\t\t\t}\n\t\t\t_, err = client.authenticateByOAuthAuthorizationCode()\n\t\t\tassertNilF(t, err)\n\t\t}\n\t\tassertEqualE(t, authCodeProvider.responseBody, \"\")\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 0)\n\t})\n\n\tt.Run(\"InvalidState\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(accessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(refreshTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"))\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{\n\t\t\ttamperWithState: true,\n\t\t\tt: t,\n\t\t}\n\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\treturn authCodeProvider\n\t\t}\n\t\t_, err = client.authenticateByOAuthAuthorizationCode()\n\t\tassertEqualE(t, err.Error(), \"invalid oauth state received\")\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tauthCodeProvider.assertResponseBodyContains(\"invalid oauth state received\")\n\t})\n\n\tt.Run(\"ErrorFromIdPWhileGettingCode\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(accessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(refreshTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/error_from_idp.json\"))\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{t: t}\n\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\treturn authCodeProvider\n\t\t}\n\t\t_, err = client.authenticateByOAuthAuthorizationCode()\n\t\tassertEqualE(t, err.Error(), \"error while getting authentication from oauth: some error. Details: some error desc\")\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tauthCodeProvider.assertResponseBodyContains(\"error while getting authentication from oauth: some error. Details: some error desc\")\n\t})\n\n\tt.Run(\"ErrorFromProviderWhileGettingCode\", func(t *testing.T) {\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{\n\t\t\ttriggerError: \"test error\",\n\t\t}\n\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\treturn authCodeProvider\n\t\t}\n\t\t_, err = client.authenticateByOAuthAuthorizationCode()\n\t\tassertEqualE(t, err.Error(), \"test error\")\n\t})\n\n\tt.Run(\"InvalidCode\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(accessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(refreshTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/invalid_code.json\"))\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{t: t}\n\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\treturn authCodeProvider\n\t\t}\n\t\t_, err = client.authenticateByOAuthAuthorizationCode()\n\t\tassertNotNilE(t, err)\n\t\tassertEqualE(t, err.(*oauth2.RetrieveError).ErrorCode, \"invalid_grant\")\n\t\tassertEqualE(t, err.(*oauth2.RetrieveError).ErrorDescription, \"The authorization code is invalid or has expired.\")\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tauthCodeProvider.assertResponseBodyContains(\"invalid_grant\")\n\t})\n\n\tt.Run(\"timeout\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(accessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(refreshTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"))\n\t\tclient.cfg.ExternalBrowserTimeout = 2 * time.Second\n\t\tauthCodeProvider := &nonInteractiveAuthorizationCodeProvider{\n\t\t\tsleepTime: 3 * time.Second,\n\t\t\ttriggerError: \"timed out\",\n\t\t\tt: t,\n\t\t}\n\t\tclient.authorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\treturn authCodeProvider\n\t\t}\n\t\t_, err = client.authenticateByOAuthAuthorizationCode()\n\t\tassertNotNilE(t, err)\n\t\tassertStringContainsE(t, err.Error(), \"timed out\")\n\t\ttime.Sleep(2 * time.Second) // awaiting timeout\n\t})\n}\n\nfunc TestUnitOAuthClientCredentials(t *testing.T) {\n\tskipOnMac(t, \"keychain requires password\")\n\tcacheTokenSpec := newOAuthAccessTokenSpec(wiremock.connectionConfig().OauthTokenRequestURL, wiremock.connectionConfig().User)\n\tcrt := newCountingRoundTripper(createTestNoRevocationTransport())\n\thttpClient := http.Client{\n\t\tTransport: crt,\n\t}\n\tcfgFactory := func() *Config {\n\t\treturn &Config{\n\t\t\tUser: \"testUser\",\n\t\t\tRole: \"ANALYST\",\n\t\t\tOauthClientID: \"testClientId\",\n\t\t\tOauthClientSecret: \"testClientSecret\",\n\t\t\tOauthTokenRequestURL: wiremock.baseURL() + \"/oauth/token\",\n\t\t\tTransporter: crt,\n\t\t\tClientStoreTemporaryCredential: ConfigBoolTrue,\n\t\t}\n\t}\n\tclient, err := newOauthClient(context.WithValue(context.Background(), oauth2.HTTPClient, httpClient), cfgFactory(), &snowflakeConn{})\n\tassertNilF(t, err)\n\n\tt.Run(\"success\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(cacheTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"))\n\t\ttoken, err := client.authenticateByOAuthClientCredentials()\n\t\tassertNilF(t, err)\n\t\tassertEqualE(t, token, \"access-token-123\")\n\t})\n\n\tt.Run(\"should store token in cache\", func(t *testing.T) {\n\t\tskipOnMissingHome(t)\n\t\tcrt.reset()\n\t\tcredentialsStorage.deleteCredential(cacheTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"))\n\t\ttoken, err := client.authenticateByOAuthClientCredentials()\n\t\tassertNilF(t, err)\n\t\tassertEqualE(t, token, \"access-token-123\")\n\n\t\tclient, err := newOauthClient(context.Background(), cfgFactory(), &snowflakeConn{})\n\t\tassertNilF(t, err)\n\t\ttoken, err = client.authenticateByOAuthClientCredentials()\n\t\tassertNilF(t, err)\n\t\tassertEqualE(t, token, \"access-token-123\")\n\n\t\tassertEqualE(t, crt.postReqCount[cfgFactory().OauthTokenRequestURL], 1)\n\t})\n\n\tt.Run(\"consecutive calls should take token from cache\", func(t *testing.T) {\n\t\tskipOnMissingHome(t)\n\t\tcrt.reset()\n\t\tcredentialsStorage.setCredential(cacheTokenSpec, \"access-token-123\")\n\t\tfor range 3 {\n\t\t\tclient, err := newOauthClient(context.Background(), cfgFactory(), &snowflakeConn{})\n\t\t\tassertNilF(t, err)\n\t\t\ttoken, err := client.authenticateByOAuthClientCredentials()\n\t\t\tassertNilF(t, err)\n\t\t\tassertEqualE(t, token, \"access-token-123\")\n\t\t}\n\t\tassertEqualE(t, crt.postReqCount[cfgFactory().OauthTokenRequestURL], 0)\n\t})\n\n\tt.Run(\"disabling cache\", func(t *testing.T) {\n\t\tskipOnMissingHome(t)\n\t\tcfg := cfgFactory()\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolFalse\n\t\tcredentialsStorage.deleteCredential(cacheTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"))\n\t\tclient, err := newOauthClient(context.Background(), cfg, &snowflakeConn{})\n\t\tassertNilF(t, err)\n\t\ttoken, err := client.authenticateByOAuthClientCredentials()\n\t\tassertNilF(t, err)\n\t\tassertEqualE(t, token, \"access-token-123\")\n\n\t\tclient, err = newOauthClient(context.Background(), cfg, &snowflakeConn{})\n\t\tassertNilF(t, err)\n\t\ttoken, err = client.authenticateByOAuthClientCredentials()\n\t\tassertNilF(t, err)\n\t\tassertEqualE(t, token, \"access-token-123\")\n\n\t\tassertEqualE(t, crt.postReqCount[cfg.OauthTokenRequestURL], 2)\n\t})\n\n\tt.Run(\"invalid_client\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(cacheTokenSpec)\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/client_credentials/invalid_client.json\"))\n\t\t_, err = client.authenticateByOAuthClientCredentials()\n\t\tassertNotNilF(t, err)\n\t\toauth2Err := err.(*oauth2.RetrieveError)\n\t\tassertEqualE(t, oauth2Err.ErrorCode, \"invalid_client\")\n\t\tassertEqualE(t, oauth2Err.ErrorDescription, \"The client secret supplied for a confidential client is invalid.\")\n\t})\n}\n\nfunc TestAuthorizationCodeFlow(t *testing.T) {\n\tif runningOnGithubAction() && runningOnLinux() {\n\t\tt.Skip(\"Github blocks writing to file system\")\n\t}\n\tskipOnMac(t, \"keychain requires password\")\n\tcurrentDefaultAuthorizationCodeProviderFactory := defaultAuthorizationCodeProviderFactory\n\tdefer func() {\n\t\tdefaultAuthorizationCodeProviderFactory = currentDefaultAuthorizationCodeProviderFactory\n\t}()\n\tdefaultAuthorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\treturn &nonInteractiveAuthorizationCodeProvider{\n\t\t\tt: t,\n\t\t\tmu: sync.Mutex{},\n\t\t}\n\t}\n\troundTripper := newCountingRoundTripper(createTestNoRevocationTransport())\n\n\tt.Run(\"successful flow\", func(t *testing.T) {\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Role = \"ANALYST\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.deleteCredential(oauthAccessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(oauthRefreshTokenSpec)\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t})\n\n\tt.Run(\"successful flow with multiple threads\", func(t *testing.T) {\n\t\tfor _, singleAuthenticationPrompt := range []ConfigBool{ConfigBoolFalse, ConfigBoolTrue, configBoolNotSet} {\n\t\t\tt.Run(\"singleAuthenticationPrompt=\"+singleAuthenticationPrompt.String(), func(t *testing.T) {\n\t\t\t\tcurrentDefaultAuthorizationCodeProviderFactory := defaultAuthorizationCodeProviderFactory\n\t\t\t\tdefer func() {\n\t\t\t\t\tdefaultAuthorizationCodeProviderFactory = currentDefaultAuthorizationCodeProviderFactory\n\t\t\t\t}()\n\t\t\t\tdefaultAuthorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\t\t\t\treturn &nonInteractiveAuthorizationCodeProvider{\n\t\t\t\t\t\tt: t,\n\t\t\t\t\t\tmu: sync.Mutex{},\n\t\t\t\t\t\tsleepTime: 500 * time.Millisecond,\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\troundTripper.reset()\n\t\t\t\twiremock.registerMappings(t,\n\t\t\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"),\n\t\t\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\t\t\tcfg := wiremock.connectionConfig()\n\t\t\t\tcfg.Role = \"ANALYST\"\n\t\t\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\t\t\tcfg.Transporter = roundTripper\n\t\t\t\tcfg.SingleAuthenticationPrompt = singleAuthenticationPrompt\n\t\t\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\t\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\t\t\tcredentialsStorage.deleteCredential(oauthAccessTokenSpec)\n\t\t\t\tcredentialsStorage.deleteCredential(oauthRefreshTokenSpec)\n\t\t\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\t\t\tdb := sql.OpenDB(connector)\n\t\t\t\tinitPoolWithSize(t, db, 20)\n\t\t\t\tprintln(roundTripper.postReqCount[cfg.OauthTokenRequestURL])\n\t\t\t\tif singleAuthenticationPrompt == ConfigBoolFalse {\n\t\t\t\t\tassertTrueE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL] > 1)\n\t\t\t\t} else {\n\t\t\t\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1)\n\t\t\t\t}\n\t\t\t})\n\t\t}\n\t})\n\n\tt.Run(\"successful flow with single-use refresh token enabled\", func(t *testing.T) {\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow_with_single_use_refresh_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Role = \"ANALYST\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\tcfg.EnableSingleUseRefreshTokens = true\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.deleteCredential(oauthAccessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(oauthRefreshTokenSpec)\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t})\n\n\tt.Run(\"should use cached access token\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Role = \"ANALYST\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.deleteCredential(oauthAccessTokenSpec)\n\t\tcredentialsStorage.deleteCredential(oauthRefreshTokenSpec)\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\tconn1, err := db.Conn(context.Background())\n\t\tassertNilF(t, err)\n\t\tdefer conn1.Close()\n\t\tconn2, err := db.Conn(context.Background())\n\t\tassertNilF(t, err)\n\t\tdefer conn2.Close()\n\t\trunSmokeQueryWithConn(t, conn1)\n\t\trunSmokeQueryWithConn(t, conn2)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1)\n\t})\n\n\tt.Run(\"should update cache with new token when the old one expired if refresh token is missing\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Role = \"ANALYST\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"expired-token\")\n\t\tcredentialsStorage.deleteCredential(oauthRefreshTokenSpec)\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1)\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t})\n\n\tt.Run(\"if access token is missing and refresh token is present, should run refresh token flow\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.OauthScope = \"session:role:ANALYST offline_access\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.deleteCredential(oauthAccessTokenSpec)\n\t\tcredentialsStorage.setCredential(oauthRefreshTokenSpec, \"refresh-token-123\")\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/refresh_token/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1) // only refresh token\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthRefreshTokenSpec), \"refresh-token-123a\")\n\t})\n\n\tt.Run(\"if access token is expired and refresh token is present, should run refresh token flow\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.OauthScope = \"session:role:ANALYST offline_access\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"expired-token\")\n\t\tcredentialsStorage.setCredential(oauthRefreshTokenSpec, \"refresh-token-123\")\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/refresh_token/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1) // only refresh token\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthRefreshTokenSpec), \"refresh-token-123a\")\n\t})\n\n\tt.Run(\"if new refresh token is not returned, should keep old one\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.OauthScope = \"session:role:ANALYST offline_access\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"expired-token\")\n\t\tcredentialsStorage.setCredential(oauthRefreshTokenSpec, \"refresh-token-123\")\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/refresh_token/successful_flow_without_new_refresh_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1) // only refresh token\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthRefreshTokenSpec), \"refresh-token-123\")\n\t})\n\n\tt.Run(\"if refreshing token failed, run normal flow\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.OauthScope = \"session:role:ANALYST offline_access\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"expired-token\")\n\t\tcredentialsStorage.setCredential(oauthRefreshTokenSpec, \"expired-refresh-token\")\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/refresh_token/invalid_refresh_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/authorization_code/successful_flow_with_offline_access.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 2) // only refresh token fails, then authorization code\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthRefreshTokenSpec), \"refresh-token-123\")\n\t})\n\n\tt.Run(\"if secure storage is disabled, run normal flow\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.OauthScope = \"session:role:ANALYST offline_access\"\n\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\tcfg.OauthRedirectURI = \"http://localhost:1234/snowflake/oauth-redirect\"\n\t\tcfg.Transporter = roundTripper\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolFalse\n\t\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"old-access-token\")\n\t\tcredentialsStorage.setCredential(oauthRefreshTokenSpec, \"old-refresh-token\")\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/oauth2/authorization_code/successful_flow_with_offline_access.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1) // only access token token\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"old-access-token\")\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthRefreshTokenSpec), \"old-refresh-token\")\n\t})\n}\n\nfunc TestClientCredentialsFlow(t *testing.T) {\n\tif runningOnGithubAction() && runningOnLinux() {\n\t\tt.Skip(\"Github blocks writing to file system\")\n\t}\n\tskipOnMac(t, \"keychain requires password\")\n\tcurrentDefaultAuthorizationCodeProviderFactory := defaultAuthorizationCodeProviderFactory\n\tdefer func() {\n\t\tdefaultAuthorizationCodeProviderFactory = currentDefaultAuthorizationCodeProviderFactory\n\t}()\n\tdefaultAuthorizationCodeProviderFactory = func() authorizationCodeProvider {\n\t\treturn &nonInteractiveAuthorizationCodeProvider{\n\t\t\tt: t,\n\t\t\tmu: sync.Mutex{},\n\t\t}\n\t}\n\troundTripper := newCountingRoundTripper(createTestNoRevocationTransport())\n\n\tcfg := wiremock.connectionConfig()\n\tcfg.Role = \"ANALYST\"\n\tcfg.Authenticator = AuthTypeOAuthClientCredentials\n\tcfg.Transporter = roundTripper\n\n\toauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\toauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)\n\n\tt.Run(\"successful flow\", func(t *testing.T) {\n\t\tcredentialsStorage.deleteCredential(oauthAccessTokenSpec)\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t})\n\n\tt.Run(\"should use cached access token\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\t\tcredentialsStorage.deleteCredential(oauthAccessTokenSpec)\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\tconn1, err := db.Conn(context.Background())\n\t\tassertNilF(t, err)\n\t\tdefer conn1.Close()\n\t\tconn2, err := db.Conn(context.Background())\n\t\tassertNilF(t, err)\n\t\tdefer conn2.Close()\n\t\trunSmokeQueryWithConn(t, conn1)\n\t\trunSmokeQueryWithConn(t, conn2)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1)\n\t})\n\n\tt.Run(\"should update cache with new token when the old one expired\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"expired-token\")\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1)\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t})\n\n\tt.Run(\"should not use refresh token, but ask for fresh access token\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"expired-token\")\n\t\tcredentialsStorage.setCredential(oauthRefreshTokenSpec, \"refresh-token-123\")\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1)\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthRefreshTokenSpec), \"refresh-token-123\")\n\t})\n\n\tt.Run(\"should not use access token if token cache is disabled\", func(t *testing.T) {\n\t\troundTripper.reset()\n\t\twiremock.registerMappings(t,\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request_with_expired_access_token.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/client_credentials/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"auth/oauth2/login_request.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"))\n\n\t\tcredentialsStorage.setCredential(oauthAccessTokenSpec, \"access-token-123\")\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolFalse\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tdb := sql.OpenDB(connector)\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, roundTripper.postReqCount[cfg.OauthTokenRequestURL], 1)\n\t\tassertEqualE(t, credentialsStorage.getCredential(oauthAccessTokenSpec), \"access-token-123\")\n\t})\n}\n\nfunc TestEligibleForDefaultClientCredentials(t *testing.T) {\n\ttests := []struct {\n\t\tname string\n\t\toauthClient *oauthClient\n\t\texpected bool\n\t}{\n\t\t{\n\t\t\tname: \"Client credentials not supplied and Snowflake as IdP\",\n\t\t\toauthClient: &oauthClient{\n\t\t\t\tcfg: &Config{\n\t\t\t\t\tHost: \"example.snowflakecomputing.com\",\n\t\t\t\t\tOauthClientID: \"\",\n\t\t\t\t\tOauthClientSecret: \"\",\n\t\t\t\t\tOauthAuthorizationURL: \"https://example.snowflakecomputing.com/oauth/authorize\",\n\t\t\t\t\tOauthTokenRequestURL: \"https://example.snowflakecomputing.com/oauth/token\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname: \"Client credentials not supplied and empty URLs (defaults to Snowflake)\",\n\t\t\toauthClient: &oauthClient{\n\t\t\t\tcfg: &Config{\n\t\t\t\t\tHost: \"example.snowflakecomputing.com\",\n\t\t\t\t\tOauthClientID: \"\",\n\t\t\t\t\tOauthClientSecret: \"\",\n\t\t\t\t\tOauthAuthorizationURL: \"\",\n\t\t\t\t\tOauthTokenRequestURL: \"\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname: \"Client credentials supplied\",\n\t\t\toauthClient: &oauthClient{\n\t\t\t\tcfg: &Config{\n\t\t\t\t\tHost: \"example.snowflakecomputing.com\",\n\t\t\t\t\tOauthClientID: \"testClientID\",\n\t\t\t\t\tOauthClientSecret: \"testClientSecret\",\n\t\t\t\t\tOauthAuthorizationURL: \"https://example.snowflakecomputing.com/oauth/authorize\",\n\t\t\t\t\tOauthTokenRequestURL: \"https://example.snowflakecomputing.com/oauth/token\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname: \"Only client ID supplied\",\n\t\t\toauthClient: &oauthClient{\n\t\t\t\tcfg: &Config{\n\t\t\t\t\tHost: \"example.snowflakecomputing.com\",\n\t\t\t\t\tOauthClientID: \"testClientID\",\n\t\t\t\t\tOauthClientSecret: \"\",\n\t\t\t\t\tOauthAuthorizationURL: \"https://example.snowflakecomputing.com/oauth/authorize\",\n\t\t\t\t\tOauthTokenRequestURL: \"https://example.snowflakecomputing.com/oauth/token\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname: \"Non-Snowflake IdP\",\n\t\t\toauthClient: &oauthClient{\n\t\t\t\tcfg: &Config{\n\t\t\t\t\tHost: \"example.snowflakecomputing.com\",\n\t\t\t\t\tOauthClientID: \"\",\n\t\t\t\t\tOauthClientSecret: \"\",\n\t\t\t\t\tOauthAuthorizationURL: \"https://example.com/oauth/authorize\",\n\t\t\t\t\tOauthTokenRequestURL: \"https://example.com/oauth/token\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpected: false,\n\t\t},\n\t}\n\n\tfor _, test := range tests {\n\t\tt.Run(test.name, func(t *testing.T) {\n\t\t\tresult := test.oauthClient.eligibleForDefaultClientCredentials()\n\t\t\tif result != test.expected {\n\t\t\t\tt.Errorf(\"expected %v, got %v\", test.expected, result)\n\t\t\t}\n\t\t})\n\t}\n}\n\ntype nonInteractiveAuthorizationCodeProvider struct {\n\tt *testing.T\n\ttamperWithState bool\n\ttriggerError string\n\tresponseBody string\n\tmu sync.Mutex\n\tsleepTime time.Duration\n}\n\nfunc (provider *nonInteractiveAuthorizationCodeProvider) run(authorizationURL string) error {\n\tif provider.sleepTime != 0 {\n\t\ttime.Sleep(provider.sleepTime)\n\t\tif provider.triggerError != \"\" {\n\t\t\treturn errors.New(provider.triggerError)\n\t\t}\n\t}\n\tif provider.triggerError != \"\" {\n\t\treturn errors.New(provider.triggerError)\n\t}\n\tgo func() {\n\t\tresp, err := http.Get(authorizationURL)\n\t\tassertNilF(provider.t, err)\n\t\tassertEqualE(provider.t, resp.StatusCode, http.StatusOK)\n\t\trespBody, err := io.ReadAll(resp.Body)\n\t\tassertNilF(provider.t, err)\n\t\tprovider.mu.Lock()\n\t\tdefer provider.mu.Unlock()\n\t\tprovider.responseBody = string(respBody)\n\t}()\n\treturn nil\n}\n\nfunc (provider *nonInteractiveAuthorizationCodeProvider) createState() string {\n\tif provider.tamperWithState {\n\t\treturn \"invalidState\"\n\t}\n\treturn \"testState\"\n}\n\nfunc (provider *nonInteractiveAuthorizationCodeProvider) createCodeVerifier() string {\n\treturn \"testCodeVerifier\"\n}\n\nfunc (provider *nonInteractiveAuthorizationCodeProvider) assertResponseBodyContains(str string) {\n\tprovider.mu.Lock()\n\tdefer provider.mu.Unlock()\n\tassertStringContainsE(provider.t, provider.responseBody, str)\n}\n"
},
{
"path": "auth_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"cmp\"\n\t\"context\"\n\t\"crypto/rand\"\n\t\"crypto/rsa\"\n\t\"database/sql\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\tsfconfig \"github.com/snowflakedb/gosnowflake/v2/internal/config\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"runtime\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/golang-jwt/jwt/v5\"\n)\n\nfunc TestUnitPostAuth(t *testing.T) {\n\tsr := &snowflakeRestful{\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t\tFuncAuthPost: postAuthTestAfterRenew,\n\t}\n\tvar err error\n\tbodyCreator := func() ([]byte, error) {\n\t\treturn []byte{0x12, 0x34}, nil\n\t}\n\t_, err = postAuth(context.Background(), sr, sr.Client, &url.Values{}, make(map[string]string), bodyCreator, 0)\n\tif err != nil {\n\t\tt.Fatalf(\"err: %v\", err)\n\t}\n\tsr.FuncAuthPost = postAuthTestError\n\t_, err = postAuth(context.Background(), sr, sr.Client, &url.Values{}, make(map[string]string), bodyCreator, 0)\n\tif err == nil {\n\t\tt.Fatal(\"should have failed to auth for unknown reason\")\n\t}\n\tsr.FuncAuthPost = postAuthTestAppBadGatewayError\n\t_, err = postAuth(context.Background(), sr, sr.Client, &url.Values{}, make(map[string]string), bodyCreator, 0)\n\tif err == nil {\n\t\tt.Fatal(\"should have failed to auth for unknown reason\")\n\t}\n\tsr.FuncAuthPost = postAuthTestAppForbiddenError\n\t_, err = postAuth(context.Background(), sr, sr.Client, &url.Values{}, make(map[string]string), bodyCreator, 0)\n\tif err == nil {\n\t\tt.Fatal(\"should have failed to auth for unknown reason\")\n\t}\n\tsr.FuncAuthPost = postAuthTestAppUnexpectedError\n\t_, err = postAuth(context.Background(), sr, sr.Client, &url.Values{}, make(map[string]string), bodyCreator, 0)\n\tif err == nil {\n\t\tt.Fatal(\"should have failed to auth for unknown reason\")\n\t}\n}\n\nfunc postAuthFailServiceIssue(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\treturn nil, &SnowflakeError{\n\t\tNumber: ErrCodeServiceUnavailable,\n\t}\n}\n\nfunc postAuthFailWrongAccount(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\treturn nil, &SnowflakeError{\n\t\tNumber: ErrCodeFailedToConnect,\n\t}\n}\n\nfunc postAuthFailUnknown(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\treturn nil, &SnowflakeError{\n\t\tNumber: ErrFailedToAuth,\n\t}\n}\n\nfunc postAuthSuccessWithErrorCode(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\treturn &authResponse{\n\t\tSuccess: false,\n\t\tCode: \"98765\",\n\t\tMessage: \"wrong!\",\n\t}, nil\n}\n\nfunc postAuthSuccessWithInvalidErrorCode(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\treturn &authResponse{\n\t\tSuccess: false,\n\t\tCode: \"abcdef\",\n\t\tMessage: \"wrong!\",\n\t}, nil\n}\n\nfunc postAuthSuccess(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckSAMLResponse(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, err := bodyCreator()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif err = json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\tif ar.Data.RawSAMLResponse == \"\" {\n\t\treturn nil, errors.New(\"SAML response is empty\")\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\n// Checks that the request body generated when authenticating with OAuth\n// contains all the necessary values.\nfunc postAuthCheckOAuth(\n\t_ context.Context,\n\t_ *snowflakeRestful,\n\t_ *http.Client,\n\t_ *url.Values, _ map[string]string,\n\tbodyCreator bodyCreatorType,\n\t_ time.Duration,\n) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\tif ar.Data.Authenticator != AuthTypeOAuth.String() {\n\t\treturn nil, errors.New(\"Authenticator is not OAUTH\")\n\t}\n\tif ar.Data.Token == \"\" {\n\t\treturn nil, errors.New(\"Token is empty\")\n\t}\n\tif ar.Data.LoginName == \"\" {\n\t\treturn nil, errors.New(\"Login name is empty\")\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckPasscode(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\tif ar.Data.Passcode != \"987654321\" || ar.Data.ExtAuthnDuoMethod != \"passcode\" {\n\t\treturn nil, fmt.Errorf(\"passcode didn't match. expected: 987654321, got: %v, duo: %v\", ar.Data.Passcode, ar.Data.ExtAuthnDuoMethod)\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckPasscodeInPassword(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\tif ar.Data.Passcode != \"\" || ar.Data.ExtAuthnDuoMethod != \"passcode\" {\n\t\treturn nil, fmt.Errorf(\"passcode must be empty, got: %v, duo: %v\", ar.Data.Passcode, ar.Data.ExtAuthnDuoMethod)\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckUsernamePasswordMfa(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif ar.Data.SessionParameters[\"CLIENT_REQUEST_MFA_TOKEN\"] != true {\n\t\treturn nil, fmt.Errorf(\"expected client_request_mfa_token to be true but was %v\", ar.Data.SessionParameters[\"CLIENT_REQUEST_MFA_TOKEN\"])\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tMfaToken: \"mockedMfaToken\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckUsernamePasswordMfaToken(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif ar.Data.Token != \"mockedMfaToken\" {\n\t\treturn nil, fmt.Errorf(\"unexpected mfa token: %v\", ar.Data.Token)\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tMfaToken: \"mockedMfaToken\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckUsernamePasswordMfaFailed(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif ar.Data.Token != \"mockedMfaToken\" {\n\t\treturn nil, fmt.Errorf(\"unexpected mfa token: %v\", ar.Data.Token)\n\t}\n\treturn &authResponse{\n\t\tSuccess: false,\n\t\tData: authResponseMain{},\n\t\tMessage: \"auth failed\",\n\t\tCode: \"260008\",\n\t}, nil\n}\n\nfunc postAuthCheckExternalBrowser(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif ar.Data.SessionParameters[\"CLIENT_STORE_TEMPORARY_CREDENTIAL\"] != true {\n\t\treturn nil, fmt.Errorf(\"expected client_store_temporary_credential to be true but was %v\", ar.Data.SessionParameters[\"CLIENT_STORE_TEMPORARY_CREDENTIAL\"])\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tIDToken: \"mockedIDToken\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckExternalBrowserToken(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif ar.Data.Token != \"mockedIDToken\" {\n\t\treturn nil, fmt.Errorf(\"unexpected mfatoken: %v\", ar.Data.Token)\n\t}\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tIDToken: \"mockedIDToken\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc postAuthCheckExternalBrowserFailed(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\tjsonBody, _ := bodyCreator()\n\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\treturn nil, err\n\t}\n\n\tif ar.Data.SessionParameters[\"CLIENT_STORE_TEMPORARY_CREDENTIAL\"] != true {\n\t\treturn nil, fmt.Errorf(\"expected client_store_temporary_credential to be true but was %v\", ar.Data.SessionParameters[\"CLIENT_STORE_TEMPORARY_CREDENTIAL\"])\n\t}\n\treturn &authResponse{\n\t\tSuccess: false,\n\t\tData: authResponseMain{},\n\t\tMessage: \"auth failed\",\n\t\tCode: \"260008\",\n\t}, nil\n}\n\ntype restfulTestWrapper struct {\n\tt *testing.T\n}\n\nfunc (rtw restfulTestWrapper) postAuthOktaWithNewToken(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\tvar ar authRequest\n\n\tcfg := &Config{\n\t\tAuthenticator: AuthTypeOkta,\n\t}\n\n\t// Retry 3 times and success\n\tclient := &fakeHTTPClient{\n\t\tcnt: 3,\n\t\tsuccess: true,\n\t\tstatusCode: 429,\n\t\tt: rtw.t,\n\t}\n\n\turlPtr, err := url.Parse(\"https://fakeaccountretrylogin.snowflakecomputing.com:443/login-request?request_guid=testguid\")\n\tif err != nil {\n\t\treturn &authResponse{}, err\n\t}\n\n\tbody := func() ([]byte, error) {\n\t\tjsonBody, _ := bodyCreator()\n\t\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn jsonBody, err\n\t}\n\n\t_, err = newRetryHTTP(context.Background(), client, emptyRequest, urlPtr, make(map[string]string), 60*time.Second, 3, defaultTimeProvider, cfg).doPost().setBodyCreator(body).execute()\n\tif err != nil {\n\t\treturn &authResponse{}, err\n\t}\n\n\treturn &authResponse{\n\t\tSuccess: true,\n\t\tData: authResponseMain{\n\t\t\tToken: \"t\",\n\t\t\tMasterToken: \"m\",\n\t\t\tMfaToken: \"mockedMfaToken\",\n\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t},\n\t\t},\n\t}, nil\n}\n\nfunc getDefaultSnowflakeConn() *snowflakeConn {\n\tsc := &snowflakeConn{\n\t\trest: &snowflakeRestful{\n\t\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t\t},\n\t\tcfg: &Config{\n\t\t\tAccount: \"a\",\n\t\t\tUser: \"u\",\n\t\t\tPassword: \"p\",\n\t\t\tDatabase: \"d\",\n\t\t\tSchema: \"s\",\n\t\t\tWarehouse: \"w\",\n\t\t\tRole: \"r\",\n\t\t\tRegion: \"\",\n\t\t\tPasscodeInPassword: false,\n\t\t\tPasscode: \"\",\n\t\t\tApplication: \"testapp\",\n\t\t},\n\t\ttelemetry: &snowflakeTelemetry{enabled: false},\n\t}\n\treturn sc\n}\n\nfunc TestUnitAuthenticateWithTokenAccessor(t *testing.T) {\n\texpectedSessionID := int64(123)\n\texpectedMasterToken := \"master_token\"\n\texpectedToken := \"auth_token\"\n\n\tta := getSimpleTokenAccessor()\n\tta.SetTokens(expectedToken, expectedMasterToken, expectedSessionID)\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeTokenAccessor\n\tsc.cfg.TokenAccessor = ta\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthFailServiceIssue,\n\t\tTokenAccessor: ta,\n\t}\n\tsc.rest = sr\n\n\t// FuncPostAuth is set to fail, but AuthTypeTokenAccessor should not even make a call to FuncPostAuth\n\tresp, err := authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"should not have failed, err %v\", err)\n\t}\n\n\tif resp.SessionID != expectedSessionID {\n\t\tt.Fatalf(\"Expected session id %v but got %v\", expectedSessionID, resp.SessionID)\n\t}\n\tif resp.Token != expectedToken {\n\t\tt.Fatalf(\"Expected token %v but got %v\", expectedToken, resp.Token)\n\t}\n\tif resp.MasterToken != expectedMasterToken {\n\t\tt.Fatalf(\"Expected master token %v but got %v\", expectedMasterToken, resp.MasterToken)\n\t}\n\tif resp.SessionInfo.DatabaseName != sc.cfg.Database {\n\t\tt.Fatalf(\"Expected database %v but got %v\", sc.cfg.Database, resp.SessionInfo.DatabaseName)\n\t}\n\tif resp.SessionInfo.WarehouseName != sc.cfg.Warehouse {\n\t\tt.Fatalf(\"Expected warehouse %v but got %v\", sc.cfg.Warehouse, resp.SessionInfo.WarehouseName)\n\t}\n\tif resp.SessionInfo.RoleName != sc.cfg.Role {\n\t\tt.Fatalf(\"Expected role %v but got %v\", sc.cfg.Role, resp.SessionInfo.RoleName)\n\t}\n\tif resp.SessionInfo.SchemaName != sc.cfg.Schema {\n\t\tt.Fatalf(\"Expected schema %v but got %v\", sc.cfg.Schema, resp.SessionInfo.SchemaName)\n\t}\n}\n\nfunc TestUnitAuthenticate(t *testing.T) {\n\tvar err error\n\tvar driverErr *SnowflakeError\n\tvar ok bool\n\n\tta := getSimpleTokenAccessor()\n\tsc := getDefaultSnowflakeConn()\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthFailServiceIssue,\n\t\tTokenAccessor: ta,\n\t}\n\tsc.rest = sr\n\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\tdriverErr, ok = err.(*SnowflakeError)\n\tif !ok || driverErr.Number != ErrCodeServiceUnavailable {\n\t\tt.Fatalf(\"Snowflake error is expected. err: %v\", driverErr)\n\t}\n\tsr.FuncPostAuth = postAuthFailWrongAccount\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\tdriverErr, ok = err.(*SnowflakeError)\n\tif !ok || driverErr.Number != ErrCodeFailedToConnect {\n\t\tt.Fatalf(\"Snowflake error is expected. err: %v\", driverErr)\n\t}\n\tsr.FuncPostAuth = postAuthFailUnknown\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\tdriverErr, ok = err.(*SnowflakeError)\n\tif !ok || driverErr.Number != ErrFailedToAuth {\n\t\tt.Fatalf(\"Snowflake error is expected. err: %v\", driverErr)\n\t}\n\tta.SetTokens(\"bad-token\", \"bad-master-token\", 1)\n\tsr.FuncPostAuth = postAuthSuccessWithErrorCode\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\tnewToken, newMasterToken, newSessionID := ta.GetTokens()\n\tif newToken != \"\" || newMasterToken != \"\" || newSessionID != -1 {\n\t\tt.Fatalf(\"failed auth should have reset tokens: %v %v %v\", newToken, newMasterToken, newSessionID)\n\t}\n\tdriverErr, ok = err.(*SnowflakeError)\n\tif !ok || driverErr.Number != 98765 {\n\t\tt.Fatalf(\"Snowflake error is expected. err: %v\", driverErr)\n\t}\n\tta.SetTokens(\"bad-token\", \"bad-master-token\", 1)\n\tsr.FuncPostAuth = postAuthSuccessWithInvalidErrorCode\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\toldToken, oldMasterToken, oldSessionID := ta.GetTokens()\n\tif oldToken != \"\" || oldMasterToken != \"\" || oldSessionID != -1 {\n\t\tt.Fatalf(\"failed auth should have reset tokens: %v %v %v\", oldToken, oldMasterToken, oldSessionID)\n\t}\n\tsr.FuncPostAuth = postAuthSuccess\n\tvar resp *authResponseMain\n\tresp, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to auth. err: %v\", err)\n\t}\n\tif resp.SessionInfo.DatabaseName != \"dbn\" {\n\t\tt.Fatalf(\"failed to get response from auth\")\n\t}\n\tnewToken, newMasterToken, newSessionID = ta.GetTokens()\n\tif newToken == oldToken {\n\t\tt.Fatalf(\"new token was not set: %v\", newToken)\n\t}\n\tif newMasterToken == oldMasterToken {\n\t\tt.Fatalf(\"new master token was not set: %v\", newMasterToken)\n\t}\n\tif newSessionID == oldSessionID {\n\t\tt.Fatalf(\"new session id was not set: %v\", newSessionID)\n\t}\n}\n\nfunc TestUnitAuthenticateSaml(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tProtocol: \"https\",\n\t\tHost: \"abc.com\",\n\t\tPort: 443,\n\t\tFuncPostAuthSAML: postAuthSAMLAuthSuccess,\n\t\tFuncPostAuthOKTA: postAuthOKTASuccess,\n\t\tFuncGetSSO: getSSOSuccess,\n\t\tFuncPostAuth: postAuthCheckSAMLResponse,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeOkta\n\tsc.cfg.OktaURL = &url.URL{\n\t\tScheme: \"https\",\n\t\tHost: \"abc.com\",\n\t}\n\tsc.rest = sr\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tassertNilF(t, err, \"failed to run.\")\n}\n\n// Unit test for OAuth.\nfunc TestUnitAuthenticateOAuth(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthCheckOAuth,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Token = \"oauthToken\"\n\tsc.cfg.Authenticator = AuthTypeOAuth\n\tsc.rest = sr\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n}\n\nfunc TestUnitAuthenticatePasscode(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthCheckPasscode,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Passcode = \"987654321\"\n\tsc.rest = sr\n\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n\tsr.FuncPostAuth = postAuthCheckPasscodeInPassword\n\tsc.rest = sr\n\tsc.cfg.PasscodeInPassword = true\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n}\n\n// Test JWT function in the local environment against the validation function in go\nfunc TestUnitAuthenticateJWT(t *testing.T) {\n\tvar err error\n\n\t// Generate a fresh private key for this unit test only\n\tlocalTestKey, err := rsa.GenerateKey(rand.Reader, 2048)\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to generate test private key: %s\", err.Error())\n\t}\n\n\t// Create custom JWT verification function that uses the local key\n\tpostAuthCheckLocalJWTToken := func(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {\n\t\tvar ar authRequest\n\t\tjsonBody, _ := bodyCreator()\n\t\tif err := json.Unmarshal(jsonBody, &ar); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tif ar.Data.Authenticator != AuthTypeJwt.String() {\n\t\t\treturn nil, errors.New(\"Authenticator is not JWT\")\n\t\t}\n\n\t\ttokenString := ar.Data.Token\n\n\t\t// Validate token using the local test key's public key\n\t\t_, err := jwt.Parse(tokenString, func(token *jwt.Token) (any, error) {\n\t\t\tif _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {\n\t\t\t\treturn nil, fmt.Errorf(\"Unexpected signing method: %v\", token.Header[\"alg\"])\n\t\t\t}\n\t\t\treturn localTestKey.Public(), nil // Use local key for verification\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\treturn &authResponse{\n\t\t\tSuccess: true,\n\t\t\tData: authResponseMain{\n\t\t\t\tToken: \"t\",\n\t\t\t\tMasterToken: \"m\",\n\t\t\t\tSessionInfo: authResponseSessionInfo{\n\t\t\t\t\tDatabaseName: \"dbn\",\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil\n\t}\n\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthCheckLocalJWTToken, // Use local verification function\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeJwt\n\tsc.cfg.JWTExpireTimeout = time.Duration(sfconfig.DefaultJWTTimeout)\n\tsc.cfg.PrivateKey = localTestKey\n\tsc.rest = sr\n\n\t// A valid JWT token should pass\n\tif _, err = authenticate(context.Background(), sc, []byte{}, []byte{}); err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n\n\t// An invalid JWT token should not pass\n\tinvalidPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)\n\tif err != nil {\n\t\tt.Error(err)\n\t}\n\tsc.cfg.PrivateKey = invalidPrivateKey\n\tif _, err = authenticate(context.Background(), sc, []byte{}, []byte{}); err == nil {\n\t\tt.Fatalf(\"invalid token passed\")\n\t}\n}\n\nfunc TestUnitAuthenticateUsernamePasswordMfa(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthCheckUsernamePasswordMfa,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeUsernamePasswordMFA\n\tsc.cfg.ClientRequestMfaToken = ConfigBoolTrue\n\tsc.rest = sr\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n\n\tsr.FuncPostAuth = postAuthCheckUsernamePasswordMfaToken\n\tsc.mfaToken = \"mockedMfaToken\"\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n\n\tsr.FuncPostAuth = postAuthCheckUsernamePasswordMfaFailed\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed\")\n\t}\n}\n\nfunc TestUnitAuthenticateWithConfigMFA(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthCheckUsernamePasswordMfa,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeUsernamePasswordMFA\n\tsc.cfg.ClientRequestMfaToken = ConfigBoolTrue\n\tsc.rest = sr\n\tsc.ctx = context.Background()\n\terr = authenticateWithConfig(sc)\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n}\n\n// This test creates two groups of scenarios:\n// a) singleAuthenticationPrompt=true - in this case, we start authenticating threads at once,\n// but due to locking mechanism only one should reach wiremock without MFA token.\n// b) singleAuthenticationPrompt=false - in this case, there is no locking, so all threads should rush,\n// but on Wiremock only first will be served with correct response (simulating a user confirming MFA only once).\n// The remaining threads should return error.\nfunc TestMfaParallelLogin(t *testing.T) {\n\tskipOnMissingHome(t)\n\tskipOnMac(t, \"interactive keyring access not available on macOS runners\")\n\tcfg := wiremock.connectionConfig()\n\ttokenSpec := newMfaTokenSpec(cfg.Host, cfg.User)\n\n\tfor _, singleAuthenticationPrompt := range []ConfigBool{ConfigBoolTrue, ConfigBoolFalse} {\n\t\tt.Run(\"starts without mfa token, singleAuthenticationPrompt=\"+singleAuthenticationPrompt.String(), func(t *testing.T) {\n\t\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/mfa/parallel_login_successful_flow.json\"),\n\t\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\t\tcfg := wiremock.connectionConfig()\n\t\t\tcfg.Authenticator = AuthTypeUsernamePasswordMFA\n\t\t\tcfg.SingleAuthenticationPrompt = singleAuthenticationPrompt\n\t\t\tcfg.ClientRequestMfaToken = ConfigBoolTrue\n\t\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\t\tdb := sql.OpenDB(connector)\n\t\t\tdefer db.Close()\n\t\t\tcredentialsStorage.deleteCredential(tokenSpec)\n\t\t\terrs := initPoolWithSizeAndReturnErrors(db, 20)\n\t\t\tif singleAuthenticationPrompt == ConfigBoolTrue {\n\t\t\t\tassertEqualE(t, len(errs), 0)\n\t\t\t} else {\n\t\t\t\t// most of for the one that actually retrieves MFA token should fail\n\t\t\t\tassertEqualE(t, len(errs), 19)\n\t\t\t}\n\t\t})\n\n\t\tt.Run(\"starts without mfa token, first attempt fails, singleAuthenticationPrompt=\"+singleAuthenticationPrompt.String(), func(t *testing.T) {\n\t\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/mfa/parallel_login_first_fails_then_successful_flow.json\"),\n\t\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\t\tcfg := wiremock.connectionConfig()\n\t\t\tcfg.Authenticator = AuthTypeUsernamePasswordMFA\n\t\t\tcfg.SingleAuthenticationPrompt = singleAuthenticationPrompt\n\t\t\tcfg.ClientRequestMfaToken = ConfigBoolTrue\n\t\t\tcredentialsStorage.deleteCredential(tokenSpec)\n\t\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\t\tdb := sql.OpenDB(connector)\n\t\t\tdefer db.Close()\n\t\t\terrs := initPoolWithSizeAndReturnErrors(db, 20)\n\t\t\tif singleAuthenticationPrompt == ConfigBoolTrue {\n\t\t\t\tassertEqualF(t, len(errs), 1)\n\t\t\t\tassertStringContainsE(t, errs[0].Error(), \"MFA with TOTP is required\")\n\t\t\t} else {\n\t\t\t\tassertEqualE(t, len(errs), 19)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestUnitAuthenticateWithConfigOkta(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tProtocol: \"https\",\n\t\tHost: \"abc.com\",\n\t\tPort: 443,\n\t\tFuncPostAuthSAML: postAuthSAMLAuthSuccess,\n\t\tFuncPostAuthOKTA: postAuthOKTASuccess,\n\t\tFuncGetSSO: getSSOSuccess,\n\t\tFuncPostAuth: postAuthCheckSAMLResponse,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeOkta\n\tsc.cfg.OktaURL = &url.URL{\n\t\tScheme: \"https\",\n\t\tHost: \"abc.com\",\n\t}\n\tsc.rest = sr\n\tsc.ctx = context.Background()\n\n\terr = authenticateWithConfig(sc)\n\tassertNilE(t, err, \"expected to have no error.\")\n\n\tsr.FuncPostAuthSAML = postAuthSAMLError\n\terr = authenticateWithConfig(sc)\n\tassertNotNilF(t, err, \"should have failed at FuncPostAuthSAML.\")\n\tassertEqualE(t, err.Error(), \"failed to get SAML response\")\n}\n\nfunc TestUnitAuthenticateWithExternalBrowserParallel(t *testing.T) {\n\tskipOnMissingHome(t)\n\tskipOnMac(t, \"interactive keyring access not available on macOS runners\")\n\tt.Run(\"no ID token cached\", func(t *testing.T) {\n\t\torigSamlResponseProvider := defaultSamlResponseProvider\n\t\tdefer func() { defaultSamlResponseProvider = origSamlResponseProvider }()\n\t\tdefaultSamlResponseProvider = func() samlResponseProvider {\n\t\t\treturn &nonInteractiveSamlResponseProvider{t: t}\n\t\t}\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/external_browser/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Authenticator = AuthTypeExternalBrowser\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolTrue\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tcredentialsStorage.deleteCredential(newIDTokenSpec(cfg.Host, cfg.User))\n\t\tdb := sql.OpenDB(connector)\n\t\tdefer db.Close()\n\t\trunSmokeQuery(t, db)\n\t\tassertEqualE(t, credentialsStorage.getCredential(newIDTokenSpec(cfg.Host, cfg.User)), \"test-id-token\")\n\t})\n\n\tt.Run(\"ID token cached\", func(t *testing.T) {\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/external_browser/successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Authenticator = AuthTypeExternalBrowser\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolTrue\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tcredentialsStorage.setCredential(newIDTokenSpec(cfg.Host, cfg.User), \"test-id-token\")\n\t\tdb := sql.OpenDB(connector)\n\t\tdefer db.Close()\n\t\trunSmokeQuery(t, db)\n\t})\n\n\tt.Run(\"first connection retrieves ID token, second request uses cached ID token\", func(t *testing.T) {\n\t\torigSamlResponseProvider := defaultSamlResponseProvider\n\t\tdefer func() { defaultSamlResponseProvider = origSamlResponseProvider }()\n\t\tdefaultSamlResponseProvider = func() samlResponseProvider {\n\t\t\treturn &nonInteractiveSamlResponseProvider{t: t}\n\t\t}\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/external_browser/parallel_login_successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Authenticator = AuthTypeExternalBrowser\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolTrue\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tcredentialsStorage.deleteCredential(newIDTokenSpec(cfg.Host, cfg.User))\n\t\tdb := sql.OpenDB(connector)\n\t\tdefer db.Close()\n\t\tconn1, err := db.Conn(context.Background())\n\t\tassertNilF(t, err)\n\t\tdefer conn1.Close()\n\t\trunSmokeQueryWithConn(t, conn1)\n\t\tconn2, err := db.Conn(context.Background())\n\t\tassertNilF(t, err)\n\t\tdefer conn2.Close()\n\t\trunSmokeQueryWithConn(t, conn2)\n\t})\n\n\tt.Run(\"first connection retrieves ID token, remaining ones wait and reuse\", func(t *testing.T) {\n\t\torigSamlResponseProvider := defaultSamlResponseProvider\n\t\tdefer func() { defaultSamlResponseProvider = origSamlResponseProvider }()\n\t\tdefaultSamlResponseProvider = func() samlResponseProvider {\n\t\t\treturn &nonInteractiveSamlResponseProvider{t: t}\n\t\t}\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/external_browser/parallel_login_successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Authenticator = AuthTypeExternalBrowser\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolTrue\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tcredentialsStorage.deleteCredential(newIDTokenSpec(cfg.Host, cfg.User))\n\t\tdb := sql.OpenDB(connector)\n\t\tdefer db.Close()\n\t\terrs := initPoolWithSizeAndReturnErrors(db, 20)\n\t\tassertEqualE(t, len(errs), 0)\n\t})\n\n\tt.Run(\"first connection fails, second retrieves ID token, remaining ones wait and reuse\", func(t *testing.T) {\n\t\torigSamlResponseProvider := defaultSamlResponseProvider\n\t\tdefer func() { defaultSamlResponseProvider = origSamlResponseProvider }()\n\t\tdefaultSamlResponseProvider = func() samlResponseProvider {\n\t\t\treturn &nonInteractiveSamlResponseProvider{t: t}\n\t\t}\n\t\twiremock.registerMappings(t, newWiremockMapping(\"auth/external_browser/parallel_login_first_fails_then_successful_flow.json\"),\n\t\t\tnewWiremockMapping(\"select1.json\"),\n\t\t\tnewWiremockMapping(\"close_session.json\"))\n\t\tcfg := wiremock.connectionConfig()\n\t\tcfg.Authenticator = AuthTypeExternalBrowser\n\t\tcfg.ClientStoreTemporaryCredential = ConfigBoolTrue\n\t\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\t\tcredentialsStorage.deleteCredential(newIDTokenSpec(cfg.Host, cfg.User))\n\t\tdb := sql.OpenDB(connector)\n\t\tdefer db.Close()\n\t\terrs := initPoolWithSizeAndReturnErrors(db, 20)\n\t\tassertEqualE(t, len(errs), 1)\n\t})\n}\n\nfunc TestUnitAuthenticateWithConfigExternalBrowserWithFailedSAMLResponse(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuthSAML: postAuthSAMLError,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeExternalBrowser\n\tsc.cfg.ExternalBrowserTimeout = time.Duration(sfconfig.DefaultExternalBrowserTimeout)\n\tsc.rest = sr\n\tsc.ctx = context.Background()\n\terr = authenticateWithConfig(sc)\n\tassertNotNilF(t, err, \"should have failed at FuncPostAuthSAML.\")\n\tassertEqualE(t, err.Error(), \"failed to get SAML response\")\n}\n\nfunc TestUnitAuthenticateExternalBrowser(t *testing.T) {\n\tvar err error\n\tsr := &snowflakeRestful{\n\t\tFuncPostAuth: postAuthCheckExternalBrowser,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeExternalBrowser\n\tsc.cfg.ClientStoreTemporaryCredential = ConfigBoolTrue\n\tsc.rest = sr\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n\n\tsr.FuncPostAuth = postAuthCheckExternalBrowserToken\n\tsc.idToken = \"mockedIDToken\"\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err != nil {\n\t\tt.Fatalf(\"failed to run. err: %v\", err)\n\t}\n\n\tsr.FuncPostAuth = postAuthCheckExternalBrowserFailed\n\t_, err = authenticate(context.Background(), sc, []byte{}, []byte{})\n\tif err == nil {\n\t\tt.Fatal(\"should have failed\")\n\t}\n}\n\n// To run this test you need to set environment variables in parameters.json to a user with MFA authentication enabled\n// Set any other snowflake_test variables needed for database, schema, role for this user\nfunc TestUsernamePasswordMfaCaching(t *testing.T) {\n\tt.Skip(\"manual test for MFA token caching\")\n\n\tconfig, err := ParseDSN(dsn)\n\tif err != nil {\n\t\tt.Fatal(\"Failed to parse dsn\")\n\t}\n\t// connect with MFA authentication\n\tuser := os.Getenv(\"SNOWFLAKE_TEST_MFA_USER\")\n\tpassword := os.Getenv(\"SNOWFLAKE_TEST_MFA_PASSWORD\")\n\tconfig.User = user\n\tconfig.Password = password\n\tconfig.Authenticator = AuthTypeUsernamePasswordMFA\n\tif runtime.GOOS == \"linux\" {\n\t\tconfig.ClientRequestMfaToken = ConfigBoolTrue\n\t}\n\tconnector := NewConnector(SnowflakeDriver{}, *config)\n\tdb := sql.OpenDB(connector)\n\tfor range 3 {\n\t\t// should only be prompted to authenticate first time around.\n\t\t_, err := db.Query(\"select current_user()\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t}\n}\n\nfunc TestUsernamePasswordMfaCachingWithPasscode(t *testing.T) {\n\tt.Skip(\"manual test for MFA token caching\")\n\n\tconfig, err := ParseDSN(dsn)\n\tif err != nil {\n\t\tt.Fatal(\"Failed to parse dsn\")\n\t}\n\t// connect with MFA authentication\n\tuser := os.Getenv(\"SNOWFLAKE_TEST_MFA_USER\")\n\tpassword := os.Getenv(\"SNOWFLAKE_TEST_MFA_PASSWORD\")\n\tconfig.User = user\n\tconfig.Password = password\n\tconfig.Passcode = \"\" // fill with your passcode from DUO app\n\tconfig.Authenticator = AuthTypeUsernamePasswordMFA\n\tif runtime.GOOS == \"linux\" {\n\t\tconfig.ClientRequestMfaToken = ConfigBoolTrue\n\t}\n\tconnector := NewConnector(SnowflakeDriver{}, *config)\n\tdb := sql.OpenDB(connector)\n\tfor range 3 {\n\t\t// should only be prompted to authenticate first time around.\n\t\t_, err := db.Query(\"select current_user()\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t}\n}\n\nfunc TestUsernamePasswordMfaCachingWithPasscodeInPassword(t *testing.T) {\n\tt.Skip(\"manual test for MFA token caching\")\n\n\tconfig, err := ParseDSN(dsn)\n\tif err != nil {\n\t\tt.Fatal(\"Failed to parse dsn\")\n\t}\n\t// connect with MFA authentication\n\tuser := os.Getenv(\"SNOWFLAKE_TEST_MFA_USER\")\n\tpassword := os.Getenv(\"SNOWFLAKE_TEST_MFA_PASSWORD\")\n\tconfig.User = user\n\tconfig.Password = password + \"\" // fill with your passcode from DUO app\n\tconfig.PasscodeInPassword = true\n\tconnector := NewConnector(SnowflakeDriver{}, *config)\n\tdb := sql.OpenDB(connector)\n\tfor range 3 {\n\t\t// should only be prompted to authenticate first time around.\n\t\t_, err := db.Query(\"select current_user()\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t}\n}\n\n// To run this test you need to set environment variables in parameters.json to a user with MFA authentication enabled\n// Set any other snowflake_test variables needed for database, schema, role for this user\nfunc TestDisableUsernamePasswordMfaCaching(t *testing.T) {\n\tt.Skip(\"manual test for disabling MFA token caching\")\n\n\tconfig, err := ParseDSN(dsn)\n\tif err != nil {\n\t\tt.Fatal(\"Failed to parse dsn\")\n\t}\n\t// connect with MFA authentication\n\tuser := os.Getenv(\"SNOWFLAKE_TEST_MFA_USER\")\n\tpassword := os.Getenv(\"SNOWFLAKE_TEST_MFA_PASSWORD\")\n\tconfig.User = user\n\tconfig.Password = password\n\tconfig.Authenticator = AuthTypeUsernamePasswordMFA\n\t// disable MFA token caching\n\tconfig.ClientRequestMfaToken = ConfigBoolFalse\n\tconnector := NewConnector(SnowflakeDriver{}, *config)\n\tdb := sql.OpenDB(connector)\n\tfor range 3 {\n\t\t// should be prompted to authenticate 3 times.\n\t\t_, err := db.Query(\"select current_user()\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t}\n}\n\n// To run this test you need to set SNOWFLAKE_TEST_EXT_BROWSER_USER environment variable to an external browser user\n// Set any other snowflake_test variables needed for database, schema, role for this user\nfunc TestExternalBrowserCaching(t *testing.T) {\n\tt.Skip(\"manual test for external browser token caching\")\n\n\tconfig, err := ParseDSN(dsn)\n\tif err != nil {\n\t\tt.Fatal(\"Failed to parse dsn\")\n\t}\n\t// connect with external browser authentication\n\tuser := os.Getenv(\"SNOWFLAKE_TEST_EXT_BROWSER_USER\")\n\tconfig.User = user\n\tconfig.Authenticator = AuthTypeExternalBrowser\n\tif runtime.GOOS == \"linux\" {\n\t\tconfig.ClientStoreTemporaryCredential = ConfigBoolTrue\n\t}\n\tconnector := NewConnector(SnowflakeDriver{}, *config)\n\tdb := sql.OpenDB(connector)\n\tfor range 3 {\n\t\t// should only be prompted to authenticate first time around.\n\t\t_, err := db.Query(\"select current_user()\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t}\n}\n\n// To run this test you need to set SNOWFLAKE_TEST_EXT_BROWSER_USER environment variable to an external browser user\n// Set any other snowflake_test variables needed for database, schema, role for this user\nfunc TestDisableExternalBrowserCaching(t *testing.T) {\n\tt.Skip(\"manual test for disabling external browser token caching\")\n\n\tconfig, err := ParseDSN(dsn)\n\tif err != nil {\n\t\tt.Fatal(\"Failed to parse dsn\")\n\t}\n\t// connect with external browser authentication\n\tuser := os.Getenv(\"SNOWFLAKE_TEST_EXT_BROWSER_USER\")\n\tconfig.User = user\n\tconfig.Authenticator = AuthTypeExternalBrowser\n\t// disable external browser token caching\n\tconfig.ClientStoreTemporaryCredential = ConfigBoolFalse\n\tconnector := NewConnector(SnowflakeDriver{}, *config)\n\tdb := sql.OpenDB(connector)\n\tfor range 3 {\n\t\t// should be prompted to authenticate 3 times.\n\t\t_, err := db.Query(\"select current_user()\")\n\t\tif err != nil {\n\t\t\tt.Fatal(err)\n\t\t}\n\t}\n}\n\nfunc TestOktaRetryWithNewToken(t *testing.T) {\n\texpectedMasterToken := \"m\"\n\texpectedToken := \"t\"\n\texpectedMfaToken := \"mockedMfaToken\"\n\texpectedDatabaseName := \"dbn\"\n\n\tsr := &snowflakeRestful{\n\t\tProtocol: \"https\",\n\t\tHost: \"abc.com\",\n\t\tPort: 443,\n\t\tFuncPostAuthSAML: postAuthSAMLAuthSuccess,\n\t\tFuncPostAuthOKTA: postAuthOKTASuccess,\n\t\tFuncGetSSO: getSSOSuccess,\n\t\tFuncPostAuth: restfulTestWrapper{t: t}.postAuthOktaWithNewToken,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\tsc := getDefaultSnowflakeConn()\n\tsc.cfg.Authenticator = AuthTypeOkta\n\tsc.cfg.OktaURL = &url.URL{\n\t\tScheme: \"https\",\n\t\tHost: \"abc.com\",\n\t}\n\tsc.rest = sr\n\tsc.ctx = context.Background()\n\n\tauthResponse, err := authenticate(context.Background(), sc, []byte{0x12, 0x34}, []byte{0x56, 0x78})\n\tassertNilF(t, err, \"should not have failed to run authenticate()\")\n\tassertEqualF(t, authResponse.MasterToken, expectedMasterToken)\n\tassertEqualF(t, authResponse.Token, expectedToken)\n\tassertEqualF(t, authResponse.MfaToken, expectedMfaToken)\n\tassertEqualF(t, authResponse.SessionInfo.DatabaseName, expectedDatabaseName)\n}\n\nfunc TestContextPropagatedToAuthWhenUsingOpen(t *testing.T) {\n\tdb, err := sql.Open(\"snowflake\", dsn)\n\tassertNilF(t, err)\n\tdefer db.Close()\n\n\tctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)\n\t_, err = db.QueryContext(ctx, \"SELECT 1\")\n\tassertNotNilF(t, err)\n\tassertStringContainsE(t, err.Error(), \"context deadline exceeded\")\n\tcancel()\n}\n\nfunc TestContextPropagatedToAuthWhenUsingOpenDB(t *testing.T) {\n\tcfg, err := ParseDSN(dsn)\n\tassertNilF(t, err)\n\tconnector := NewConnector(&SnowflakeDriver{}, *cfg)\n\tdb := sql.OpenDB(connector)\n\tdefer db.Close()\n\n\tctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)\n\t_, err = db.QueryContext(ctx, \"SELECT 1\")\n\tassertNotNilF(t, err)\n\tassertStringContainsE(t, err.Error(), \"context deadline exceeded\")\n\tcancel()\n}\n\nfunc TestPatSuccessfulFlow(t *testing.T) {\n\tcfg := wiremock.connectionConfig()\n\tcfg.Authenticator = AuthTypePat\n\tcfg.Token = \"some PAT\"\n\twiremock.registerMappings(t,\n\t\twiremockMapping{filePath: \"auth/pat/successful_flow.json\"},\n\t\twiremockMapping{filePath: \"select1.json\"},\n\t)\n\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\tdb := sql.OpenDB(connector)\n\trows, err := db.Query(\"SELECT 1\")\n\tassertNilF(t, err)\n\tvar v int\n\tassertTrueE(t, rows.Next())\n\tassertNilF(t, rows.Scan(&v))\n\tassertEqualE(t, v, 1)\n}\n\nfunc TestPatTokenRotation(t *testing.T) {\n\tdir := t.TempDir()\n\ttokenFilePath := filepath.Join(dir, \"tokenFile\")\n\tassertNilF(t, os.WriteFile(tokenFilePath, []byte(\"some PAT\"), 0644))\n\n\tcfg := wiremock.connectionConfig()\n\tcfg.Authenticator = AuthTypePat\n\tcfg.TokenFilePath = tokenFilePath\n\twiremock.registerMappings(t,\n\t\twiremockMapping{filePath: \"auth/pat/reading_fresh_token.json\"},\n\t)\n\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\tdb := sql.OpenDB(connector)\n\t_, err := db.Conn(context.Background())\n\tassertNilF(t, err)\n\n\tassertNilF(t, os.WriteFile(tokenFilePath, []byte(\"some PAT 2\"), 0644))\n\t_, err = db.Conn(context.Background())\n\tassertNilF(t, err)\n}\n\nfunc TestPatInvalidToken(t *testing.T) {\n\twiremock.registerMappings(t,\n\t\twiremockMapping{filePath: \"auth/pat/invalid_token.json\"},\n\t)\n\tcfg := wiremock.connectionConfig()\n\tcfg.Authenticator = AuthTypePat\n\tcfg.Token = \"some PAT\"\n\tconnector := NewConnector(SnowflakeDriver{}, *cfg)\n\tdb := sql.OpenDB(connector)\n\t_, err := db.Query(\"SELECT 1\")\n\tassertNotNilF(t, err)\n\tvar se *SnowflakeError\n\tassertErrorsAsF(t, err, &se)\n\tassertEqualE(t, se.Number, 394400)\n\tassertEqualE(t, se.Message, \"Programmatic access token is invalid.\")\n}\n\nfunc TestWithOauthAuthorizationCodeFlowManual(t *testing.T) {\n\tt.Skip(\"manual test\")\n\tfor _, provider := range []string{\"OKTA\", \"SNOWFLAKE\"} {\n\t\tt.Run(provider, func(t *testing.T) {\n\t\t\tcfg, err := GetConfigFromEnv([]*ConfigParam{\n\t\t\t\t{Name: \"OAuthClientId\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_CLIENT_ID\", FailOnMissing: true},\n\t\t\t\t{Name: \"OAuthClientSecret\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_CLIENT_SECRET\", FailOnMissing: true},\n\t\t\t\t{Name: \"OAuthAuthorizationURL\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_AUTHORIZATION_URL\", FailOnMissing: false},\n\t\t\t\t{Name: \"OAuthTokenRequestURL\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_TOKEN_REQUEST_URL\", FailOnMissing: false},\n\t\t\t\t{Name: \"OAuthRedirectURI\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_REDIRECT_URI\", FailOnMissing: false},\n\t\t\t\t{Name: \"OAuthScope\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_SCOPE\", FailOnMissing: false},\n\t\t\t\t{Name: \"User\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_USER\", FailOnMissing: true},\n\t\t\t\t{Name: \"Role\", EnvName: \"SNOWFLAKE_TEST_OAUTH_\" + provider + \"_ROLE\", FailOnMissing: true},\n\t\t\t\t{Name: \"Account\", EnvName: \"SNOWFLAKE_TEST_ACCOUNT\", FailOnMissing: true},\n\t\t\t})\n\t\t\tassertNilF(t, err)\n\t\t\tcfg.Authenticator = AuthTypeOAuthAuthorizationCode\n\t\t\ttokenRequestURL := cmp.Or(cfg.OauthTokenRequestURL, fmt.Sprintf(\"https://%v.snowflakecomputing.com:443/oauth/token-request\", cfg.Account))\n\t\t\tcredentialsStorage.deleteCredential(newOAuthAccessTokenSpec(tokenRequestURL, cfg.User))\n\t\t\tcredentialsStorage.deleteCredential(newOAuthRefreshTokenSpec(tokenRequestURL, cfg.User))\n\t\t\tconnector := NewConnector(&SnowflakeDriver{}, *cfg)\n\t\t\tdb := sql.OpenDB(connector)\n\t\t\tdefer db.Close()\n\t\t\tconn1, err := db.Conn(context.Background())\n\t\t\tassertNilF(t, err)\n\t\t\tdefer conn1.Close()\n\t\t\trunSmokeQueryWithConn(t, conn1)\n\t\t\tconn2, err := db.Conn(context.Background())\n\t\t\tassertNilF(t, err)\n\t\t\tdefer conn2.Close()\n\t\t\trunSmokeQueryWithConn(t, conn2)\n\t\t\tcredentialsStorage.setCredential(newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User), \"expired-token\")\n\t\t\tconn3, err := db.Conn(context.Background())\n\t\t\tassertNilF(t, err)\n\t\t\tdefer conn3.Close()\n\t\t\trunSmokeQueryWithConn(t, conn3)\n\t\t})\n\t}\n}\n\nfunc TestWithOAuthClientCredentialsFlowManual(t *testing.T) {\n\tt.Skip(\"manual test\")\n\tcfg, err := GetConfigFromEnv([]*ConfigParam{\n\t\t{Name: \"OAuthClientId\", EnvName: \"SNOWFLAKE_TEST_OAUTH_OKTA_CLIENT_ID\", FailOnMissing: true},\n\t\t{Name: \"OAuthClientSecret\", EnvName: \"SNOWFLAKE_TEST_OAUTH_OKTA_CLIENT_SECRET\", FailOnMissing: true},\n\t\t{Name: \"OAuthTokenRequestURL\", EnvName: \"SNOWFLAKE_TEST_OAUTH_OKTA_TOKEN_REQUEST_URL\", FailOnMissing: true},\n\t\t{Name: \"Role\", EnvName: \"SNOWFLAKE_TEST_OAUTH_OKTA_ROLE\", FailOnMissing: true},\n\t\t{Name: \"Account\", EnvName: \"SNOWFLAKE_TEST_ACCOUNT\", FailOnMissing: true},\n\t})\n\tassertNilF(t, err)\n\tcfg.Authenticator = AuthTypeOAuthClientCredentials\n\tconnector := NewConnector(&SnowflakeDriver{}, *cfg)\n\tdb := sql.OpenDB(connector)\n\tdefer db.Close()\n\trunSmokeQuery(t, db)\n}\n"
},
{
"path": "auth_wif.go",
"content": "package gosnowflake\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"crypto/sha256\"\n\t\"encoding/base64\"\n\t\"encoding/hex\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"net/http\"\n\t\"os\"\n\t\"strings\"\n\t\"time\"\n\n\t\"github.com/aws/aws-sdk-go-v2/aws\"\n\tv4 \"github.com/aws/aws-sdk-go-v2/aws/signer/v4\"\n\t\"github.com/aws/aws-sdk-go-v2/config\"\n\t\"github.com/aws/aws-sdk-go-v2/credentials\"\n\t\"github.com/aws/aws-sdk-go-v2/service/sts\"\n\t\"github.com/golang-jwt/jwt/v5\"\n\tsfconfig \"github.com/snowflakedb/gosnowflake/v2/internal/config\"\n)\n\nconst (\n\tawsWif wifProviderType = \"AWS\"\n\tgcpWif wifProviderType = \"GCP\"\n\tazureWif wifProviderType = \"AZURE\"\n\toidcWif wifProviderType = \"OIDC\"\n\n\tgcpMetadataFlavorHeaderName = \"Metadata-Flavor\"\n\tgcpMetadataFlavor = \"Google\"\n\tdefaultMetadataServiceBase = \"http://169.254.169.254\"\n\tdefaultGcpIamCredentialsBase = \"https://iamcredentials.googleapis.com\"\n\tsnowflakeAudience = \"snowflakecomputing.com\"\n)\n\ntype wifProviderType string\n\ntype wifAttestation struct {\n\tProviderType string `json:\"providerType\"`\n\tCredential string `json:\"credential\"`\n\tMetadata map[string]string `json:\"metadata\"`\n}\n\ntype wifAttestationCreator interface {\n\tcreateAttestation() (*wifAttestation, error)\n}\n\ntype wifAttestationProvider struct {\n\tcontext context.Context\n\tcfg *Config\n\tawsCreator wifAttestationCreator\n\tgcpCreator wifAttestationCreator\n\tazureCreator wifAttestationCreator\n\toidcCreator wifAttestationCreator\n}\n\nfunc createWifAttestationProvider(ctx context.Context, cfg *Config, telemetry *snowflakeTelemetry) *wifAttestationProvider {\n\treturn &wifAttestationProvider{\n\t\tcontext: ctx,\n\t\tcfg: cfg,\n\t\tawsCreator: &awsIdentityAttestationCreator{\n\t\t\tcfg: cfg,\n\t\t\tattestationServiceFactory: createDefaultAwsAttestationMetadataProvider,\n\t\t\tctx: ctx,\n\t\t},\n\t\tgcpCreator: &gcpIdentityAttestationCreator{\n\t\t\tcfg: cfg,\n\t\t\ttelemetry: telemetry,\n\t\t\tmetadataServiceBaseURL: defaultMetadataServiceBase,\n\t\t\tiamCredentialsURL: defaultGcpIamCredentialsBase,\n\t\t},\n\t\tazureCreator: &azureIdentityAttestationCreator{\n\t\t\tazureAttestationMetadataProvider: &defaultAzureAttestationMetadataProvider{},\n\t\t\tcfg: cfg,\n\t\t\ttelemetry: telemetry,\n\t\t\tworkloadIdentityEntraResource: determineEntraResource(cfg),\n\t\t\tazureMetadataServiceBaseURL: defaultMetadataServiceBase,\n\t\t},\n\t\toidcCreator: &oidcIdentityAttestationCreator{token: func() (string, error) { return sfconfig.GetToken(cfg) }},\n\t}\n}\n\nfunc (p *wifAttestationProvider) getAttestation(identityProvider string) (*wifAttestation, error) {\n\tswitch strings.ToUpper(identityProvider) {\n\tcase string(awsWif):\n\t\treturn p.awsCreator.createAttestation()\n\tcase string(gcpWif):\n\t\treturn p.gcpCreator.createAttestation()\n\tcase string(azureWif):\n\t\treturn p.azureCreator.createAttestation()\n\tcase string(oidcWif):\n\t\treturn p.oidcCreator.createAttestation()\n\tdefault:\n\t\treturn nil, fmt.Errorf(\"unknown WorkloadIdentityProvider specified: %s. Valid values are: %s, %s, %s, %s\", identityProvider, awsWif, gcpWif, azureWif, oidcWif)\n\t}\n}\n\ntype awsAttestastationMetadataProviderFactory func(ctx context.Context, cfg *Config) awsAttestationMetadataProvider\n\ntype awsIdentityAttestationCreator struct {\n\tcfg *Config\n\tattestationServiceFactory awsAttestastationMetadataProviderFactory\n\tctx context.Context\n}\n\ntype gcpIdentityAttestationCreator struct {\n\tcfg *Config\n\ttelemetry *snowflakeTelemetry\n\tmetadataServiceBaseURL string\n\tiamCredentialsURL string\n}\n\ntype oidcIdentityAttestationCreator struct {\n\ttoken func() (string, error)\n}\n\ntype awsAttestationMetadataProvider interface {\n\tawsCredentials() (aws.Credentials, error)\n\tawsCredentialsViaRoleChaining() (aws.Credentials, error)\n\tawsRegion() string\n}\n\ntype defaultAwsAttestationMetadataProvider struct {\n\tctx context.Context\n\tcfg *Config\n\tawsCfg aws.Config\n}\n\nfunc createDefaultAwsAttestationMetadataProvider(ctx context.Context, cfg *Config) awsAttestationMetadataProvider {\n\tawsCfg, err := config.LoadDefaultConfig(ctx, config.WithEC2IMDSRegion())\n\tif err != nil {\n\t\tlogger.Debugf(\"Unable to load AWS config: %v\", err)\n\t\treturn nil\n\t}\n\treturn &defaultAwsAttestationMetadataProvider{\n\t\tawsCfg: awsCfg,\n\t\tcfg: cfg,\n\t\tctx: ctx,\n\t}\n}\n\nfunc (s *defaultAwsAttestationMetadataProvider) awsCredentials() (aws.Credentials, error) {\n\treturn s.awsCfg.Credentials.Retrieve(s.ctx)\n}\n\nfunc (s *defaultAwsAttestationMetadataProvider) awsCredentialsViaRoleChaining() (aws.Credentials, error) {\n\tcreds, err := s.awsCredentials()\n\tif err != nil {\n\t\treturn aws.Credentials{}, err\n\t}\n\tfor _, roleArn := range s.cfg.WorkloadIdentityImpersonationPath {\n\t\tif creds, err = s.assumeRole(creds, roleArn); err != nil {\n\t\t\treturn aws.Credentials{}, err\n\t\t}\n\t}\n\treturn creds, nil\n}\n\nfunc (s *defaultAwsAttestationMetadataProvider) assumeRole(creds aws.Credentials, roleArn string) (aws.Credentials, error) {\n\tlogger.Debugf(\"assuming role %v\", roleArn)\n\tawsCfg := s.awsCfg\n\tawsCfg.Credentials = credentials.StaticCredentialsProvider{Value: creds}\n\tawsCfg.Region = s.awsRegion()\n\tstsClient := sts.NewFromConfig(awsCfg)\n\n\trole, err := stsClient.AssumeRole(s.ctx, &sts.AssumeRoleInput{\n\t\tRoleArn: aws.String(roleArn),\n\t\tRoleSessionName: aws.String(\"identity-federation-session\"),\n\t})\n\tif err != nil {\n\t\tlogger.Debugf(\"failed to assume role %v: %v\", roleArn, err)\n\t\treturn aws.Credentials{}, err\n\t}\n\n\treturn aws.Credentials{\n\t\tAccessKeyID: *role.Credentials.AccessKeyId,\n\t\tSecretAccessKey: *role.Credentials.SecretAccessKey,\n\t\tSessionToken: *role.Credentials.SessionToken,\n\t\tExpires: *role.Credentials.Expiration,\n\t}, nil\n}\n\nfunc (s *defaultAwsAttestationMetadataProvider) awsRegion() string {\n\treturn s.awsCfg.Region\n}\n\nfunc (c *awsIdentityAttestationCreator) createAttestation() (*wifAttestation, error) {\n\tlogger.Debug(\"Creating AWS identity attestation...\")\n\n\tattestationService := c.attestationServiceFactory(c.ctx, c.cfg)\n\tif attestationService == nil {\n\t\treturn nil, errors.New(\"AWS attestation service could not be created\")\n\t}\n\n\tvar creds aws.Credentials\n\tvar err error\n\n\tif len(c.cfg.WorkloadIdentityImpersonationPath) == 0 {\n\t\tif creds, err = attestationService.awsCredentials(); err != nil {\n\t\t\tlogger.Debugf(\"error while getting for aws credentials. %v\", err)\n\t\t\treturn nil, err\n\t\t}\n\t} else {\n\t\tif creds, err = attestationService.awsCredentialsViaRoleChaining(); err != nil {\n\t\t\tlogger.Debugf(\"error while getting for aws credentials via role chaining. %v\", err)\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tif creds.AccessKeyID == \"\" || creds.SecretAccessKey == \"\" {\n\t\treturn nil, fmt.Errorf(\"no AWS credentials were found\")\n\t}\n\n\tregion := attestationService.awsRegion()\n\tif region == \"\" {\n\t\treturn nil, fmt.Errorf(\"no AWS region was found\")\n\t}\n\n\tstsHostname := stsHostname(region)\n\treq, err := c.createStsRequest(stsHostname)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\terr = c.signRequestWithSigV4(c.ctx, req, creds, region)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tcredential, err := c.createBase64EncodedRequestCredential(req)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &wifAttestation{\n\t\tProviderType: string(awsWif),\n\t\tCredential: credential,\n\t\tMetadata: map[string]string{},\n\t}, nil\n}\n\nfunc stsHostname(region string) string {\n\tvar domain string\n\tif strings.HasPrefix(region, \"cn-\") {\n\t\tdomain = \"amazonaws.com.cn\"\n\t} else {\n\t\tdomain = \"amazonaws.com\"\n\t}\n\treturn fmt.Sprintf(\"sts.%s.%s\", region, domain)\n}\n\nfunc (c *awsIdentityAttestationCreator) createStsRequest(hostname string) (*http.Request, error) {\n\turl := fmt.Sprintf(\"https://%s?Action=GetCallerIdentity&Version=2011-06-15\", hostname)\n\treq, err := http.NewRequest(\"POST\", url, nil)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treq.Header.Set(\"Host\", hostname)\n\treq.Header.Set(\"X-Snowflake-Audience\", \"snowflakecomputing.com\")\n\treturn req, nil\n}\n\nfunc (c *awsIdentityAttestationCreator) signRequestWithSigV4(ctx context.Context, req *http.Request, creds aws.Credentials, region string) error {\n\tsigner := v4.NewSigner()\n\t// as per docs of SignHTTP, the payload hash must be present even if the payload is empty\n\tpayloadHash := hex.EncodeToString(sha256.New().Sum(nil))\n\treturn signer.SignHTTP(ctx, creds, req, payloadHash, \"sts\", region, time.Now())\n}\n\nfunc (c *awsIdentityAttestationCreator) createBase64EncodedRequestCredential(req *http.Request) (string, error) {\n\theaders := make(map[string]string)\n\tfor key, values := range req.Header {\n\t\theaders[key] = values[0]\n\t}\n\n\tassertion := map[string]any{\n\t\t\"url\": req.URL.String(),\n\t\t\"method\": req.Method,\n\t\t\"headers\": headers,\n\t}\n\n\tassertionJSON, err := json.Marshal(assertion)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn base64.StdEncoding.EncodeToString(assertionJSON), nil\n}\n\nfunc (c *gcpIdentityAttestationCreator) createAttestation() (*wifAttestation, error) {\n\tlogger.Debugf(\"Creating GCP identity attestation...\")\n\tif len(c.cfg.WorkloadIdentityImpersonationPath) == 0 {\n\t\treturn c.createGcpIdentityTokenFromMetadataService()\n\t}\n\treturn c.createGcpIdentityViaImpersonation()\n}\n\nfunc (c *gcpIdentityAttestationCreator) createGcpIdentityTokenFromMetadataService() (*wifAttestation, error) {\n\treq, err := c.createTokenRequest()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to create GCP token request: %w\", err)\n\t}\n\ttoken := fetchTokenFromMetadataService(req, c.cfg, c.telemetry)\n\tif token == \"\" {\n\t\treturn nil, fmt.Errorf(\"no GCP token was found\")\n\t}\n\tsub, _, err := extractSubIssWithoutVerifyingSignature(token)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"could not extract claims from token: %v\", err)\n\t}\n\treturn &wifAttestation{\n\t\tProviderType: string(gcpWif),\n\t\tCredential: token,\n\t\tMetadata: map[string]string{\"sub\": sub},\n\t}, nil\n}\n\nfunc (c *gcpIdentityAttestationCreator) createTokenRequest() (*http.Request, error) {\n\turi := fmt.Sprintf(\"%s/computeMetadata/v1/instance/service-accounts/default/identity?audience=%s\",\n\t\tc.metadataServiceBaseURL, snowflakeAudience)\n\treq, err := http.NewRequest(\"GET\", uri, nil)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to create HTTP request: %v\", err)\n\t}\n\treq.Header.Set(gcpMetadataFlavorHeaderName, gcpMetadataFlavor)\n\treturn req, nil\n}\n\nfunc (c *gcpIdentityAttestationCreator) createGcpIdentityViaImpersonation() (*wifAttestation, error) {\n\t// initialize transport\n\ttransport, err := newTransportFactory(c.cfg, c.telemetry).createTransport(transportConfigFor(transportTypeWIF))\n\tif err != nil {\n\t\tlogger.Debugf(\"Failed to create HTTP transport: %v\", err)\n\t\treturn nil, err\n\t}\n\tclient := &http.Client{Transport: transport}\n\n\t// fetch access token for impersonation\n\taccessToken, err := c.fetchServiceToken(client)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// map paths to full service account paths\n\tvar fullServiceAccountPaths []string\n\tfor _, path := range c.cfg.WorkloadIdentityImpersonationPath {\n\t\tfullServiceAccountPaths = append(fullServiceAccountPaths, fmt.Sprintf(\"projects/-/serviceAccounts/%s\", path))\n\t}\n\ttargetServiceAccount := fullServiceAccountPaths[len(fullServiceAccountPaths)-1]\n\tdelegates := fullServiceAccountPaths[:len(fullServiceAccountPaths)-1]\n\n\t// fetch impersonated token\n\timpersonationToken, err := c.fetchImpersonatedToken(targetServiceAccount, delegates, accessToken, client)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// create attestation\n\tsub, _, err := extractSubIssWithoutVerifyingSignature(impersonationToken)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"could not extract claims from token: %v\", err)\n\t}\n\treturn &wifAttestation{\n\t\tProviderType: string(gcpWif),\n\t\tCredential: impersonationToken,\n\t\tMetadata: map[string]string{\"sub\": sub},\n\t}, nil\n}\n\nfunc (c *gcpIdentityAttestationCreator) fetchServiceToken(client *http.Client) (string, error) {\n\t// initialize and do request\n\treq, err := http.NewRequest(\"GET\", c.metadataServiceBaseURL+\"/computeMetadata/v1/instance/service-accounts/default/token\", nil)\n\tif err != nil {\n\t\tlogger.Debugf(\"cannot create token request for impersonation. %v\", err)\n\t\treturn \"\", err\n\t}\n\treq.Header.Set(gcpMetadataFlavorHeaderName, gcpMetadataFlavor)\n\tresp, err := client.Do(req)\n\tif err != nil {\n\t\tlogger.Debugf(\"cannot fetch token for impersonation. %v\", err)\n\t\treturn \"\", err\n\t}\n\tdefer func(body io.ReadCloser) {\n\t\tif err = body.Close(); err != nil {\n\t\t\tlogger.Debugf(\"cannot close token response body for impersonation. %v\", err)\n\t\t}\n\t}(resp.Body)\n\n\t// if it is not 200, do not parse the response\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn \"\", fmt.Errorf(\"token response status is %v, not parsing\", resp.StatusCode)\n\t}\n\n\t// parse response and extract access token\n\taccessTokenResponse := struct {\n\t\tAccessToken string `json:\"access_token\"`\n\t}{}\n\tif err = json.NewDecoder(resp.Body).Decode(&accessTokenResponse); err != nil {\n\t\tlogger.Debugf(\"cannot decode token for impersonation. %v\", err)\n\t\treturn \"\", err\n\t}\n\taccessToken := accessTokenResponse.AccessToken\n\treturn accessToken, nil\n}\n\nfunc (c *gcpIdentityAttestationCreator) fetchImpersonatedToken(targetServiceAccount string, delegates []string, accessToken string, client *http.Client) (string, error) {\n\t// prepare the request\n\turl := fmt.Sprintf(\"%v/v1/%v:generateIdToken\", c.iamCredentialsURL, targetServiceAccount)\n\tbody := struct {\n\t\tDelegates []string `json:\"delegates,omitempty\"`\n\t\tAudience string `json:\"audience\"`\n\t}{\n\t\tDelegates: delegates,\n\t\tAudience: snowflakeAudience,\n\t}\n\tpayload := new(bytes.Buffer)\n\tif err := json.NewEncoder(payload).Encode(body); err != nil {\n\t\tlogger.Debugf(\"cannot encode impersonation request body. %v\", err)\n\t\treturn \"\", err\n\t}\n\treq, err := http.NewRequest(\"POST\", url, payload)\n\tif err != nil {\n\t\tlogger.Debugf(\"cannot create token request for impersonation. %v\", err)\n\t\treturn \"\", err\n\t}\n\treq.Header.Set(\"Authorization\", \"Bearer \"+accessToken)\n\treq.Header.Set(\"Content-Type\", \"application/json\")\n\n\t// send the request\n\tresp, err := client.Do(req)\n\tif err != nil {\n\t\tlogger.Debugf(\"cannot call impersonation service. %v\", err)\n\t\treturn \"\", err\n\t}\n\tdefer func(body io.ReadCloser) {\n\t\tif err = body.Close(); err != nil {\n\t\t\tlogger.Debugf(\"cannot close token response body for impersonation. %v\", err)\n\t\t}\n\t}(resp.Body)\n\n\t// handle the response\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn \"\", fmt.Errorf(\"response status is %v, not parsing\", resp.StatusCode)\n\t}\n\ttokenResponse := struct {\n\t\tToken string `json:\"token\"`\n\t}{}\n\tif err = json.NewDecoder(resp.Body).Decode(&tokenResponse); err != nil {\n\t\tlogger.Debugf(\"cannot decode token response. %v\", err)\n\t\treturn \"\", err\n\t}\n\treturn tokenResponse.Token, nil\n}\n\nfunc fetchTokenFromMetadataService(req *http.Request, cfg *Config, telemetry *snowflakeTelemetry) string {\n\ttransport, err := newTransportFactory(cfg, telemetry).createTransport(transportConfigFor(transportTypeWIF))\n\tif err != nil {\n\t\tlogger.Debugf(\"Failed to create HTTP transport: %v\", err)\n\t\treturn \"\"\n\t}\n\tclient := &http.Client{Transport: transport}\n\tresp, err := client.Do(req)\n\tif err != nil {\n\t\tlogger.Debugf(\"Metadata server request was not successful: %v\", err)\n\t\treturn \"\"\n\t}\n\tdefer func() {\n\t\tif err = resp.Body.Close(); err != nil {\n\t\t\tlogger.Debugf(\"Failed to close response body: %v\", err)\n\t\t}\n\t}()\n\tbody, err := io.ReadAll(resp.Body)\n\tif err != nil {\n\t\tlogger.Debugf(\"Failed to read response body: %v\", err)\n\t\treturn \"\"\n\t}\n\treturn string(body)\n}\n\nfunc extractSubIssWithoutVerifyingSignature(token string) (subject string, issuer string, err error) {\n\tclaims, err := extractClaimsMap(token)\n\tif err != nil {\n\t\treturn \"\", \"\", err\n\t}\n\tissuerClaim, ok := claims[\"iss\"]\n\tif !ok {\n\t\treturn \"\", \"\", errors.New(\"missing issuer claim in JWT token\")\n\t}\n\tsubjectClaim, ok := claims[\"sub\"]\n\tif !ok {\n\t\treturn \"\", \"\", errors.New(\"missing sub claim in JWT token\")\n\t}\n\tsubject, ok = subjectClaim.(string)\n\tif !ok {\n\t\treturn \"\", \"\", errors.New(\"sub claim is not a string in JWT token\")\n\t}\n\tissuer, ok = issuerClaim.(string)\n\tif !ok {\n\t\treturn \"\", \"\", errors.New(\"iss claim is not a string in JWT token\")\n\t}\n\treturn\n}\n\n// extractClaimsMap parses a JWT token and returns its claims as a map.\n// It does not verify the token signature.\nfunc extractClaimsMap(token string) (map[string]any, error) {\n\tparser := jwt.NewParser(jwt.WithoutClaimsValidation())\n\tclaims := jwt.MapClaims{}\n\t_, _, err := parser.ParseUnverified(token, claims)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"unable to extract JWT claims from token: %w\", err)\n\t}\n\treturn claims, nil\n}\n\nfunc (c *oidcIdentityAttestationCreator) createAttestation() (*wifAttestation, error) {\n\tlogger.Debugf(\"Creating OIDC identity attestation...\")\n\ttoken, err := c.token()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to get OIDC token: %w\", err)\n\t}\n\tif token == \"\" {\n\t\treturn nil, fmt.Errorf(\"no OIDC token was specified\")\n\t}\n\tsub, iss, err := extractSubIssWithoutVerifyingSignature(token)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif sub == \"\" || iss == \"\" {\n\t\treturn nil, errors.New(\"missing sub or iss claim in JWT token\")\n\t}\n\treturn &wifAttestation{\n\t\tProviderType: string(oidcWif),\n\t\tCredential: token,\n\t\tMetadata: map[string]string{\"sub\": sub},\n\t}, nil\n}\n\n// azureAttestationMetadataProvider defines the interface for Azure attestation services\ntype azureAttestationMetadataProvider interface {\n\tidentityEndpoint() string\n\tidentityHeader() string\n\tclientID() string\n}\n\ntype defaultAzureAttestationMetadataProvider struct{}\n\nfunc (p *defaultAzureAttestationMetadataProvider) identityEndpoint() string {\n\treturn os.Getenv(\"IDENTITY_ENDPOINT\")\n}\n\nfunc (p *defaultAzureAttestationMetadataProvider) identityHeader() string {\n\treturn os.Getenv(\"IDENTITY_HEADER\")\n}\n\nfunc (p *defaultAzureAttestationMetadataProvider) clientID() string {\n\treturn os.Getenv(\"MANAGED_IDENTITY_CLIENT_ID\")\n}\n\ntype azureIdentityAttestationCreator struct {\n\tazureAttestationMetadataProvider azureAttestationMetadataProvider\n\tcfg *Config\n\ttelemetry *snowflakeTelemetry\n\tworkloadIdentityEntraResource string\n\tazureMetadataServiceBaseURL string\n}\n\n// createAttestation creates an attestation using Azure identity\nfunc (a *azureIdentityAttestationCreator) createAttestation() (*wifAttestation, error) {\n\tlogger.Debug(\"Creating Azure identity attestation...\")\n\n\tidentityEndpoint := a.azureAttestationMetadataProvider.identityEndpoint()\n\tvar request *http.Request\n\tvar err error\n\n\tif identityEndpoint == \"\" {\n\t\trequest, err = a.azureVMIdentityRequest()\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to create Azure VM identity request: %v\", err)\n\t\t}\n\t} else {\n\t\tidentityHeader := a.azureAttestationMetadataProvider.identityHeader()\n\t\tif identityHeader == \"\" {\n\t\t\treturn nil, fmt.Errorf(\"managed identity is not enabled on this Azure function\")\n\t\t}\n\t\trequest, err = a.azureFunctionsIdentityRequest(\n\t\t\tidentityEndpoint,\n\t\t\tidentityHeader,\n\t\t\ta.azureAttestationMetadataProvider.clientID(),\n\t\t)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"failed to create Azure Functions identity request: %v\", err)\n\t\t}\n\t}\n\n\ttokenJSON := fetchTokenFromMetadataService(request, a.cfg, a.telemetry)\n\tif tokenJSON == \"\" {\n\t\treturn nil, fmt.Errorf(\"could not fetch Azure token\")\n\t}\n\n\ttoken, err := extractTokenFromJSON(tokenJSON)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to extract token from JSON: %v\", err)\n\t}\n\tif token == \"\" {\n\t\treturn nil, fmt.Errorf(\"no access token found in Azure response\")\n\t}\n\n\tsub, iss, err := extractSubIssWithoutVerifyingSignature(token)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to extract sub and iss claims from token: %v\", err)\n\t}\n\tif sub == \"\" || iss == \"\" {\n\t\treturn nil, fmt.Errorf(\"missing sub or iss claim in JWT token\")\n\t}\n\n\treturn &wifAttestation{\n\t\tProviderType: string(azureWif),\n\t\tCredential: token,\n\t\tMetadata: map[string]string{\"sub\": sub, \"iss\": iss},\n\t}, nil\n}\n\nfunc determineEntraResource(config *Config) string {\n\tif config != nil && config.WorkloadIdentityEntraResource != \"\" {\n\t\treturn config.WorkloadIdentityEntraResource\n\t}\n\t// default resource if none specified\n\treturn \"api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad\"\n}\n\nfunc extractTokenFromJSON(tokenJSON string) (string, error) {\n\tvar response struct {\n\t\tAccessToken string `json:\"access_token\"`\n\t}\n\n\terr := json.Unmarshal([]byte(tokenJSON), &response)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn response.AccessToken, nil\n}\n\nfunc (a *azureIdentityAttestationCreator) azureFunctionsIdentityRequest(identityEndpoint, identityHeader, managedIdentityClientID string) (*http.Request, error) {\n\tqueryParams := fmt.Sprintf(\"api-version=2019-08-01&resource=%s\", a.workloadIdentityEntraResource)\n\tif managedIdentityClientID != \"\" {\n\t\tqueryParams += fmt.Sprintf(\"&client_id=%s\", managedIdentityClientID)\n\t}\n\n\turl := fmt.Sprintf(\"%s?%s\", identityEndpoint, queryParams)\n\treq, err := http.NewRequest(\"GET\", url, nil)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treq.Header.Add(\"X-IDENTITY-HEADER\", identityHeader)\n\n\treturn req, nil\n}\n\nfunc (a *azureIdentityAttestationCreator) azureVMIdentityRequest() (*http.Request, error) {\n\turlWithoutQuery := a.azureMetadataServiceBaseURL + \"/metadata/identity/oauth2/token?\"\n\tqueryParams := fmt.Sprintf(\"api-version=2018-02-01&resource=%s\", a.workloadIdentityEntraResource)\n\n\turl := urlWithoutQuery + queryParams\n\treq, err := http.NewRequest(\"GET\", url, nil)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treq.Header.Add(\"Metadata\", \"true\")\n\n\treturn req, nil\n}\n"
},
{
"path": "auth_wif_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"context\"\n\t\"database/sql\"\n\t\"encoding/base64\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"os\"\n\t\"os/exec\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/aws/aws-sdk-go-v2/aws\"\n)\n\ntype mockWifAttestationCreator struct {\n\tproviderType wifProviderType\n\treturnError error\n}\n\nfunc (m *mockWifAttestationCreator) createAttestation() (*wifAttestation, error) {\n\tif m.returnError != nil {\n\t\treturn nil, m.returnError\n\t}\n\treturn &wifAttestation{\n\t\tProviderType: string(m.providerType),\n\t}, nil\n}\n\nfunc TestGetAttestation(t *testing.T) {\n\tawsError := errors.New(\"aws attestation error\")\n\tgcpError := errors.New(\"gcp attestation error\")\n\tazureError := errors.New(\"azure attestation error\")\n\toidcError := errors.New(\"oidc attestation error\")\n\n\tprovider := &wifAttestationProvider{\n\t\tcontext: context.Background(),\n\t\tawsCreator: &mockWifAttestationCreator{providerType: awsWif},\n\t\tgcpCreator: &mockWifAttestationCreator{providerType: gcpWif},\n\t\tazureCreator: &mockWifAttestationCreator{providerType: azureWif},\n\t\toidcCreator: &mockWifAttestationCreator{providerType: oidcWif},\n\t}\n\n\tproviderWithErrors := &wifAttestationProvider{\n\t\tcontext: context.Background(),\n\t\tawsCreator: &mockWifAttestationCreator{providerType: awsWif, returnError: awsError},\n\t\tgcpCreator: &mockWifAttestationCreator{providerType: gcpWif, returnError: gcpError},\n\t\tazureCreator: &mockWifAttestationCreator{providerType: azureWif, returnError: azureError},\n\t\toidcCreator: &mockWifAttestationCreator{providerType: oidcWif, returnError: oidcError},\n\t}\n\n\ttests := []struct {\n\t\tname string\n\t\tprovider *wifAttestationProvider\n\t\tidentityProvider string\n\t\texpectedResult *wifAttestation\n\t\texpectedError error\n\t}{\n\t\t{\n\t\t\tname: \"AWS success\",\n\t\t\tprovider: provider,\n\t\t\tidentityProvider: \"AWS\",\n\t\t\texpectedResult: &wifAttestation{ProviderType: string(awsWif)},\n\t\t\texpectedError: nil,\n\t\t},\n\t\t{\n\t\t\tname: \"AWS error\",\n\t\t\tprovider: providerWithErrors,\n\t\t\tidentityProvider: \"AWS\",\n\t\t\texpectedResult: nil,\n\t\t\texpectedError: awsError,\n\t\t},\n\t\t{\n\t\t\tname: \"GCP success\",\n\t\t\tprovider: provider,\n\t\t\tidentityProvider: \"GCP\",\n\t\t\texpectedResult: &wifAttestation{ProviderType: string(gcpWif)},\n\t\t\texpectedError: nil,\n\t\t},\n\t\t{\n\t\t\tname: \"GCP error\",\n\t\t\tprovider: providerWithErrors,\n\t\t\tidentityProvider: \"GCP\",\n\t\t\texpectedResult: nil,\n\t\t\texpectedError: gcpError,\n\t\t},\n\t\t{\n\t\t\tname: \"AZURE success\",\n\t\t\tprovider: provider,\n\t\t\tidentityProvider: \"AZURE\",\n\t\t\texpectedResult: &wifAttestation{ProviderType: string(azureWif)},\n\t\t\texpectedError: nil,\n\t\t},\n\t\t{\n\t\t\tname: \"AZURE error\",\n\t\t\tprovider: providerWithErrors,\n\t\t\tidentityProvider: \"AZURE\",\n\t\t\texpectedResult: nil,\n\t\t\texpectedError: azureError,\n\t\t},\n\t\t{\n\t\t\tname: \"OIDC success\",\n\t\t\tprovider: provider,\n\t\t\tidentityProvider: \"OIDC\",\n\t\t\texpectedResult: &wifAttestation{ProviderType: string(oidcWif)},\n\t\t\texpectedError: nil,\n\t\t},\n\t\t{\n\t\t\tname: \"OIDC error\",\n\t\t\tprovider: providerWithErrors,\n\t\t\tidentityProvider: \"OIDC\",\n\t\t\texpectedResult: nil,\n\t\t\texpectedError: oidcError,\n\t\t},\n\t\t{\n\t\t\tname: \"Unknown provider\",\n\t\t\tprovider: provider,\n\t\t\tidentityProvider: \"UNKNOWN\",\n\t\t\texpectedResult: nil,\n\t\t\texpectedError: errors.New(\"unknown WorkloadIdentityProvider specified: UNKNOWN. Valid values are: AWS, GCP, AZURE, OIDC\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Empty provider\",\n\t\t\tprovider: provider,\n\t\t\tidentityProvider: \"\",\n\t\t\texpectedResult: nil,\n\t\t\texpectedError: errors.New(\"unknown WorkloadIdentityProvider specified: . Valid values are: AWS, GCP, AZURE, OIDC\"),\n\t\t},\n\t}\n\n\tfor _, test := range tests {\n\t\tt.Run(test.name, func(t *testing.T) {\n\t\t\tattestation, err := test.provider.getAttestation(test.identityProvider)\n\t\t\tif test.expectedError != nil {\n\t\t\t\tassertNilE(t, attestation)\n\t\t\t\tassertNotNilF(t, err)\n\t\t\t\tassertEqualE(t, test.expectedError.Error(), err.Error())\n\t\t\t} else if test.expectedResult != nil {\n\t\t\t\tassertNilE(t, err)\n\t\t\t\tassertNotNilF(t, attestation)\n\t\t\t\tassertEqualE(t, test.expectedResult.ProviderType, attestation.ProviderType)\n\t\t\t} else {\n\t\t\t\tt.Fatal(\"test case must specify either expectedError or expectedResult\")\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestAwsIdentityAttestationCreator(t *testing.T) {\n\ttests := []struct {\n\t\tname string\n\t\tconfig Config\n\t\tattestationSvc awsAttestationMetadataProvider\n\t\texpectedError error\n\t\texpectedProvider string\n\t\texpectedStsHost string\n\t}{\n\t\t{\n\t\t\tname: \"No attestation service\",\n\t\t\tattestationSvc: nil,\n\t\t\texpectedError: fmt.Errorf(\"AWS attestation service could not be created\"),\n\t\t},\n\t\t{\n\t\t\tname: \"No AWS credentials\",\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: aws.Credentials{},\n\t\t\t\tregion: \"us-west-2\",\n\t\t\t},\n\t\t\texpectedError: fmt.Errorf(\"no AWS credentials were found\"),\n\t\t},\n\t\t{\n\t\t\tname: \"No AWS region\",\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tregion: \"\",\n\t\t\t},\n\t\t\texpectedError: fmt.Errorf(\"no AWS region was found\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Successful attestation\",\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tregion: \"us-west-2\",\n\t\t\t},\n\t\t\texpectedProvider: \"AWS\",\n\t\t\texpectedStsHost: \"sts.us-west-2.amazonaws.com\",\n\t\t},\n\t\t{\n\t\t\tname: \"Successful attestation for CN region\",\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tregion: \"cn-northwest-1\",\n\t\t\t},\n\t\t\texpectedProvider: \"AWS\",\n\t\t\texpectedStsHost: \"sts.cn-northwest-1.amazonaws.com.cn\",\n\t\t},\n\t\t{\n\t\t\tname: \"Successful attestation with single role chaining\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\"arn:aws:iam::123456789012:role/test-role\"},\n\t\t\t},\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tchainingCreds: aws.Credentials{\n\t\t\t\t\tAccessKeyID: \"chainedAccessKey\",\n\t\t\t\t\tSecretAccessKey: \"chainedSecretKey\",\n\t\t\t\t\tSessionToken: \"chainedSessionToken\",\n\t\t\t\t},\n\t\t\t\tregion: \"us-east-1\",\n\t\t\t\tuseRoleChaining: true,\n\t\t\t},\n\t\t\texpectedProvider: \"AWS\",\n\t\t\texpectedStsHost: \"sts.us-east-1.amazonaws.com\",\n\t\t},\n\t\t{\n\t\t\tname: \"Successful attestation with multiple role chaining\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\n\t\t\t\t\t\"arn:aws:iam::123456789012:role/role1\",\n\t\t\t\t\t\"arn:aws:iam::123456789012:role/role2\",\n\t\t\t\t\t\"arn:aws:iam::123456789012:role/role3\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tchainingCreds: aws.Credentials{\n\t\t\t\t\tAccessKeyID: \"finalRoleAccessKey\",\n\t\t\t\t\tSecretAccessKey: \"finalRoleSecretKey\",\n\t\t\t\t\tSessionToken: \"finalRoleSessionToken\",\n\t\t\t\t},\n\t\t\t\tregion: \"us-west-2\",\n\t\t\t\tuseRoleChaining: true,\n\t\t\t},\n\t\t\texpectedProvider: \"AWS\",\n\t\t\texpectedStsHost: \"sts.us-west-2.amazonaws.com\",\n\t\t},\n\t\t{\n\t\t\tname: \"Role chaining with no credentials\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\"arn:aws:iam::123456789012:role/test-role\"},\n\t\t\t},\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: aws.Credentials{},\n\t\t\t\tregion: \"us-west-2\",\n\t\t\t\tuseRoleChaining: true,\n\t\t\t},\n\t\t\texpectedError: fmt.Errorf(\"no AWS credentials were found\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Role chaining with no region\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\"arn:aws:iam::123456789012:role/test-role\"},\n\t\t\t},\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: aws.Credentials{\n\t\t\t\t\tAccessKeyID: \"chainedAccessKey\",\n\t\t\t\t\tSecretAccessKey: \"chainedSecretKey\",\n\t\t\t\t\tSessionToken: \"chainedSessionToken\",\n\t\t\t\t},\n\t\t\t\tregion: \"\",\n\t\t\t\tuseRoleChaining: true,\n\t\t\t},\n\t\t\texpectedError: fmt.Errorf(\"no AWS region was found\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Role chaining failure\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\"arn:aws:iam::123456789012:role/test-role\"},\n\t\t\t},\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tregion: \"us-west-2\",\n\t\t\t\tchainingError: fmt.Errorf(\"failed to assume role: AccessDenied\"),\n\t\t\t\tuseRoleChaining: true,\n\t\t\t},\n\t\t\texpectedError: fmt.Errorf(\"failed to assume role: AccessDenied\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Cross-account role chaining\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\n\t\t\t\t\t\"arn:aws:iam::111111111111:role/cross-account-role\",\n\t\t\t\t\t\"arn:aws:iam::222222222222:role/target-role\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tchainingCreds: aws.Credentials{\n\t\t\t\t\tAccessKeyID: \"crossAccountAccessKey\",\n\t\t\t\t\tSecretAccessKey: \"crossAccountSecretKey\",\n\t\t\t\t\tSessionToken: \"crossAccountSessionToken\",\n\t\t\t\t},\n\t\t\t\tregion: \"us-east-1\",\n\t\t\t\tuseRoleChaining: true,\n\t\t\t},\n\t\t\texpectedProvider: \"AWS\",\n\t\t\texpectedStsHost: \"sts.us-east-1.amazonaws.com\",\n\t\t},\n\t\t{\n\t\t\tname: \"Role chaining in CN region\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\"arn:aws-cn:iam::123456789012:role/cn-role\"},\n\t\t\t},\n\t\t\tattestationSvc: &mockAwsAttestationMetadataProvider{\n\t\t\t\tcreds: mockCreds,\n\t\t\t\tchainingCreds: aws.Credentials{\n\t\t\t\t\tAccessKeyID: \"cnRoleAccessKey\",\n\t\t\t\t\tSecretAccessKey: \"cnRoleSecretKey\",\n\t\t\t\t\tSessionToken: \"cnRoleSessionToken\",\n\t\t\t\t},\n\t\t\t\tregion: \"cn-north-1\",\n\t\t\t\tuseRoleChaining: true,\n\t\t\t},\n\t\t\texpectedProvider: \"AWS\",\n\t\t\texpectedStsHost: \"sts.cn-north-1.amazonaws.com.cn\",\n\t\t},\n\t}\n\n\tfor _, test := range tests {\n\t\tt.Run(test.name, func(t *testing.T) {\n\t\t\tcreator := &awsIdentityAttestationCreator{\n\t\t\t\tattestationServiceFactory: func(ctx context.Context, cfg *Config) awsAttestationMetadataProvider {\n\t\t\t\t\treturn test.attestationSvc\n\t\t\t\t},\n\t\t\t\tctx: context.Background(),\n\t\t\t\tcfg: &test.config,\n\t\t\t}\n\t\t\tattestation, err := creator.createAttestation()\n\t\t\tif test.expectedError != nil {\n\t\t\t\tassertNilF(t, attestation)\n\t\t\t\tassertNotNilE(t, err)\n\t\t\t\tassertEqualE(t, test.expectedError.Error(), err.Error())\n\t\t\t} else {\n\t\t\t\tassertNilE(t, err)\n\t\t\t\tassertNotNilE(t, attestation)\n\t\t\t\tassertNotNilE(t, attestation.Credential)\n\t\t\t\tassertEqualE(t, test.expectedProvider, attestation.ProviderType)\n\t\t\t\tdecoded, err := base64.StdEncoding.DecodeString(attestation.Credential)\n\t\t\t\tif err != nil {\n\t\t\t\t\tt.Fatalf(\"Failed to decode credential: %v\", err)\n\t\t\t\t}\n\t\t\t\tvar credentialMap map[string]any\n\t\t\t\tif err := json.Unmarshal(decoded, &credentialMap); err != nil {\n\t\t\t\t\tt.Fatalf(\"Failed to unmarshal credential JSON: %v\", err)\n\t\t\t\t}\n\t\t\t\tassertEqualE(t, fmt.Sprintf(\"https://%s?Action=GetCallerIdentity&Version=2011-06-15\", test.expectedStsHost), credentialMap[\"url\"])\n\t\t\t}\n\t\t})\n\t}\n}\n\ntype mockAwsAttestationMetadataProvider struct {\n\tcreds aws.Credentials\n\tregion string\n\tchainingCreds aws.Credentials\n\tchainingError error\n\tuseRoleChaining bool\n}\n\nvar mockCreds = aws.Credentials{\n\tAccessKeyID: \"mockAccessKey\",\n\tSecretAccessKey: \"mockSecretKey\",\n\tSessionToken: \"mockSessionToken\",\n}\n\nfunc (m *mockAwsAttestationMetadataProvider) awsCredentials() (aws.Credentials, error) {\n\treturn m.creds, nil\n}\n\nfunc (m *mockAwsAttestationMetadataProvider) awsCredentialsViaRoleChaining() (aws.Credentials, error) {\n\tif m.chainingError != nil {\n\t\treturn aws.Credentials{}, m.chainingError\n\t}\n\tif m.chainingCreds.AccessKeyID != \"\" {\n\t\treturn m.chainingCreds, nil\n\t}\n\treturn m.creds, nil\n}\n\nfunc (m *mockAwsAttestationMetadataProvider) awsRegion() string {\n\treturn m.region\n}\n\nfunc TestGcpIdentityAttestationCreator(t *testing.T) {\n\ttests := []struct {\n\t\tname string\n\t\twiremockMappingPath string\n\t\tconfig Config\n\t\texpectedError error\n\t\texpectedSub string\n\t}{\n\t\t{\n\t\t\tname: \"Successful flow\",\n\t\t\twiremockMappingPath: \"auth/wif/gcp/successful_flow.json\",\n\t\t\texpectedError: nil,\n\t\t\texpectedSub: \"some-subject\",\n\t\t},\n\t\t{\n\t\t\tname: \"Successful impersonation flow\",\n\t\t\twiremockMappingPath: \"auth/wif/gcp/successful_impersionation_flow.json\",\n\t\t\tconfig: Config{\n\t\t\t\tWorkloadIdentityImpersonationPath: []string{\n\t\t\t\t\t\"delegate1\",\n\t\t\t\t\t\"delegate2\",\n\t\t\t\t\t\"targetServiceAccount\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectedError: nil,\n\t\t\texpectedSub: \"some-impersonated-subject\",\n\t\t},\n\t\t{\n\t\t\tname: \"No GCP credential - http error\",\n\t\t\twiremockMappingPath: \"auth/wif/gcp/http_error.json\",\n\t\t\texpectedError: fmt.Errorf(\"no GCP token was found\"),\n\t\t\texpectedSub: \"\",\n\t\t},\n\t\t{\n\t\t\tname: \"missing issuer claim\",\n\t\t\twiremockMappingPath: \"auth/wif/gcp/missing_issuer_claim.json\",\n\t\t\texpectedError: fmt.Errorf(\"could not extract claims from token: missing issuer claim in JWT token\"),\n\t\t\texpectedSub: \"\",\n\t\t},\n\t\t{\n\t\t\tname: \"missing sub claim\",\n\t\t\twiremockMappingPath: \"auth/wif/gcp/missing_sub_claim.json\",\n\t\t\texpectedError: fmt.Errorf(\"could not extract claims from token: missing sub claim in JWT token\"),\n\t\t\texpectedSub: \"\",\n\t\t},\n\t\t{\n\t\t\tname: \"unparsable token\",\n\t\t\twiremockMappingPath: \"auth/wif/gcp/unparsable_token.json\",\n\t\t\texpectedError: fmt.Errorf(\"could not extract claims from token: unable to extract JWT claims from token: token is malformed: token contains an invalid number of segments\"),\n\t\t\texpectedSub: \"\",\n\t\t},\n\t}\n\n\tfor _, test := range tests {\n\t\tt.Run(test.name, func(t *testing.T) {\n\t\t\tcreator := &gcpIdentityAttestationCreator{\n\t\t\t\tcfg: &test.config,\n\t\t\t\tmetadataServiceBaseURL: wiremock.baseURL(),\n\t\t\t\tiamCredentialsURL: wiremock.baseURL(),\n\t\t\t}\n\t\t\twiremock.registerMappings(t, wiremockMapping{filePath: test.wiremockMappingPath})\n\t\t\tattestation, err := creator.createAttestation()\n\n\t\t\tif test.expectedError != nil {\n\t\t\t\tassertNilF(t, attestation)\n\t\t\t\tassertNotNilF(t, err)\n\t\t\t\tassertEqualE(t, test.expectedError.Error(), err.Error())\n\t\t\t} else {\n\t\t\t\tassertNilF(t, err)\n\t\t\t\tassertNotNilF(t, attestation)\n\t\t\t\tassertEqualE(t, string(gcpWif), attestation.ProviderType)\n\t\t\t\tassertEqualE(t, test.expectedSub, attestation.Metadata[\"sub\"])\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestOidcIdentityAttestationCreator(t *testing.T) {\n\tconst (\n\t\t/*\n\t\t * {\n\t\t * \"sub\": \"some-subject\",\n\t\t * \"iat\": 1743761213,\n\t\t * \"exp\": 1743764813,\n\t\t * \"aud\": \"www.example.com\"\n\t\t * }\n\t\t */\n\t\tmissingIssuerClaimToken = \"eyJ0eXAiOiJhdCtqd3QiLCJhbGciOiJFUzI1NiIsImtpZCI6ImU2M2I5NzA1OTRiY2NmZTAxMDlkOTg4OWM2MDk3OWEwIn0.eyJzdWIiOiJzb21lLXN1YmplY3QiLCJpYXQiOjE3NDM3NjEyMTMsImV4cCI6MTc0Mzc2NDgxMywiYXVkIjoid3d3LmV4YW1wbGUuY29tIn0.H6sN6kjA82EuijFcv-yCJTqau5qvVTCsk0ZQ4gvFQMkB7c71XPs4lkwTa7ZlNNlx9e6TpN1CVGnpCIRDDAZaDw\" // pragma: allowlist secret\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://accounts.google.com\",\n\t\t * \"iat\": 1743761213,\n\t\t * \"exp\": 1743764813,\n\t\t * \"aud\": \"www.example.com\"\n\t\t * }\n\t\t */\n\t\tmissingSubClaimToken = \"eyJ0eXAiOiJhdCtqd3QiLCJhbGciOiJFUzI1NiIsImtpZCI6ImU2M2I5NzA1OTRiY2NmZTAxMDlkOTg4OWM2MDk3OWEwIn0.eyJpc3MiOiJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20iLCJpYXQiOjE3NDM3NjEyMTMsImV4cCI6MTc0Mzc2NDgxMywiYXVkIjoid3d3LmV4YW1wbGUuY29tIn0.w0njdpfWFETVK8Ktq9GdvuKRQJjvhOplcSyvQ_zHHwBUSMapqO1bjEWBx5VhGkdECZIGS1VY7db_IOqT45yOMA\" // pragma: allowlist secret\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://oidc.eks.us-east-2.amazonaws.com/id/3B869BC5D12CEB5515358621D8085D58\",\n\t\t * \"iat\": 1743692017,\n\t\t * \"exp\": 1775228014,\n\t\t * \"aud\": \"www.example.com\",\n\t\t * \"sub\": \"system:serviceaccount:poc-namespace:oidc-sa\"\n\t\t * }\n\t\t */\n\t\tvalidToken = \"eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJodHRwczovL29pZGMuZWtzLnVzLWVhc3QtMi5hbWF6b25hd3MuY29tL2lkLzNCODY5QkM1RDEyQ0VCNTUxNTM1ODYyMUQ4MDg1RDU4IiwiaWF0IjoxNzQ0Mjg3ODc4LCJleHAiOjE3NzU4MjM4NzgsImF1ZCI6Ind3dy5leGFtcGxlLmNvbSIsInN1YiI6InN5c3RlbTpzZXJ2aWNlYWNjb3VudDpwb2MtbmFtZXNwYWNlOm9pZGMtc2EifQ.a8H6KRIF1XmM8lkqL6kR8ccInr7wAzQrbKd3ZHFgiEg\" // pragma: allowlist secret\n\t\tunparsableToken = \"unparsable_token\"\n\t\temptyToken = \"\"\n\t)\n\n\ttype testCase struct {\n\t\tname string\n\t\ttoken string\n\t\texpectedError error\n\t\texpectedSub string\n\t}\n\n\ttests := []testCase{\n\t\t{\n\t\t\tname: \"no token input\",\n\t\t\ttoken: emptyToken,\n\t\t\texpectedError: fmt.Errorf(\"no OIDC token was specified\"),\n\t\t},\n\t\t{\n\t\t\tname: \"valid token returns proper attestation\",\n\t\t\ttoken: validToken,\n\t\t\texpectedError: nil,\n\t\t\texpectedSub: \"system:serviceaccount:poc-namespace:oidc-sa\",\n\t\t},\n\t\t{\n\t\t\tname: \"missing issuer returns error\",\n\t\t\ttoken: missingIssuerClaimToken,\n\t\t\texpectedError: errors.New(\"missing issuer claim in JWT token\"),\n\t\t},\n\t\t{\n\t\t\tname: \"missing sub returns error\",\n\t\t\ttoken: missingSubClaimToken,\n\t\t\texpectedError: errors.New(\"missing sub claim in JWT token\"),\n\t\t},\n\t\t{\n\t\t\tname: \"unparsable token returns error\",\n\t\t\ttoken: unparsableToken,\n\t\t\texpectedError: errors.New(\"unable to extract JWT claims from token: token is malformed: token contains an invalid number of segments\"),\n\t\t},\n\t}\n\n\tfor _, test := range tests {\n\t\tt.Run(test.name, func(t *testing.T) {\n\t\t\tcreator := &oidcIdentityAttestationCreator{token: func() (string, error) {\n\t\t\t\treturn test.token, nil\n\t\t\t}}\n\t\t\tattestation, err := creator.createAttestation()\n\n\t\t\tif test.expectedError != nil {\n\t\t\t\tassertNotNilE(t, err)\n\t\t\t\tassertNilF(t, attestation)\n\t\t\t\tassertEqualE(t, test.expectedError.Error(), err.Error())\n\t\t\t} else {\n\t\t\t\tassertNilE(t, err)\n\t\t\t\tassertNotNilE(t, attestation)\n\t\t\t\tassertEqualE(t, string(oidcWif), attestation.ProviderType)\n\t\t\t\tassertEqualE(t, test.expectedSub, attestation.Metadata[\"sub\"])\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestAzureIdentityAttestationCreator(t *testing.T) {\n\ttests := []struct {\n\t\tname string\n\t\twiremockMappingPath string\n\t\tmetadataProvider *mockAzureAttestationMetadataProvider\n\t\tcfg *Config\n\t\texpectedIss string\n\t\texpectedError error\n\t}{\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t * \"sub\": \"77213E30-E8CB-4595-B1B6-5F050E8308FD\"\n\t\t * }\n\t\t */\n\t\t{\n\t\t\tname: \"Successful flow\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/successful_flow_basic.json\",\n\t\t\tmetadataProvider: azureVMMetadataProvider(),\n\t\t\texpectedIss: \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t\texpectedError: nil,\n\t\t},\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t * \"sub\": \"77213E30-E8CB-4595-B1B6-5F050E8308FD\"\n\t\t * }\n\t\t */\n\t\t{\n\t\t\tname: \"Successful flow v2 issuer\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/successful_flow_v2_issuer.json\",\n\t\t\tmetadataProvider: azureVMMetadataProvider(),\n\t\t\texpectedIss: \"https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t\texpectedError: nil,\n\t\t},\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t * \"sub\": \"77213E30-E8CB-4595-B1B6-5F050E8308FD\"\n\t\t * }\n\t\t */\n\t\t{\n\t\t\tname: \"Successful flow azure functions\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/successful_flow_azure_functions.json\",\n\t\t\tmetadataProvider: azureFunctionsMetadataProvider(),\n\t\t\texpectedIss: \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t\texpectedError: nil,\n\t\t},\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t * \"sub\": \"77213E30-E8CB-4595-B1B6-5F050E8308FD\"\n\t\t * }\n\t\t */\n\t\t{\n\t\t\tname: \"Successful flow azure functions v2 issuer\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/successful_flow_azure_functions_v2_issuer.json\",\n\t\t\tmetadataProvider: azureFunctionsMetadataProvider(),\n\t\t\texpectedIss: \"https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t\texpectedError: nil,\n\t\t},\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t * \"sub\": \"77213E30-E8CB-4595-B1B6-5F050E8308FD\"\n\t\t * }\n\t\t */\n\t\t{\n\t\t\tname: \"Successful flow azure functions no client ID\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/successful_flow_azure_functions_no_client_id.json\",\n\t\t\tmetadataProvider: &mockAzureAttestationMetadataProvider{\n\t\t\t\tidentityEndpointValue: wiremock.baseURL() + \"/metadata/identity/endpoint/from/env\",\n\t\t\t\tidentityHeaderValue: \"some-identity-header-from-env\",\n\t\t\t\tclientIDValue: \"\",\n\t\t\t},\n\t\t\texpectedIss: \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t\texpectedError: nil,\n\t\t},\n\t\t/*\n\t\t * {\n\t\t * \"iss\": \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t * \"sub\": \"77213E30-E8CB-4595-B1B6-5F050E8308FD\"\n\t\t * }\n\t\t */\n\t\t{\n\t\t\tname: \"Successful flow azure functions custom Entra resource\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/successful_flow_azure_functions_custom_entra_resource.json\",\n\t\t\tmetadataProvider: azureFunctionsMetadataProvider(),\n\t\t\tcfg: &Config{WorkloadIdentityEntraResource: \"api://1111111-2222-3333-44444-55555555\"},\n\t\t\texpectedIss: \"https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/\",\n\t\t\texpectedError: nil,\n\t\t},\n\t\t{\n\t\t\tname: \"Non-json response\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/non_json_response.json\",\n\t\t\tmetadataProvider: azureVMMetadataProvider(),\n\t\t\texpectedError: fmt.Errorf(\"failed to extract token from JSON: invalid character 'o' in literal null (expecting 'u')\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Identity endpoint but no identity header\",\n\t\t\tmetadataProvider: &mockAzureAttestationMetadataProvider{\n\t\t\t\tidentityEndpointValue: wiremock.baseURL() + \"/metadata/identity/endpoint/from/env\",\n\t\t\t\tidentityHeaderValue: \"\",\n\t\t\t\tclientIDValue: \"managed-client-id-from-env\",\n\t\t\t},\n\t\t\texpectedError: fmt.Errorf(\"managed identity is not enabled on this Azure function\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Unparsable token\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/unparsable_token.json\",\n\t\t\tmetadataProvider: azureVMMetadataProvider(),\n\t\t\texpectedError: fmt.Errorf(\"failed to extract sub and iss claims from token: unable to extract JWT claims from token: token is malformed: token contains an invalid number of segments\"),\n\t\t},\n\t\t{\n\t\t\tname: \"HTTP error\",\n\t\t\tmetadataProvider: azureVMMetadataProvider(),\n\t\t\twiremockMappingPath: \"auth/wif/azure/http_error.json\",\n\t\t\texpectedError: fmt.Errorf(\"could not fetch Azure token\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Missing sub or iss claim\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/missing_issuer_claim.json\",\n\t\t\tmetadataProvider: azureVMMetadataProvider(),\n\t\t\texpectedError: fmt.Errorf(\"failed to extract sub and iss claims from token: missing issuer claim in JWT token\"),\n\t\t},\n\t\t{\n\t\t\tname: \"Missing sub claim\",\n\t\t\twiremockMappingPath: \"auth/wif/azure/missing_sub_claim.json\",\n\t\t\tmetadataProvider: azureVMMetadataProvider(),\n\t\t\texpectedError: fmt.Errorf(\"failed to extract sub and iss claims from token: missing sub claim in JWT token\"),\n\t\t},\n\t}\n\n\tfor _, test := range tests {\n\t\tt.Run(test.name, func(t *testing.T) {\n\t\t\tif test.wiremockMappingPath != \"\" {\n\t\t\t\twiremock.registerMappings(t, wiremockMapping{filePath: test.wiremockMappingPath})\n\t\t\t}\n\t\t\tcreator := &azureIdentityAttestationCreator{\n\t\t\t\tcfg: test.cfg,\n\t\t\t\tazureMetadataServiceBaseURL: wiremock.baseURL(),\n\t\t\t\tazureAttestationMetadataProvider: test.metadataProvider,\n\t\t\t\tworkloadIdentityEntraResource: determineEntraResource(test.cfg),\n\t\t\t}\n\t\t\tattestation, err := creator.createAttestation()\n\n\t\t\tif test.expectedError != nil {\n\t\t\t\tassertNilF(t, attestation)\n\t\t\t\tassertNotNilE(t, err)\n\t\t\t\tassertEqualE(t, test.expectedError.Error(), err.Error())\n\t\t\t} else {\n\t\t\t\tassertNilF(t, err)\n\t\t\t\tassertNotNilF(t, attestation)\n\t\t\t\tassertEqualE(t, string(azureWif), attestation.ProviderType)\n\t\t\t\tassertEqualE(t, test.expectedIss, attestation.Metadata[\"iss\"])\n\t\t\t\tassertEqualE(t, \"77213E30-E8CB-4595-B1B6-5F050E8308FD\", attestation.Metadata[\"sub\"])\n\t\t\t}\n\t\t})\n\t}\n}\n\ntype mockAzureAttestationMetadataProvider struct {\n\tidentityEndpointValue string\n\tidentityHeaderValue string\n\tclientIDValue string\n}\n\nfunc (m *mockAzureAttestationMetadataProvider) identityEndpoint() string {\n\treturn m.identityEndpointValue\n}\n\nfunc (m *mockAzureAttestationMetadataProvider) identityHeader() string {\n\treturn m.identityHeaderValue\n}\n\nfunc (m *mockAzureAttestationMetadataProvider) clientID() string {\n\treturn m.clientIDValue\n}\n\nfunc azureFunctionsMetadataProvider() *mockAzureAttestationMetadataProvider {\n\treturn &mockAzureAttestationMetadataProvider{\n\t\tidentityEndpointValue: wiremock.baseURL() + \"/metadata/identity/endpoint/from/env\",\n\t\tidentityHeaderValue: \"some-identity-header-from-env\",\n\t\tclientIDValue: \"managed-client-id-from-env\",\n\t}\n}\n\nfunc azureVMMetadataProvider() *mockAzureAttestationMetadataProvider {\n\treturn &mockAzureAttestationMetadataProvider{\n\t\tidentityEndpointValue: \"\",\n\t\tidentityHeaderValue: \"\",\n\t\tclientIDValue: \"\",\n\t}\n}\n\n// Running this test locally:\n// * Push branch to repository\n// * Set PARAMETERS_SECRET\n// * Run ci/test_wif.sh\nfunc TestWorkloadIdentityAuthOnCloudVM(t *testing.T) {\n\taccount := os.Getenv(\"SNOWFLAKE_TEST_WIF_ACCOUNT\")\n\thost := os.Getenv(\"SNOWFLAKE_TEST_WIF_HOST\")\n\tprovider := os.Getenv(\"SNOWFLAKE_TEST_WIF_PROVIDER\")\n\tprintln(\"provider = \" + provider)\n\tif account == \"\" || host == \"\" || provider == \"\" {\n\t\tt.Skip(\"Test can run only on cloud VM with env variables set\")\n\t}\n\ttestCases := []struct {\n\t\tname string\n\t\tskip func() (bool, string)\n\t\tsetupCfg func(*testing.T, *Config)\n\t\texpectedUsername string\n\t}{\n\t\t{\n\t\t\tname: \"provider=\" + provider,\n\t\t\tsetupCfg: func(_ *testing.T, config *Config) {\n\t\t\t\tif provider != \"GCP+OIDC\" {\n\t\t\t\t\tconfig.WorkloadIdentityProvider = provider\n\t\t\t\t} else {\n\t\t\t\t\tconfig.WorkloadIdentityProvider = \"OIDC\"\n\t\t\t\t\tconfig.Token = func() string {\n\t\t\t\t\t\tcmd := exec.Command(\"wget\", \"-O\", \"-\", \"--header=Metadata-Flavor: Google\", \"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience=snowflakecomputing.com\")\n\t\t\t\t\t\toutput, err := cmd.Output()\n\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\tt.Fatalf(\"error executing GCP metadata request: %v\", err)\n\t\t\t\t\t\t}\n\t\t\t\t\t\ttoken := strings.TrimSpace(string(output))\n\t\t\t\t\t\tif token == \"\" {\n\t\t\t\t\t\t\tt.Fatal(\"failed to retrieve GCP access token: empty response\")\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn token\n\t\t\t\t\t}()\n\t\t\t\t}\n\t\t\t},\n\t\t\texpectedUsername: os.Getenv(\"SNOWFLAKE_TEST_WIF_USERNAME\"),\n\t\t},\n\t\t{\n\t\t\tname: \"provider=\" + provider + \",impersonation\",\n\t\t\tskip: func() (bool, string) {\n\t\t\t\tif provider != \"AWS\" && provider != \"GCP\" {\n\t\t\t\t\treturn true, \"Impersonation is supported only on AWS and GCP\"\n\t\t\t\t}\n\t\t\t\treturn false, \"\"\n\t\t\t},\n\t\t\tsetupCfg: func(t *testing.T, config *Config) {\n\t\t\t\tconfig.WorkloadIdentityProvider = provider\n\t\t\t\timpersonationPath := os.Getenv(\"SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH\")\n\t\t\t\tassertNotEqualF(t, impersonationPath, \"\", \"SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH is not set\")\n\t\t\t\tconfig.WorkloadIdentityImpersonationPath = strings.Split(impersonationPath, \",\")\n\t\t\t\tassertNotEqualF(t, os.Getenv(\"SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION\"), \"\", \"SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION is not set\")\n\t\t\t},\n\t\t\texpectedUsername: os.Getenv(\"SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION\"),\n\t\t},\n\t}\n\tfor _, tc := range testCases {\n\t\tt.Run(tc.name, func(t *testing.T) {\n\t\t\tif tc.skip != nil {\n\t\t\t\tif skip, msg := tc.skip(); skip {\n\t\t\t\t\tt.Skip(msg)\n\t\t\t\t}\n\t\t\t}\n\t\t\tconfig := &Config{\n\t\t\t\tAccount: account,\n\t\t\t\tHost: host,\n\t\t\t\tAuthenticator: AuthTypeWorkloadIdentityFederation,\n\t\t\t}\n\t\t\ttc.setupCfg(t, config)\n\t\t\tconnector := NewConnector(SnowflakeDriver{}, *config)\n\t\t\tdb := sql.OpenDB(connector)\n\t\t\tdefer db.Close()\n\t\t\tcurrentUser := runSelectCurrentUser(t, db)\n\t\t\tassertEqualE(t, currentUser, tc.expectedUsername)\n\t\t})\n\t}\n}\n"
},
{
"path": "auth_with_external_browser_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"context\"\n\t\"database/sql\"\n\t\"fmt\"\n\t\"log\"\n\t\"os/exec\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc TestExternalBrowserSuccessful(t *testing.T) {\n\tcfg := setupExternalBrowserTest(t)\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.Success, cfg.User, cfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n}\n\nfunc TestExternalBrowserFailed(t *testing.T) {\n\tcfg := setupExternalBrowserTest(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(10) * time.Second\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.Fail, \"FakeAccount\", \"NotARealPassword\")\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNotNilF(t, err)\n\t\tassertEqualE(t, err.Error(), \"authentication timed out\")\n\t}()\n\twg.Wait()\n}\n\nfunc TestExternalBrowserTimeout(t *testing.T) {\n\tcfg := setupExternalBrowserTest(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.Timeout, cfg.User, cfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNotNilF(t, err)\n\t\tassertEqualE(t, err.Error(), \"authentication timed out\")\n\t}()\n\twg.Wait()\n}\n\nfunc TestExternalBrowserMismatchUser(t *testing.T) {\n\tcfg := setupExternalBrowserTest(t)\n\tcorrectUsername := cfg.User\n\tcfg.User = \"fakeAccount\"\n\tvar wg sync.WaitGroup\n\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.Success, correctUsername, cfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tvar snowflakeErr *SnowflakeError\n\t\tassertErrorsAsF(t, err, &snowflakeErr)\n\t\tassertEqualE(t, snowflakeErr.Number, 390191, fmt.Sprintf(\"Expected 390191, but got %v\", snowflakeErr.Number))\n\t}()\n\twg.Wait()\n}\n\nfunc TestClientStoreCredentials(t *testing.T) {\n\tcfg := setupExternalBrowserTest(t)\n\tcfg.ClientStoreTemporaryCredential = 1\n\tcfg.ExternalBrowserTimeout = time.Duration(10) * time.Second\n\n\tt.Run(\"Obtains the ID token from the server and saves it on the local storage\", func(t *testing.T) {\n\t\tcleanupBrowserProcesses(t)\n\t\tvar wg sync.WaitGroup\n\t\twg.Add(2)\n\t\tgo func() {\n\t\t\tdefer wg.Done()\n\t\t\tprovideExternalBrowserCredentials(t, externalBrowserType.Success, cfg.User, cfg.Password)\n\t\t}()\n\t\tgo func() {\n\t\t\tdefer wg.Done()\n\t\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed: err %v\", err))\n\t\t}()\n\t\twg.Wait()\n\t})\n\n\tt.Run(\"Verify validation of ID token if option enabled\", func(t *testing.T) {\n\t\tcleanupBrowserProcesses(t)\n\t\tcfg.ClientStoreTemporaryCredential = 1\n\t\tdb := getDbHandlerFromConfig(t, cfg)\n\t\tconn, err := db.Conn(context.Background())\n\t\tassertNilE(t, err, fmt.Sprintf(\"Failed to connect to Snowflake. err: %v\", err))\n\t\tdefer conn.Close()\n\t\trows, err := conn.QueryContext(context.Background(), \"SELECT 1\")\n\t\tassertNilE(t, err, fmt.Sprintf(\"Failed to run a query. err: %v\", err))\n\t\trows.Close()\n\t})\n\n\tt.Run(\"Verify validation of idToken if option disabled\", func(t *testing.T) {\n\t\tcleanupBrowserProcesses(t)\n\t\tcfg.ClientStoreTemporaryCredential = 0\n\t\tdb := getDbHandlerFromConfig(t, cfg)\n\t\t_, err := db.Conn(context.Background())\n\t\tassertNotNilF(t, err)\n\t\tassertEqualE(t, err.Error(), \"authentication timed out\", fmt.Sprintf(\"Expected timeout, but got %v\", err))\n\t})\n}\n\ntype ExternalBrowserProcessResult struct {\n\tSuccess string\n\tFail string\n\tTimeout string\n\tOauthOktaSuccess string\n\tOauthSnowflakeSuccess string\n}\n\nvar externalBrowserType = ExternalBrowserProcessResult{\n\tSuccess: \"success\",\n\tFail: \"fail\",\n\tTimeout: \"timeout\",\n\tOauthOktaSuccess: \"externalOauthOktaSuccess\",\n\tOauthSnowflakeSuccess: \"internalOauthSnowflakeSuccess\",\n}\n\nfunc cleanupBrowserProcesses(t *testing.T) {\n\tif isTestRunningInDockerContainer() {\n\t\tconst cleanBrowserProcessesPath = \"/externalbrowser/cleanBrowserProcesses.js\"\n\t\t_, err := exec.Command(\"node\", cleanBrowserProcessesPath).CombinedOutput()\n\t\tassertNilE(t, err, fmt.Sprintf(\"failed to execute command: %v\", err))\n\t}\n}\n\nfunc provideExternalBrowserCredentials(t *testing.T, ExternalBrowserProcess string, user string, password string) {\n\tif isTestRunningInDockerContainer() {\n\t\tconst provideBrowserCredentialsPath = \"/externalbrowser/provideBrowserCredentials.js\"\n\t\toutput, err := exec.Command(\"node\", provideBrowserCredentialsPath, ExternalBrowserProcess, user, password).CombinedOutput()\n\t\tlog.Printf(\"Output: %s\\n\", output)\n\t\tassertNilE(t, err, fmt.Sprintf(\"failed to execute command: %v\", err))\n\t}\n}\n\nfunc verifyConnectionToSnowflakeAuthTests(t *testing.T, cfg *Config) (err error) {\n\tdsn, err := DSN(cfg)\n\tassertNilE(t, err, \"failed to create DSN from Config\")\n\n\tdb, err := sql.Open(\"snowflake\", dsn)\n\tassertNilE(t, err, \"failed to open Snowflake DB connection\")\n\tdefer db.Close()\n\n\trows, err := db.Query(\"SELECT 1\")\n\tif err != nil {\n\t\tlog.Printf(\"failed to run a query. 'SELECT 1', err: %v\", err)\n\t\treturn err\n\t}\n\tdefer rows.Close()\n\tassertTrueE(t, rows.Next(), \"failed to get result\", \"There were no results for query: \")\n\n\treturn err\n}\n\nfunc setupExternalBrowserTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping External Browser tests\")\n\tcleanupBrowserProcesses(t)\n\tcfg, err := getAuthTestsConfig(t, AuthTypeExternalBrowser)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get config: %v\", err))\n\treturn cfg\n}\n"
},
{
"path": "auth_with_keypair_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"crypto/rsa\"\n\t\"fmt\"\n\t\"golang.org/x/crypto/ssh\"\n\t\"os\"\n\t\"testing\"\n)\n\nfunc TestKeypairSuccessful(t *testing.T) {\n\tcfg := setupKeyPairTest(t)\n\tcfg.PrivateKey = loadRsaPrivateKeyForKeyPair(t, \"SNOWFLAKE_AUTH_TEST_PRIVATE_KEY_PATH\")\n\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"failed to connect. err: %v\", err))\n}\n\nfunc TestKeypairInvalidKey(t *testing.T) {\n\tcfg := setupKeyPairTest(t)\n\tcfg.PrivateKey = loadRsaPrivateKeyForKeyPair(t, \"SNOWFLAKE_AUTH_TEST_INVALID_PRIVATE_KEY_PATH\")\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 390144, fmt.Sprintf(\"Expected 390144, but got %v\", snowflakeErr.Number))\n}\n\nfunc setupKeyPairTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping KeyPair tests\")\n\tcfg, err := getAuthTestsConfig(t, AuthTypeJwt)\n\tassertEqualE(t, err, nil, fmt.Sprintf(\"failed to get config: %v\", err))\n\n\treturn cfg\n}\n\nfunc loadRsaPrivateKeyForKeyPair(t *testing.T, envName string) *rsa.PrivateKey {\n\tfilePath, err := GetFromEnv(envName, true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get env: %v\", err))\n\n\tbytes, err := os.ReadFile(filePath)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to read file: %v\", err))\n\n\tkey, err := ssh.ParseRawPrivateKey(bytes)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to parse private key: %v\", err))\n\n\treturn key.(*rsa.PrivateKey)\n}\n"
},
{
"path": "auth_with_mfa_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"log\"\n\t\"os/exec\"\n\t\"strings\"\n\t\"testing\"\n)\n\nfunc TestMfaSuccessful(t *testing.T) {\n\tcfg := setupMfaTest(t)\n\n\t// Enable MFA token caching\n\tcfg.ClientRequestMfaToken = ConfigBoolTrue\n\n\t//Provide your own TOTP code/codes here, to test manually\n\t//totpKeys := []string{\"222222\", \"333333\", \"444444\"}\n\n\ttotpKeys := getTOPTcodes(t)\n\n\tverifyConnectionToSnowflakeUsingTotpCodes(t, cfg, totpKeys)\n\tlog.Printf(\"Testing MFA token caching with second connection...\")\n\n\t// Clear the passcode to force use of cached MFA token\n\tcfg.Passcode = \"\"\n\n\t// Attempt to connect using cached MFA token\n\tcacheErr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilF(t, cacheErr, \"Failed to connect with cached MFA token\")\n}\n\nfunc setupMfaTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping MFA tests\")\n\tcfg, err := getAuthTestsConfig(t, AuthTypeUsernamePasswordMFA)\n\tassertNilF(t, err, \"failed to get config\")\n\n\tcfg.User, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_MFA_USER\", true)\n\tassertNilF(t, err, \"failed to get MFA user from environment\")\n\n\tcfg.Password, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_MFA_PASSWORD\", true)\n\tassertNilF(t, err, \"failed to get MFA password from environment\")\n\n\treturn cfg\n}\n\nfunc getTOPTcodes(t *testing.T) []string {\n\tif isTestRunningInDockerContainer() {\n\t\tconst provideTotpPath = \"/externalbrowser/totpGenerator.js\"\n\t\toutput, err := exec.Command(\"node\", provideTotpPath).CombinedOutput()\n\t\tassertNilF(t, err, fmt.Sprintf(\"failed to execute command: %v\", err))\n\t\ttotpCodes := strings.Fields(string(output))\n\t\treturn totpCodes\n\t}\n\treturn []string{}\n}\n\nfunc verifyConnectionToSnowflakeUsingTotpCodes(t *testing.T, cfg *Config, totpKeys []string) {\n\tif len(totpKeys) == 0 {\n\t\tt.Fatalf(\"no TOTP codes provided\")\n\t}\n\n\tvar lastError error\n\n\tfor i, totpKey := range totpKeys {\n\t\tcfg.Passcode = totpKey\n\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tif err == nil {\n\t\t\treturn\n\t\t}\n\n\t\tlastError = err\n\t\terrorMsg := err.Error()\n\n\t\tlog.Printf(\"TOTP code %d failed: %v\", i+1, errorMsg)\n\n\t\tvar snowflakeErr *SnowflakeError\n\t\tif errors.As(err, &snowflakeErr) && (snowflakeErr.Number == 394633 || snowflakeErr.Number == 394507) {\n\t\t\tlog.Printf(\"MFA error detected (%d), trying next code...\", snowflakeErr.Number)\n\t\t\tcontinue\n\t\t} else {\n\t\t\tlog.Printf(\"Non-MFA error detected: %v\", errorMsg)\n\t\t\tbreak\n\t\t}\n\t}\n\n\tassertNilF(t, lastError, \"failed to connect with any TOTP code\")\n}\n"
},
{
"path": "auth_with_oauth_okta_authorization_code_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"fmt\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc TestOauthOktaAuthorizationCodeSuccessful(t *testing.T) {\n\tcfg := setupOauthOktaAuthorizationCodeTest(t)\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthOktaSuccess, cfg.User, cfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n}\n\nfunc TestOauthOktaAuthorizationCodeMismatchedUsername(t *testing.T) {\n\tcfg := setupOauthOktaAuthorizationCodeTest(t)\n\tuser := cfg.User\n\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthOktaSuccess, user, cfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tcfg.User = \"fakeUser@snowflake.com\"\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tvar snowflakeErr *SnowflakeError\n\t\tassertErrorsAsF(t, err, &snowflakeErr)\n\t\tassertEqualE(t, snowflakeErr.Number, 390309, fmt.Sprintf(\"Expected 390309, but got %v\", snowflakeErr.Number))\n\t}()\n\twg.Wait()\n}\n\nfunc TestOauthOktaAuthorizationCodeOktaTimeout(t *testing.T) {\n\tcfg := setupOauthOktaAuthorizationCodeTest(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNotNilF(t, err, \"should failed due to timeout\")\n\tassertEqualE(t, err.Error(), \"authentication via browser timed out\", fmt.Sprintf(\"Expecteed timeout, but got %v\", err))\n}\n\nfunc TestOauthOktaAuthorizationCodeUsingTokenCache(t *testing.T) {\n\tcfg := setupOauthOktaAuthorizationCodeTest(t)\n\tcfg.ClientStoreTemporaryCredential = 1\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthOktaSuccess, cfg.User, cfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n\n\tcleanupBrowserProcesses(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n}\n\nfunc setupOauthOktaAuthorizationCodeTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping Okta Authorization Code tests\")\n\tcfg, err := getAuthTestsConfig(t, AuthTypeOAuthAuthorizationCode)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get config: %v\", err))\n\n\tcleanupBrowserProcesses(t)\n\n\tcfg.OauthClientID, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthClientSecret, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_SECRET\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthRedirectURI, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_REDIRECT_URI\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthAuthorizationURL, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_AUTH_URL\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthTokenRequestURL, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_TOKEN\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.Role, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_ROLE\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\treturn cfg\n\n}\n"
},
{
"path": "auth_with_oauth_okta_client_credentials_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n)\n\nfunc TestOauthOktaClientCredentialsSuccessful(t *testing.T) {\n\tcfg := setupOauthOktaClientCredentialsTest(t)\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"failed to connect. err: %v\", err))\n}\n\nfunc TestOauthOktaClientCredentialsMismatchedUsername(t *testing.T) {\n\tcfg := setupOauthOktaClientCredentialsTest(t)\n\tcfg.User = \"invalidUser\"\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 390309, fmt.Sprintf(\"Expected 390309, but got %v\", snowflakeErr.Number))\n}\n\nfunc TestOauthOktaClientCredentialsUnauthorized(t *testing.T) {\n\tcfg := setupOauthOktaClientCredentialsTest(t)\n\tcfg.OauthClientID = \"invalidClientID\"\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNotNilF(t, err, \"Expected an error but got nil\")\n\tassertTrueF(t, strings.Contains(err.Error(), \"invalid_client\"), fmt.Sprintf(\"Expected error to contain 'invalid_client', but got: %v\", err.Error()))\n}\n\nfunc setupOauthOktaClientCredentialsTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping Okta Client Credentials tests\")\n\n\tcfg, err := getAuthTestsConfig(t, AuthTypeOAuthClientCredentials)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get config: %v\", err))\n\n\tcfg.OauthClientID, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthClientSecret, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_SECRET\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthTokenRequestURL, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_TOKEN\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.User, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.Role, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_ROLE\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\treturn cfg\n}\n"
},
{
"path": "auth_with_oauth_snowflake_authorization_code_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"fmt\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc TestOauthSnowflakeAuthorizationCodeSuccessful(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeTest(t)\n\tbrowserCfg, err := getOauthSnowflakeAuthorizationCodeTestCredentials()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get browser config: %v\", err))\n\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthSnowflakeSuccess, browserCfg.User, browserCfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n}\n\nfunc TestOauthSnowflakeAuthorizationCodeMismatchedUsername(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeTest(t)\n\tbrowserCfg, err := getOauthSnowflakeAuthorizationCodeTestCredentials()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get browser config: %v\", err))\n\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthSnowflakeSuccess, browserCfg.User, browserCfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tcfg.User = \"fakeUser@snowflake.com\"\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tvar snowflakeErr *SnowflakeError\n\t\tassertErrorsAsF(t, err, &snowflakeErr)\n\t\tassertEqualE(t, snowflakeErr.Number, 390309, fmt.Sprintf(\"Expected 390309, but got %v\", snowflakeErr.Number))\n\t}()\n\twg.Wait()\n}\n\nfunc TestOauthSnowflakeAuthorizationCodeTimeout(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeTest(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNotNilF(t, err, \"should failed due to timeout\")\n\tassertEqualE(t, err.Error(), \"authentication via browser timed out\", fmt.Sprintf(\"Expecteed timeout, but got %v\", err))\n}\n\nfunc TestOauthSnowflakeAuthorizationCodeUsingTokenCache(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeTest(t)\n\tbrowserCfg, err := getOauthSnowflakeAuthorizationCodeTestCredentials()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get browser config: %v\", err))\n\n\tcfg.ClientStoreTemporaryCredential = 1\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthSnowflakeSuccess, browserCfg.User, browserCfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n\n\tcleanupBrowserProcesses(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\n\terr = verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n}\n\nfunc TestOauthSnowflakeAuthorizationCodeWithoutTokenCache(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeTest(t)\n\tbrowserCfg, err := getOauthSnowflakeAuthorizationCodeTestCredentials()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get browser config: %v\", err))\n\tcfg.ClientStoreTemporaryCredential = 2\n\n\tvar wg sync.WaitGroup\n\tcfg.DisableQueryContextCache = true\n\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthSnowflakeSuccess, browserCfg.User, browserCfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n\n\tcleanupBrowserProcesses(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\n\terr = verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNotNilF(t, err, \"Expected an error but got nil\")\n\tassertEqualE(t, err.Error(), \"authentication via browser timed out\", fmt.Sprintf(\"Expecteed timeout, but got %v\", err))\n}\n\nfunc setupOauthSnowflakeAuthorizationCodeTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping Snowflake Authorization Code tests\")\n\n\tcfg, err := getAuthTestsConfig(t, AuthTypeOAuthAuthorizationCode)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get config: %v\", err))\n\n\tcleanupBrowserProcesses(t)\n\n\tcfg.OauthClientID, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_CLIENT_ID\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthClientSecret, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_CLIENT_SECRET\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthRedirectURI, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_REDIRECT_URI\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.User, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.Role, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_ROLE\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.ClientStoreTemporaryCredential = 2\n\treturn cfg\n}\n\nfunc getOauthSnowflakeAuthorizationCodeTestCredentials() (*Config, error) {\n\treturn GetConfigFromEnv([]*ConfigParam{\n\t\t{Name: \"User\", EnvName: \"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID\", FailOnMissing: true},\n\t\t{Name: \"Password\", EnvName: \"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_USER_PASSWORD\", FailOnMissing: true},\n\t})\n}\n"
},
{
"path": "auth_with_oauth_snowflake_authorization_code_wildcards_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"fmt\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc TestOauthSnowflakeAuthorizationCodeWildcardsSuccessful(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeWildcardsTest(t)\n\tbrowserCfg, err := getOauthSnowflakeAuthorizationCodeTestCredentials()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get browser config: %v\", err))\n\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthSnowflakeSuccess, browserCfg.User, browserCfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n}\n\nfunc TestOauthSnowflakeAuthorizationCodeWildcardsMismatchedUsername(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeWildcardsTest(t)\n\tbrowserCfg, err := getOauthSnowflakeAuthorizationCodeTestCredentials()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get browser config: %v\", err))\n\n\tvar wg sync.WaitGroup\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthSnowflakeSuccess, browserCfg.User, browserCfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tcfg.User = \"fakeUser@snowflake.com\"\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tvar snowflakeErr *SnowflakeError\n\t\tassertErrorsAsF(t, err, &snowflakeErr)\n\t\tassertEqualE(t, snowflakeErr.Number, 390309, fmt.Sprintf(\"Expected 390309, but got %v\", snowflakeErr.Number))\n\t}()\n\twg.Wait()\n}\n\nfunc TestOauthSnowflakeAuthorizationWildcardsCodeTimeout(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeWildcardsTest(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNotNilF(t, err, \"should failed due to timeout\")\n\tassertEqualE(t, err.Error(), \"authentication via browser timed out\", fmt.Sprintf(\"Expecteed timeout, but got %v\", err))\n}\n\nfunc TestOauthSnowflakeAuthorizationCodeWildcardsWithoutTokenCache(t *testing.T) {\n\tcfg := setupOauthSnowflakeAuthorizationCodeWildcardsTest(t)\n\tbrowserCfg, err := getOauthSnowflakeAuthorizationCodeTestCredentials()\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get browser config: %v\", err))\n\tcfg.ClientStoreTemporaryCredential = 2\n\n\tvar wg sync.WaitGroup\n\tcfg.DisableQueryContextCache = true\n\n\twg.Add(2)\n\tgo func() {\n\t\tdefer wg.Done()\n\t\tprovideExternalBrowserCredentials(t, externalBrowserType.OauthSnowflakeSuccess, browserCfg.User, browserCfg.Password)\n\t}()\n\tgo func() {\n\t\tdefer wg.Done()\n\t\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\t\tassertNilE(t, err, fmt.Sprintf(\"Connection failed due to %v\", err))\n\t}()\n\twg.Wait()\n\n\tcleanupBrowserProcesses(t)\n\tcfg.ExternalBrowserTimeout = time.Duration(1) * time.Second\n\n\terr = verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNotNilF(t, err, \"Expected an error but got nil\")\n\tassertEqualE(t, err.Error(), \"authentication via browser timed out\", fmt.Sprintf(\"Expecteed timeout, but got %v\", err))\n}\n\nfunc setupOauthSnowflakeAuthorizationCodeWildcardsTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping Snowflake Authorization Code tests\")\n\n\tcfg, err := getAuthTestsConfig(t, AuthTypeOAuthAuthorizationCode)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get config: %v\", err))\n\n\tcleanupBrowserProcesses(t)\n\n\tcfg.OauthClientID, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_WILDCARDS_CLIENT_ID\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.OauthClientSecret, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_WILDCARDS_CLIENT_SECRET\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.User, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.Role, err = GetFromEnv(\"SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_ROLE\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to setup config: %v\", err))\n\n\tcfg.ClientStoreTemporaryCredential = 2\n\treturn cfg\n}\n"
},
{
"path": "auth_with_oauth_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"strings\"\n\t\"testing\"\n)\n\nfunc TestOauthSuccessful(t *testing.T) {\n\tcfg := setupOauthTest(t)\n\ttoken, err := getOauthTestToken(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"failed to get token. err: %v\", err))\n\tcfg.Token = token\n\terr = verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"failed to connect. err: %v\", err))\n}\n\nfunc TestOauthInvalidToken(t *testing.T) {\n\tcfg := setupOauthTest(t)\n\tcfg.Token = \"invalid_token\"\n\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 390303, fmt.Sprintf(\"Expected 390303, but got %v\", snowflakeErr.Number))\n}\n\nfunc TestOauthMismatchedUser(t *testing.T) {\n\tcfg := setupOauthTest(t)\n\ttoken, err := getOauthTestToken(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"failed to get token. err: %v\", err))\n\tcfg.Token = token\n\tcfg.User = \"fakeaccount\"\n\n\terr = verifyConnectionToSnowflakeAuthTests(t, cfg)\n\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 390309, fmt.Sprintf(\"Expected 390309, but got %v\", snowflakeErr.Number))\n}\n\nfunc setupOauthTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping OAuth tests\")\n\tcfg, err := getAuthTestsConfig(t, AuthTypeOAuth)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to connect. err: %v\", err))\n\n\treturn cfg\n}\n\nfunc getOauthTestToken(t *testing.T, cfg *Config) (string, error) {\n\n\tclient := &http.Client{}\n\n\tauthURL, err := GetFromEnv(\"SNOWFLAKE_AUTH_TEST_OAUTH_URL\", true)\n\tassertNilF(t, err, \"SNOWFLAKE_AUTH_TEST_OAUTH_URL is not set\")\n\n\toauthClientID, err := GetFromEnv(\"SNOWFLAKE_AUTH_TEST_OAUTH_CLIENT_ID\", true)\n\tassertNilF(t, err, \"SNOWFLAKE_AUTH_TEST_OAUTH_CLIENT_ID is not set\")\n\n\toauthClientSecret, err := GetFromEnv(\"SNOWFLAKE_AUTH_TEST_OAUTH_CLIENT_SECRET\", true)\n\tassertNilF(t, err, \"SNOWFLAKE_AUTH_TEST_OAUTH_CLIENT_SECRET is not set\")\n\n\tinputData := formData(cfg)\n\n\treq, err := http.NewRequest(\"POST\", authURL, strings.NewReader(inputData.Encode()))\n\tassertNilF(t, err, fmt.Sprintf(\"Request failed %v\", err))\n\treq.Header.Set(\"Content-Type\", \"application/x-www-form-urlencoded;charset=UTF-8\")\n\treq.SetBasicAuth(oauthClientID, oauthClientSecret)\n\tresp, err := client.Do(req)\n\n\tassertNilF(t, err, fmt.Sprintf(\"Response failed %v\", err))\n\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn \"\", fmt.Errorf(\"failed to get access token, status code: %d\", resp.StatusCode)\n\t}\n\n\tdefer resp.Body.Close()\n\n\tvar response OAuthTokenResponse\n\tif err := json.NewDecoder(resp.Body).Decode(&response); err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to decode response: %v\", err)\n\t}\n\n\treturn response.Token, err\n}\n\nfunc formData(cfg *Config) url.Values {\n\tdata := url.Values{}\n\tdata.Set(\"username\", cfg.User)\n\tdata.Set(\"password\", cfg.Password)\n\tdata.Set(\"grant_type\", \"password\")\n\tdata.Set(\"scope\", fmt.Sprintf(\"session:role:%s\", strings.ToLower(cfg.Role)))\n\n\treturn data\n\n}\n\ntype OAuthTokenResponse struct {\n\tType string `json:\"token_type\"`\n\tExpiration int `json:\"expires_in\"`\n\tToken string `json:\"access_token\"`\n\tScope string `json:\"scope\"`\n}\n"
},
{
"path": "auth_with_okta_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"fmt\"\n\t\"net/url\"\n\t\"testing\"\n)\n\nfunc TestOktaSuccessful(t *testing.T) {\n\tcfg := setupOktaTest(t)\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"failed to connect. err: %v\", err))\n}\n\nfunc TestOktaWrongCredentials(t *testing.T) {\n\tcfg := setupOktaTest(t)\n\tcfg.Password = \"fakePassword\"\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 261006, fmt.Sprintf(\"Expected 261006, but got %v\", snowflakeErr.Number))\n}\n\nfunc TestOktaWrongAuthenticator(t *testing.T) {\n\tcfg := setupOktaTest(t)\n\tinvalidAddress, err := url.Parse(\"https://fake-account-0000.okta.com\")\n\tassertNilF(t, err, fmt.Sprintf(\"failed to parse: %v\", err))\n\n\tcfg.OktaURL = invalidAddress\n\terr = verifyConnectionToSnowflakeAuthTests(t, cfg)\n\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 390139, fmt.Sprintf(\"Expected 390139, but got %v\", snowflakeErr.Number))\n}\n\nfunc setupOktaTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping Okta tests\")\n\turlEnv, err := GetFromEnv(\"SNOWFLAKE_AUTH_TEST_OKTA_AUTH\", true)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get env: %v\", err))\n\n\tcfg, err := getAuthTestsConfig(t, AuthTypeOkta)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to get config: %v\", err))\n\n\tcfg.OktaURL, err = url.Parse(urlEnv)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to parse: %v\", err))\n\n\treturn cfg\n}\n"
},
{
"path": "auth_with_pat_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"database/sql\"\n\t\"fmt\"\n\t\"log\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n)\n\ntype PatToken struct {\n\tName string\n\tValue string\n}\n\nfunc TestEndToEndPatSuccessful(t *testing.T) {\n\tcfg := setupEndToEndPatTest(t)\n\tpatToken := createEndToEndPatToken(t)\n\tdefer removeEndToEndPatToken(t, patToken.Name)\n\tcfg.Token = patToken.Value\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tassertNilE(t, err, fmt.Sprintf(\"failed to connect. err: %v\", err))\n}\n\nfunc TestEndToEndPatMismatchedUser(t *testing.T) {\n\tcfg := setupEndToEndPatTest(t)\n\tpatToken := createEndToEndPatToken(t)\n\tdefer removeEndToEndPatToken(t, patToken.Name)\n\tcfg.Token = patToken.Value\n\tcfg.User = \"invalidUser\"\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 394400, fmt.Sprintf(\"Expected 394400, but got %v\", snowflakeErr.Number))\n}\n\nfunc TestEndToEndPatInvalidToken(t *testing.T) {\n\tcfg := setupEndToEndPatTest(t)\n\tcfg.Token = \"invalidToken\"\n\terr := verifyConnectionToSnowflakeAuthTests(t, cfg)\n\tvar snowflakeErr *SnowflakeError\n\tassertErrorsAsF(t, err, &snowflakeErr)\n\tassertEqualE(t, snowflakeErr.Number, 394400, fmt.Sprintf(\"Expected 394400, but got %v\", snowflakeErr.Number))\n}\n\nfunc setupEndToEndPatTest(t *testing.T) *Config {\n\tskipAuthTests(t, \"Skipping PAT tests\")\n\tcfg, err := getAuthTestsConfig(t, AuthTypePat)\n\tassertNilF(t, err, fmt.Sprintf(\"failed to parse: %v\", err))\n\n\treturn cfg\n\n}\n\nfunc getEndToEndPatSetupCommandVariables() (*Config, error) {\n\treturn GetConfigFromEnv([]*ConfigParam{\n\t\t{Name: \"User\", EnvName: \"SNOWFLAKE_AUTH_TEST_SNOWFLAKE_USER\", FailOnMissing: true},\n\t\t{Name: \"Role\", EnvName: \"SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_ROLE\", FailOnMissing: true},\n\t})\n}\n\nfunc createEndToEndPatToken(t *testing.T) *PatToken {\n\tcfg := setupOktaTest(t)\n\tpatTokenName := fmt.Sprintf(\"PAT_GOLANG_%s\", strings.ReplaceAll(time.Now().Format(\"20060102150405.000\"), \".\", \"\"))\n\tpatCommandVariables, err := getEndToEndPatSetupCommandVariables()\n\tassertNilE(t, err, \"failed to get PAT command variables\")\n\n\tquery := fmt.Sprintf(\n\t\t\"alter user %s add programmatic access token %s ROLE_RESTRICTION = '%s' DAYS_TO_EXPIRY=1;\",\n\t\tpatCommandVariables.User,\n\t\tpatTokenName,\n\t\tpatCommandVariables.Role,\n\t)\n\n\tpatToken, err := connectUsingOktaConnectionAndExecuteCustomCommand(t, cfg, query, true)\n\tassertNilE(t, err, \"failed to create PAT command variables\")\n\n\treturn patToken\n\n}\n\nfunc removeEndToEndPatToken(t *testing.T, patTokenName string) {\n\tcfg := setupOktaTest(t)\n\tcfg.Role = \"analyst\"\n\tpatCommandVariables, err := getEndToEndPatSetupCommandVariables()\n\tassertNilE(t, err, \"failed to get PAT command variables\")\n\n\tquery := fmt.Sprintf(\n\t\t\"alter user %s remove programmatic access token %s;\",\n\t\tpatCommandVariables.User,\n\t\tpatTokenName,\n\t)\n\n\t_, err = connectUsingOktaConnectionAndExecuteCustomCommand(t, cfg, query, false)\n\tassertNilE(t, err, \"failed to remove PAT command variables\")\n}\n\nfunc connectUsingOktaConnectionAndExecuteCustomCommand(t *testing.T, cfg *Config, query string, returnToken bool) (*PatToken, error) {\n\tdsn, err := DSN(cfg)\n\tassertNilE(t, err, \"failed to create DSN from Config\")\n\n\tdb, err := sql.Open(\"snowflake\", dsn)\n\tassertNilE(t, err, \"failed to open Snowflake DB connection\")\n\tdefer db.Close()\n\n\trows, err := db.Query(query)\n\tif err != nil {\n\t\tlog.Printf(\"failed to run a query: %v, err: %v\", query, err)\n\t\treturn nil, err\n\n\t}\n\n\tvar patTokenName, patTokenValue string\n\tif returnToken && rows.Next() {\n\t\tif err := rows.Scan(&patTokenName, &patTokenValue); err != nil {\n\t\t\tt.Fatalf(\"failed to scan token: %v\", err)\n\t\t}\n\n\t\treturn &PatToken{Name: patTokenName, Value: patTokenValue}, nil\n\t}\n\n\tif returnToken {\n\t\tt.Fatalf(\"no results found for query: %s\", query)\n\t}\n\n\treturn nil, err\n}\n"
},
{
"path": "authexternalbrowser.go",
"content": "package gosnowflake\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/base64\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\terrors2 \"github.com/snowflakedb/gosnowflake/v2/internal/errors\"\n\t\"io\"\n\t\"log\"\n\t\"net\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"strconv\"\n\t\"strings\"\n\t\"time\"\n\n\t\"github.com/pkg/browser\"\n)\n\nconst (\n\tsamlSuccessHTML = `\nSAML Response for Snowflake\n\nYour identity was confirmed and propagated to Snowflake %v.\nYou can close this window now and go back where you started from.\n`\n\n\tbufSize = 8192\n)\n\n// Builds a response to show to the user after successfully\n// getting a response from Snowflake.\nfunc buildResponse(body string) (bytes.Buffer, error) {\n\tt := &http.Response{\n\t\tStatus: \"200 OK\",\n\t\tStatusCode: 200,\n\t\tProto: \"HTTP/1.1\",\n\t\tProtoMajor: 1,\n\t\tProtoMinor: 1,\n\t\tBody: io.NopCloser(bytes.NewBufferString(body)),\n\t\tContentLength: int64(len(body)),\n\t\tRequest: nil,\n\t\tHeader: make(http.Header),\n\t}\n\tvar b bytes.Buffer\n\terr := t.Write(&b)\n\treturn b, err\n}\n\n// This opens a socket that listens on all available unicast\n// and any anycast IP addresses locally. By specifying \"0\", we are\n// able to bind to a free port.\nfunc createLocalTCPListener(port int) (*net.TCPListener, error) {\n\tlogger.Debugf(\"creating local TCP listener on port %v\", port)\n\tallAddressesListener, err := net.Listen(\"tcp\", fmt.Sprintf(\"0.0.0.0:%v\", port))\n\tif err != nil {\n\t\tlogger.Warnf(\"error while setting up 0.0.0.0 listener: %v\", err)\n\t\treturn nil, err\n\t}\n\tlogger.Debug(\"Closing 0.0.0.0 tcp listener\")\n\tif err := allAddressesListener.Close(); err != nil {\n\t\tlogger.Errorf(\"error while closing TCP listener. %v\", err)\n\t\treturn nil, err\n\t}\n\n\tl, err := net.Listen(\"tcp\", fmt.Sprintf(\"localhost:%v\", port))\n\tif err != nil {\n\t\tlogger.Warnf(\"error while setting up listener: %v\", err)\n\t\treturn nil, err\n\t}\n\n\ttcpListener, ok := l.(*net.TCPListener)\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"failed to assert type as *net.TCPListener\")\n\t}\n\n\treturn tcpListener, nil\n}\n\n// Opens a browser window (or new tab) with the configured login Url.\n// This can / will fail if running inside a shell with no display, ie\n// ssh'ing into a box attempting to authenticate via external browser.\nfunc openBrowser(browserURL string) error {\n\tparsedURL, err := url.ParseRequestURI(browserURL)\n\tif err != nil {\n\t\tlogger.Errorf(\"error parsing url %v, err: %v\", browserURL, err)\n\t\treturn err\n\t}\n\tif parsedURL.Scheme != \"http\" && parsedURL.Scheme != \"https\" {\n\t\treturn fmt.Errorf(\"invalid browser URL: %v\", browserURL)\n\t}\n\terr = browser.OpenURL(browserURL)\n\tif err != nil {\n\t\tlogger.Errorf(\"failed to open a browser. err: %v\", err)\n\t\treturn err\n\t}\n\treturn nil\n}\n\n// Gets the IDP Url and Proof Key from Snowflake.\n// Note: FuncPostAuthSaml will return a fully qualified error if\n// there is something wrong getting data from Snowflake.\nfunc getIdpURLProofKey(\n\tctx context.Context,\n\tsr *snowflakeRestful,\n\tauthenticator string,\n\tapplication string,\n\taccount string,\n\tuser string,\n\tcallbackPort int) (string, string, error) {\n\n\theaders := make(map[string]string)\n\theaders[httpHeaderContentType] = headerContentTypeApplicationJSON\n\theaders[httpHeaderAccept] = headerContentTypeApplicationJSON\n\theaders[httpHeaderUserAgent] = userAgent\n\n\tclientEnvironment := newAuthRequestClientEnvironment()\n\tclientEnvironment.Application = application\n\n\trequestMain := authRequestData{\n\t\tClientAppID: clientType,\n\t\tClientAppVersion: SnowflakeGoDriverVersion,\n\t\tAccountName: account,\n\t\tLoginName: user,\n\t\tClientEnvironment: clientEnvironment,\n\t\tAuthenticator: authenticator,\n\t\tBrowserModeRedirectPort: strconv.Itoa(callbackPort),\n\t}\n\n\tauthRequest := authRequest{\n\t\tData: requestMain,\n\t}\n\n\tjsonBody, err := json.Marshal(authRequest)\n\tif err != nil {\n\t\tlogger.WithContext(ctx).Errorf(\"failed to serialize json. err: %v\", err)\n\t\treturn \"\", \"\", err\n\t}\n\n\trespd, err := sr.FuncPostAuthSAML(ctx, sr, headers, jsonBody, sr.LoginTimeout)\n\tif err != nil {\n\t\treturn \"\", \"\", err\n\t}\n\tif !respd.Success {\n\t\tlogger.WithContext(ctx).Error(\"Authentication FAILED\")\n\t\tsr.TokenAccessor.SetTokens(\"\", \"\", -1)\n\t\tcode, err := strconv.Atoi(respd.Code)\n\t\tif err != nil {\n\t\t\treturn \"\", \"\", err\n\t\t}\n\t\treturn \"\", \"\", &SnowflakeError{\n\t\t\tNumber: code,\n\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\tMessage: respd.Message,\n\t\t}\n\t}\n\treturn respd.Data.SSOURL, respd.Data.ProofKey, nil\n}\n\n// Gets the login URL for multiple SAML\nfunc getLoginURL(sr *snowflakeRestful, user string, callbackPort int) (string, string, error) {\n\tproofKey := generateProofKey()\n\n\tparams := &url.Values{}\n\tparams.Add(\"login_name\", user)\n\tparams.Add(\"browser_mode_redirect_port\", strconv.Itoa(callbackPort))\n\tparams.Add(\"proof_key\", proofKey)\n\turl := sr.getFullURL(consoleLoginRequestPath, params)\n\n\treturn url.String(), proofKey, nil\n}\n\nfunc generateProofKey() string {\n\trandomness := getSecureRandom(32)\n\treturn base64.StdEncoding.WithPadding(base64.StdPadding).EncodeToString(randomness)\n}\n\n// The response returned from Snowflake looks like so:\n// GET /?token=encodedSamlToken\n// Host: localhost:54001\n// Connection: keep-alive\n// Upgrade-Insecure-Requests: 1\n// User-Agent: userAgentStr\n// Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8\n// Referer: https://myaccount.snowflakecomputing.com/fed/login\n// Accept-Encoding: gzip, deflate, br\n// Accept-Language: en-US,en;q=0.9\n// This extracts the token portion of the response.\nfunc getTokenFromResponse(response string) (string, error) {\n\tstart := \"GET /?token=\"\n\tarr := strings.Split(response, \"\\r\\n\")\n\tif !strings.HasPrefix(arr[0], start) {\n\t\tlogger.Errorf(\"response is malformed. \")\n\t\treturn \"\", &SnowflakeError{\n\t\t\tNumber: ErrFailedToParseResponse,\n\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\tMessage: errors2.ErrMsgFailedToParseResponse,\n\t\t\tMessageArgs: []any{response},\n\t\t}\n\t}\n\ttoken := strings.TrimPrefix(arr[0], start)\n\ttoken = strings.Split(token, \" \")[0]\n\treturn token, nil\n}\n\ntype authenticateByExternalBrowserResult struct {\n\tescapedSamlResponse []byte\n\tproofKey []byte\n\terr error\n}\n\nfunc authenticateByExternalBrowser(ctx context.Context, sr *snowflakeRestful, authenticator string, application string,\n\taccount string, user string, externalBrowserTimeout time.Duration, disableConsoleLogin ConfigBool) ([]byte, []byte, error) {\n\tresultChan := make(chan authenticateByExternalBrowserResult, 1)\n\tgo GoroutineWrapper(\n\t\tctx,\n\t\tfunc() {\n\t\t\tresultChan <- doAuthenticateByExternalBrowser(ctx, sr, authenticator, application, account, user, disableConsoleLogin)\n\t\t},\n\t)\n\tselect {\n\tcase <-time.After(externalBrowserTimeout):\n\t\treturn nil, nil, errors.New(\"authentication timed out\")\n\tcase result := <-resultChan:\n\t\treturn result.escapedSamlResponse, result.proofKey, result.err\n\t}\n}\n\n// Authentication by an external browser takes place via the following:\n// - the golang snowflake driver communicates to Snowflake that the user wishes to\n// authenticate via external browser\n// - snowflake sends back the IDP Url configured at the Snowflake side for the\n// provided account, or use the multiple SAML way via console login\n// - the default browser is opened to that URL\n// - user authenticates at the IDP, and is redirected to Snowflake\n// - Snowflake directs the user back to the driver\n// - authenticate is complete!\nfunc doAuthenticateByExternalBrowser(ctx context.Context, sr *snowflakeRestful, authenticator string, application string, account string, user string, disableConsoleLogin ConfigBool) authenticateByExternalBrowserResult {\n\tl, err := createLocalTCPListener(0)\n\tif err != nil {\n\t\treturn authenticateByExternalBrowserResult{nil, nil, err}\n\t}\n\tdefer func() {\n\t\tif err = l.Close(); err != nil {\n\t\t\tlogger.Errorf(\"error while closing TCP listener for external browser (%v). %v\", l.Addr().String(), err)\n\t\t}\n\t}()\n\n\tcallbackPort := l.Addr().(*net.TCPAddr).Port\n\n\tvar loginURL string\n\tvar proofKey string\n\tif disableConsoleLogin == ConfigBoolTrue {\n\t\t// Gets the IDP URL and Proof Key from Snowflake\n\t\tloginURL, proofKey, err = getIdpURLProofKey(ctx, sr, authenticator, application, account, user, callbackPort)\n\t} else {\n\t\t// Multiple SAML way to do authentication via console login\n\t\tloginURL, proofKey, err = getLoginURL(sr, user, callbackPort)\n\t}\n\n\tif err != nil {\n\t\treturn authenticateByExternalBrowserResult{nil, nil, err}\n\t}\n\n\tif err = defaultSamlResponseProvider().run(loginURL); err != nil {\n\t\treturn authenticateByExternalBrowserResult{nil, nil, err}\n\t}\n\n\tencodedSamlResponseChan := make(chan string)\n\terrChan := make(chan error)\n\n\tvar encodedSamlResponse string\n\tvar errFromGoroutine error\n\tconn, err := l.Accept()\n\tif err != nil {\n\t\tlogger.WithContext(ctx).Errorf(\"unable to accept connection. err: %v\", err)\n\t\tlog.Fatal(err)\n\t}\n\tgo func(c net.Conn) {\n\t\tvar buf bytes.Buffer\n\t\ttotal := 0\n\t\tencodedSamlResponse := \"\"\n\t\tvar errAccept error\n\t\tfor {\n\t\t\tb := make([]byte, bufSize)\n\t\t\tn, err := c.Read(b)\n\t\t\tif err != nil {\n\t\t\t\tif err != io.EOF {\n\t\t\t\t\tlogger.WithContext(ctx).Infof(\"error reading from socket. err: %v\", err)\n\t\t\t\t\terrAccept = &SnowflakeError{\n\t\t\t\t\t\tNumber: ErrFailedToGetExternalBrowserResponse,\n\t\t\t\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\t\t\t\tMessage: errors2.ErrMsgFailedToGetExternalBrowserResponse,\n\t\t\t\t\t\tMessageArgs: []any{err},\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\tbreak\n\t\t\t}\n\t\t\ttotal += n\n\t\t\tbuf.Write(b)\n\t\t\tif n < bufSize {\n\t\t\t\t// We successfully read all data\n\t\t\t\ts := string(buf.Bytes()[:total])\n\t\t\t\tencodedSamlResponse, errAccept = getTokenFromResponse(s)\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tbuf.Grow(bufSize)\n\t\t}\n\t\tif encodedSamlResponse != \"\" {\n\t\t\tbody := fmt.Sprintf(samlSuccessHTML, application)\n\t\t\thttpResponse, err := buildResponse(body)\n\t\t\tif err != nil && errAccept == nil {\n\t\t\t\terrAccept = err\n\t\t\t}\n\t\t\tif _, err = c.Write(httpResponse.Bytes()); err != nil && errAccept == nil {\n\t\t\t\terrAccept = err\n\t\t\t}\n\t\t}\n\t\tif err := c.Close(); err != nil {\n\t\t\tlogger.Warnf(\"error while closing browser connection. %v\", err)\n\t\t}\n\t\tencodedSamlResponseChan <- encodedSamlResponse\n\t\terrChan <- errAccept\n\t}(conn)\n\n\tencodedSamlResponse = <-encodedSamlResponseChan\n\terrFromGoroutine = <-errChan\n\n\tif errFromGoroutine != nil {\n\t\treturn authenticateByExternalBrowserResult{nil, nil, errFromGoroutine}\n\t}\n\n\tescapedSamlResponse, err := url.QueryUnescape(encodedSamlResponse)\n\tif err != nil {\n\t\tlogger.WithContext(ctx).Errorf(\"unable to unescape saml response. err: %v\", err)\n\t\treturn authenticateByExternalBrowserResult{nil, nil, err}\n\t}\n\treturn authenticateByExternalBrowserResult{[]byte(escapedSamlResponse), []byte(proofKey), nil}\n}\n\ntype samlResponseProvider interface {\n\trun(url string) error\n}\n\ntype externalBrowserSamlResponseProvider struct {\n}\n\nfunc (e externalBrowserSamlResponseProvider) run(url string) error {\n\treturn openBrowser(url)\n}\n\nvar defaultSamlResponseProvider = func() samlResponseProvider {\n\treturn &externalBrowserSamlResponseProvider{}\n}\n"
},
{
"path": "authexternalbrowser_test.go",
"content": "package gosnowflake\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\tsfconfig \"github.com/snowflakedb/gosnowflake/v2/internal/config\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n)\n\nfunc TestGetTokenFromResponseFail(t *testing.T) {\n\tresponse := \"GET /?fakeToken=fakeEncodedSamlToken HTTP/1.1\\r\\n\" +\n\t\t\"Host: localhost:54001\\r\\n\" +\n\t\t\"Connection: keep-alive\\r\\n\" +\n\t\t\"Upgrade-Insecure-Requests: 1\\r\\n\" +\n\t\t\"User-Agent: userAgentStr\\r\\n\" +\n\t\t\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8\\r\\n\" +\n\t\t\"Referer: https://myaccount.snowflakecomputing.com/fed/login\\r\\n\" +\n\t\t\"Accept-Encoding: gzip, deflate, br\\r\\n\" +\n\t\t\"Accept-Language: en-US,en;q=0.9\\r\\n\\r\\n\"\n\n\t_, err := getTokenFromResponse(response)\n\tif err == nil {\n\t\tt.Errorf(\"Should have failed parsing the malformed response.\")\n\t}\n}\n\nfunc TestGetTokenFromResponse(t *testing.T) {\n\tresponse := \"GET /?token=GETtokenFromResponse HTTP/1.1\\r\\n\" +\n\t\t\"Host: localhost:54001\\r\\n\" +\n\t\t\"Connection: keep-alive\\r\\n\" +\n\t\t\"Upgrade-Insecure-Requests: 1\\r\\n\" +\n\t\t\"User-Agent: userAgentStr\\r\\n\" +\n\t\t\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8\\r\\n\" +\n\t\t\"Referer: https://myaccount.snowflakecomputing.com/fed/login\\r\\n\" +\n\t\t\"Accept-Encoding: gzip, deflate, br\\r\\n\" +\n\t\t\"Accept-Language: en-US,en;q=0.9\\r\\n\\r\\n\"\n\n\texpected := \"GETtokenFromResponse\"\n\n\ttoken, err := getTokenFromResponse(response)\n\tif err != nil {\n\t\tt.Errorf(\"Failed to get the token. Err: %#v\", err)\n\t}\n\tif token != expected {\n\t\tt.Errorf(\"Expected: %s, found: %s\", expected, token)\n\t}\n}\n\nfunc TestBuildResponse(t *testing.T) {\n\tresp, err := buildResponse(fmt.Sprintf(samlSuccessHTML, \"Go\"))\n\tassertNilF(t, err)\n\tbytes := resp.Bytes()\n\trespStr := string(bytes[:])\n\tif !strings.Contains(respStr, \"Your identity was confirmed and propagated to Snowflake Go.\\nYou can close this window now and go back where you started from.\") {\n\t\tt.Fatalf(\"failed to build response\")\n\t}\n}\n\nfunc postAuthExternalBrowserError(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {\n\treturn &authResponse{}, errors.New(\"failed to get SAML response\")\n}\n\nfunc postAuthExternalBrowserErrorDelayed(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {\n\ttime.Sleep(2 * time.Second)\n\treturn &authResponse{}, errors.New(\"failed to get SAML response\")\n}\n\nfunc postAuthExternalBrowserFail(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {\n\treturn &authResponse{\n\t\tSuccess: false,\n\t\tMessage: \"external browser auth failed\",\n\t}, nil\n}\n\nfunc postAuthExternalBrowserFailWithCode(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {\n\treturn &authResponse{\n\t\tSuccess: false,\n\t\tMessage: \"failed to connect to db\",\n\t\tCode: \"260008\",\n\t}, nil\n}\n\nfunc TestUnitAuthenticateByExternalBrowser(t *testing.T) {\n\tauthenticator := \"externalbrowser\"\n\tapplication := \"testapp\"\n\taccount := \"testaccount\"\n\tuser := \"u\"\n\ttimeout := sfconfig.DefaultExternalBrowserTimeout\n\tsr := &snowflakeRestful{\n\t\tProtocol: \"https\",\n\t\tHost: \"abc.com\",\n\t\tPort: 443,\n\t\tFuncPostAuthSAML: postAuthExternalBrowserError,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\t_, _, err := authenticateByExternalBrowser(context.Background(), sr, authenticator, application, account, user, timeout, ConfigBoolTrue)\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\tsr.FuncPostAuthSAML = postAuthExternalBrowserFail\n\t_, _, err = authenticateByExternalBrowser(context.Background(), sr, authenticator, application, account, user, timeout, ConfigBoolTrue)\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\tsr.FuncPostAuthSAML = postAuthExternalBrowserFailWithCode\n\t_, _, err = authenticateByExternalBrowser(context.Background(), sr, authenticator, application, account, user, timeout, ConfigBoolTrue)\n\tif err == nil {\n\t\tt.Fatal(\"should have failed.\")\n\t}\n\tdriverErr, ok := err.(*SnowflakeError)\n\tif !ok {\n\t\tt.Fatalf(\"should be snowflake error. err: %v\", err)\n\t}\n\tif driverErr.Number != ErrCodeFailedToConnect {\n\t\tt.Fatalf(\"unexpected error code. expected: %v, got: %v\", ErrCodeFailedToConnect, driverErr.Number)\n\t}\n}\n\nfunc TestAuthenticationTimeout(t *testing.T) {\n\tauthenticator := \"externalbrowser\"\n\tapplication := \"testapp\"\n\taccount := \"testaccount\"\n\tuser := \"u\"\n\ttimeout := 1 * time.Second\n\tsr := &snowflakeRestful{\n\t\tProtocol: \"https\",\n\t\tHost: \"abc.com\",\n\t\tPort: 443,\n\t\tFuncPostAuthSAML: postAuthExternalBrowserErrorDelayed,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\t_, _, err := authenticateByExternalBrowser(context.Background(), sr, authenticator, application, account, user, timeout, ConfigBoolTrue)\n\tassertEqualE(t, err.Error(), \"authentication timed out\", err.Error())\n}\n\nfunc Test_createLocalTCPListener(t *testing.T) {\n\tlistener, err := createLocalTCPListener(0)\n\tif err != nil {\n\t\tt.Fatalf(\"createLocalTCPListener() failed: %v\", err)\n\t}\n\tif listener == nil {\n\t\tt.Fatal(\"createLocalTCPListener() returned nil listener\")\n\t}\n\n\t// Close the listener after the test.\n\tdefer listener.Close()\n}\n\nfunc TestUnitGetLoginURL(t *testing.T) {\n\texpectedScheme := \"https\"\n\texpectedHost := \"abc.com:443\"\n\tuser := \"u\"\n\tcallbackPort := 123\n\tsr := &snowflakeRestful{\n\t\tProtocol: \"https\",\n\t\tHost: \"abc.com\",\n\t\tPort: 443,\n\t\tTokenAccessor: getSimpleTokenAccessor(),\n\t}\n\n\tloginURL, proofKey, err := getLoginURL(sr, user, callbackPort)\n\tassertNilF(t, err, \"failed to get login URL\")\n\tassertNotNilF(t, len(proofKey), \"proofKey should be non-empty string\")\n\n\turlPtr, err := url.Parse(loginURL)\n\tassertNilF(t, err, \"failed to parse the login URL\")\n\tassertEqualF(t, urlPtr.Scheme, expectedScheme)\n\tassertEqualF(t, urlPtr.Host, expectedHost)\n\tassertEqualF(t, urlPtr.Path, consoleLoginRequestPath)\n\tassertStringContainsF(t, urlPtr.RawQuery, \"login_name\")\n\tassertStringContainsF(t, urlPtr.RawQuery, \"browser_mode_redirect_port\")\n\tassertStringContainsF(t, urlPtr.RawQuery, \"proof_key\")\n}\n\ntype nonInteractiveSamlResponseProvider struct {\n\tt *testing.T\n}\n\nfunc (provider *nonInteractiveSamlResponseProvider) run(url string) error {\n\tgo func() {\n\t\tresp, err := http.Get(url)\n\t\tassertNilF(provider.t, err)\n\t\tassertEqualE(provider.t, resp.StatusCode, http.StatusOK)\n\t}()\n\treturn nil\n}\n"
},
{
"path": "authokta.go",
"content": "package gosnowflake\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"github.com/snowflakedb/gosnowflake/v2/internal/errors\"\n\t\"html\"\n\t\"io\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"strconv\"\n\t\"time\"\n)\n\ntype authOKTARequest struct {\n\tUsername string `json:\"username\"`\n\tPassword string `json:\"password\"`\n}\n\ntype authOKTAResponse struct {\n\tCookieToken string `json:\"cookieToken\"`\n\tSessionToken string `json:\"sessionToken\"`\n}\n\n/*\nauthenticateBySAML authenticates a user by SAML\nSAML Authentication\n 1. query GS to obtain IDP token and SSO url\n 2. IMPORTANT Client side validation:\n validate both token url and sso url contains same prefix\n (protocol + host + port) as the given authenticator url.\n Explanation:\n This provides a way for the user to 'authenticate' the IDP it is\n sending his/her credentials to. Without such a check, the user could\n be coerced to provide credentials to an IDP impersonator.\n 3. query IDP token url to authenticate and retrieve access token\n 4. given access token, query IDP URL snowflake app to get SAML response\n 5. IMPORTANT Client side validation:\n validate the post back url come back with the SAML response\n contains the same prefix as the Snowflake's server url, which is the\n intended destination url to Snowflake.\n\nExplanation:\n\n\tThis emulates the behavior of IDP initiated login flow in the user\n\tbrowser where the IDP instructs the browser to POST the SAML\n\tassertion to the specific SP endpoint. This is critical in\n\tpreventing a SAML assertion issued to one SP from being sent to\n\tanother SP.\n*/\nfunc authenticateBySAML(\n\tctx context.Context,\n\tsr *snowflakeRestful,\n\toktaURL *url.URL,\n\tapplication string,\n\taccount string,\n\tuser string,\n\tpassword string,\n\tdisableSamlURLCheck ConfigBool,\n) (samlResponse []byte, err error) {\n\tlogger.WithContext(ctx).Info(\"step 1: query GS to obtain IDP token and SSO url\")\n\theaders := make(map[string]string)\n\theaders[httpHeaderContentType] = headerContentTypeApplicationJSON\n\theaders[httpHeaderAccept] = headerContentTypeApplicationJSON\n\theaders[httpHeaderUserAgent] = userAgent\n\n\tclientEnvironment := newAuthRequestClientEnvironment()\n\tclientEnvironment.Application = application\n\trequestMain := authRequestData{\n\t\tClientAppID: clientType,\n\t\tClientAppVersion: SnowflakeGoDriverVersion,\n\t\tAccountName: account,\n\t\tClientEnvironment: clientEnvironment,\n\t\tAuthenticator: oktaURL.String(),\n\t}\n\tauthRequest := authRequest{\n\t\tData: requestMain,\n\t}\n\tparams := &url.Values{}\n\tjsonBody, err := json.Marshal(authRequest)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tlogger.WithContext(ctx).Infof(\"PARAMS for Auth: %v, %v\", params, sr)\n\trespd, err := sr.FuncPostAuthSAML(ctx, sr, headers, jsonBody, sr.LoginTimeout)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif !respd.Success {\n\t\tlogger.WithContext(ctx).Error(\"Authentication FAILED\")\n\t\tsr.TokenAccessor.SetTokens(\"\", \"\", -1)\n\t\tcode, err := strconv.Atoi(respd.Code)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn nil, &SnowflakeError{\n\t\t\tNumber: code,\n\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\tMessage: respd.Message,\n\t\t}\n\t}\n\tlogger.WithContext(ctx).Info(\"step 2: validate Token and SSO URL has the same prefix as oktaURL\")\n\tvar tokenURL *url.URL\n\tvar ssoURL *url.URL\n\tif tokenURL, err = url.Parse(respd.Data.TokenURL); err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to parse token URL. %v\", respd.Data.TokenURL)\n\t}\n\tif ssoURL, err = url.Parse(respd.Data.SSOURL); err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to parse SSO URL. %v\", respd.Data.SSOURL)\n\t}\n\tif !isPrefixEqual(oktaURL, ssoURL) || !isPrefixEqual(oktaURL, tokenURL) {\n\t\treturn nil, &SnowflakeError{\n\t\t\tNumber: ErrCodeIdpConnectionError,\n\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\tMessage: errors.ErrMsgIdpConnectionError,\n\t\t\tMessageArgs: []any{oktaURL, respd.Data.TokenURL, respd.Data.SSOURL},\n\t\t}\n\t}\n\tlogger.WithContext(ctx).Info(\"step 3: query IDP token url to authenticate and retrieve access token\")\n\tjsonBody, err = json.Marshal(authOKTARequest{\n\t\tUsername: user,\n\t\tPassword: password,\n\t})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\trespa, err := sr.FuncPostAuthOKTA(ctx, sr, headers, jsonBody, respd.Data.TokenURL, sr.LoginTimeout)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tlogger.WithContext(ctx).Info(\"step 4: query IDP URL snowflake app to get SAML response\")\n\tparams = &url.Values{}\n\tparams.Add(\"RelayState\", \"/some/deep/link\")\n\tvar oneTimeToken string\n\tif respa.SessionToken != \"\" {\n\t\toneTimeToken = respa.SessionToken\n\t} else {\n\t\toneTimeToken = respa.CookieToken\n\t}\n\tparams.Add(\"onetimetoken\", oneTimeToken)\n\n\theaders = make(map[string]string)\n\theaders[httpHeaderAccept] = \"*/*\"\n\tbd, err := sr.FuncGetSSO(ctx, sr, params, headers, respd.Data.SSOURL, sr.LoginTimeout)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif disableSamlURLCheck == ConfigBoolFalse {\n\t\tlogger.WithContext(ctx).Info(\"step 5: validate post_back_url matches Snowflake URL\")\n\t\ttgtURL, err := postBackURL(bd)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\n\t\tfullURL := sr.getURL()\n\t\tlogger.WithContext(ctx).Infof(\"tgtURL: %v, origURL: %v\", tgtURL, fullURL)\n\t\tif !isPrefixEqual(tgtURL, fullURL) {\n\t\t\treturn nil, &SnowflakeError{\n\t\t\t\tNumber: ErrCodeSSOURLNotMatch,\n\t\t\t\tSQLState: SQLStateConnectionRejected,\n\t\t\t\tMessage: errors.ErrMsgSSOURLNotMatch,\n\t\t\t\tMessageArgs: []any{tgtURL, fullURL},\n\t\t\t}\n\t\t}\n\t}\n\treturn bd, nil\n}\n\nfunc postBackURL(htmlData []byte) (url *url.URL, err error) {\n\tidx0 := bytes.Index(htmlData, []byte(\"