├── .github ├── CODEOWNERS ├── codeql-config.yml ├── dependabot.yml ├── ISSUE_TEMPLATE.md ├── workflows │ ├── add-to-project.yml │ ├── static_analysis.yml │ ├── codeql-analysis.yml │ ├── publish-docs.yml │ ├── windows-tests.yml │ ├── tests.yml │ ├── template-sync.yml │ └── release.yml └── PULL_REQUEST_TEMPLATE.md ├── docs ├── database.md ├── credentials.md ├── img │ ├── favicon.ico │ └── prefect-logo-mark.png ├── gen_home_page.py ├── stylesheets │ └── extra.css ├── gen_blocks_catalog.py └── gen_examples_catalog.py ├── .gitattributes ├── requirements.txt ├── prefect_snowflake ├── __init__.py ├── credentials.py ├── _version.py └── database.py ├── requirements-dev.txt ├── MANIFEST.in ├── .cruft.json ├── tests ├── test_block_standards.py ├── test_data │ ├── test_cert_no_pass.p8 │ ├── test_cert.p8 │ └── test_cert_malformed_format.p8 ├── conftest.py ├── test_database.py └── test_credentials.py ├── .pre-commit-config.yaml ├── setup.cfg ├── setup.py ├── .gitignore ├── mkdocs.yml ├── MAINTAINERS.md ├── README.md └── LICENSE /.github/CODEOWNERS: -------------------------------------------------------------------------------- 1 | * @PrefectHQ/open-source 2 | -------------------------------------------------------------------------------- /docs/database.md: -------------------------------------------------------------------------------- 1 | ::: prefect_snowflake.database 2 | -------------------------------------------------------------------------------- /docs/credentials.md: -------------------------------------------------------------------------------- 1 | ::: prefect_snowflake.credentials 2 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | prefect_snowflake/_version.py export-subst 2 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | prefect>=2.13.5 2 | snowflake-connector-python>=2.7.6 3 | -------------------------------------------------------------------------------- /docs/img/favicon.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PrefectHQ/prefect-snowflake/HEAD/docs/img/favicon.ico -------------------------------------------------------------------------------- /docs/img/prefect-logo-mark.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PrefectHQ/prefect-snowflake/HEAD/docs/img/prefect-logo-mark.png -------------------------------------------------------------------------------- /.github/codeql-config.yml: -------------------------------------------------------------------------------- 1 | paths-ignore: 2 | - tests/**/test_*.py 3 | - versioneer.py 4 | - prefect_snowflake/_version.py -------------------------------------------------------------------------------- /prefect_snowflake/__init__.py: -------------------------------------------------------------------------------- 1 | from . import _version 2 | from prefect_snowflake.credentials import SnowflakeCredentials # noqa 3 | from prefect_snowflake.database import SnowflakeConnector # noqa 4 | 5 | __version__ = _version.get_versions()["version"] 6 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | black 3 | flake8 4 | mypy 5 | mkdocs 6 | mkdocs-material 7 | mkdocstrings_python_legacy 8 | isort 9 | pre-commit 10 | pytest-asyncio 11 | mock; python_version < '3.8' 12 | mkdocs-gen-files 13 | interrogate 14 | coverage 15 | pillow 16 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | 4 | - package-ecosystem: "pip" 5 | directory: "/" 6 | schedule: 7 | interval: "daily" 8 | 9 | - package-ecosystem: "github-actions" 10 | directory: "/" 11 | schedule: 12 | interval: "daily" -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | # Expectation / Proposal 4 | 5 | # Traceback / Example 6 | 7 | - [ ] I would like to [help contribute](https://PrefectHQ.github.io/prefect-snowflake/#contributing) a pull request to resolve this! 8 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | # Things to always exclude 2 | global-exclude .git* 3 | global-exclude .ipynb_checkpoints 4 | global-exclude *.py[co] 5 | global-exclude __pycache__/** 6 | 7 | # Top-level Config 8 | include versioneer.py 9 | include prefect_snowflake/_version.py 10 | include LICENSE 11 | include MANIFEST.in 12 | include setup.cfg 13 | include requirements.txt 14 | include requirements-dev.txt 15 | -------------------------------------------------------------------------------- /docs/gen_home_page.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copies README.md to index.md. 3 | """ 4 | 5 | from pathlib import Path 6 | 7 | import mkdocs_gen_files 8 | 9 | # Home page 10 | 11 | readme_path = Path("README.md") 12 | docs_index_path = Path("index.md") 13 | 14 | with open(readme_path, "r") as readme: 15 | with mkdocs_gen_files.open(docs_index_path, "w") as generated_file: 16 | for line in readme: 17 | if line.startswith("Visit the full docs [here]("): 18 | continue # prevent linking to itself 19 | generated_file.write(line) 20 | 21 | mkdocs_gen_files.set_edit_path(Path(docs_index_path), readme_path) 22 | -------------------------------------------------------------------------------- /.github/workflows/add-to-project.yml: -------------------------------------------------------------------------------- 1 | name: Add issues to integrations board 2 | 3 | on: 4 | issues: 5 | types: 6 | - opened 7 | 8 | jobs: 9 | 10 | add-to-project: 11 | name: Add issue to project 12 | runs-on: ubuntu-latest 13 | steps: 14 | - uses: tibdex/github-app-token@v1 15 | id: generate-token 16 | name: Generate GitHub token 17 | with: 18 | app_id: ${{ secrets.SYNC_APP_ID }} 19 | private_key: ${{ secrets.SYNC_APP_PRIVATE_KEY }} 20 | 21 | - uses: actions/add-to-project@v0.4.0 22 | with: 23 | project-url: ${{ secrets.ADD_TO_PROJECT_URL }} 24 | github-token: ${{ steps.generate-token.outputs.token }} 25 | -------------------------------------------------------------------------------- /.github/workflows/static_analysis.yml: -------------------------------------------------------------------------------- 1 | name: Static analysis 2 | 3 | on: [pull_request] 4 | 5 | jobs: 6 | pre-commit-checks: 7 | name: Pre-commit checks 8 | runs-on: ubuntu-latest 9 | 10 | steps: 11 | - uses: actions/checkout@v4 12 | with: 13 | persist-credentials: false 14 | 15 | - name: Set up Python 16 | uses: actions/setup-python@v4 17 | with: 18 | python-version: 3.9 19 | 20 | - name: Install pre-commit 21 | run: | 22 | python -m pip install --upgrade pip 23 | pip install pre-commit 24 | 25 | - name: Run pre-commit 26 | run: | 27 | pre-commit run --show-diff-on-failure --color=always --all-files 28 | -------------------------------------------------------------------------------- /.cruft.json: -------------------------------------------------------------------------------- 1 | { 2 | "template": "https://github.com/PrefectHQ/prefect-collection-template", 3 | "commit": "f848dbf33334d22ade07b01ed17269bfa74d0f9c", 4 | "context": { 5 | "cookiecutter": { 6 | "full_name": "Prefect Technologies, Inc.", 7 | "email": "help@prefect.io", 8 | "github_organization": "PrefectHQ", 9 | "collection_name": "prefect-snowflake", 10 | "collection_slug": "prefect_snowflake", 11 | "collection_short_description": "Prefect integrations with Snowflake.", 12 | "_copy_without_render": [ 13 | ".github/workflows/*.yml" 14 | ], 15 | "_template": "https://github.com/PrefectHQ/prefect-collection-template" 16 | } 17 | }, 18 | "directory": null, 19 | "checkout": null 20 | } 21 | -------------------------------------------------------------------------------- /tests/test_block_standards.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from prefect.blocks.core import Block 3 | from prefect.testing.standard_test_suites import BlockStandardTestSuite 4 | from prefect.utilities.dispatch import get_registry_for_type 5 | from prefect.utilities.importtools import to_qualified_name 6 | 7 | 8 | def find_module_blocks(): 9 | blocks = get_registry_for_type(Block) 10 | module_blocks = [ 11 | block 12 | for block in blocks.values() 13 | if to_qualified_name(block).startswith("prefect_snowflake") 14 | ] 15 | return module_blocks 16 | 17 | 18 | @pytest.mark.parametrize("block", find_module_blocks()) 19 | class TestAllBlocksAdhereToStandards(BlockStandardTestSuite): 20 | @pytest.fixture 21 | def block(self, block): 22 | return block 23 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pycqa/isort 3 | rev: 5.12.0 4 | hooks: 5 | - id: isort 6 | language_version: python3 7 | - repo: https://github.com/psf/black 8 | rev: 22.3.0 9 | hooks: 10 | - id: black 11 | language_version: python3 12 | - repo: https://github.com/pycqa/flake8 13 | rev: 4.0.1 14 | hooks: 15 | - id: flake8 16 | - repo: https://github.com/econchick/interrogate 17 | rev: 1.5.0 18 | hooks: 19 | - id: interrogate 20 | args: [-vv] 21 | pass_filenames: false 22 | - repo: https://github.com/fsouza/autoflake8 23 | rev: v0.3.2 24 | hooks: 25 | - id: autoflake8 26 | language_version: python3 27 | args: [ 28 | '--in-place', 29 | ] 30 | -------------------------------------------------------------------------------- /.github/workflows/codeql-analysis.yml: -------------------------------------------------------------------------------- 1 | name: CodeQL 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | 8 | jobs: 9 | analyze: 10 | name: Analyze 11 | runs-on: ubuntu-latest 12 | permissions: 13 | actions: read 14 | contents: read 15 | security-events: write 16 | 17 | strategy: 18 | fail-fast: false 19 | matrix: 20 | language: 21 | - python 22 | 23 | steps: 24 | - name: Checkout repository 25 | uses: actions/checkout@v4 26 | 27 | - name: Initialize CodeQL 28 | uses: github/codeql-action/init@v2 29 | with: 30 | languages: ${{ matrix.language }} 31 | config-file: ./.github/codeql-config.yml 32 | queries: security-and-quality 33 | 34 | - name: Perform CodeQL Analysis 35 | uses: github/codeql-action/analyze@v2 36 | -------------------------------------------------------------------------------- /.github/workflows/publish-docs.yml: -------------------------------------------------------------------------------- 1 | name: Publish docs 2 | 3 | on: 4 | workflow_dispatch 5 | 6 | jobs: 7 | build-and-publish-docs: 8 | name: Build and publish docs 9 | runs-on: ubuntu-latest 10 | 11 | steps: 12 | - uses: actions/checkout@v4 13 | 14 | - name: Set up Python 3.10 15 | uses: actions/setup-python@v4 16 | with: 17 | python-version: "3.10" 18 | cache: pip 19 | cache-dependency-path: requirements*.txt 20 | 21 | - name: Install dependencies 22 | run: | 23 | python -m pip install --upgrade pip 24 | python -m pip install --upgrade --upgrade-strategy eager -e ".[dev]" 25 | mkdocs build 26 | 27 | - name: Publish docs 28 | uses: JamesIves/github-pages-deploy-action@v4.4.1 29 | with: 30 | branch: docs 31 | folder: site 32 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [flake8] 2 | exclude = .git,__pycache__,build,dist 3 | per-file-ignores = 4 | setup.py:E501 5 | # Match black line-length 6 | max-line-length = 88 7 | extend-ignore = 8 | E203, 9 | 10 | [isort] 11 | skip = __init__.py 12 | profile = black 13 | skip_gitignore = True 14 | multi_line_output = 3 15 | 16 | [versioneer] 17 | VCS = git 18 | style = pep440 19 | versionfile_source = prefect_snowflake/_version.py 20 | versionfile_build = prefect_snowflake/_version.py 21 | tag_prefix = v 22 | parentdir_prefix = 23 | 24 | [tool:interrogate] 25 | ignore-init-module = True 26 | ignore_init_method = True 27 | exclude = prefect_snowflake/_version.py, tests, setup.py, versioneer.py, docs, site 28 | fail-under = 95 29 | omit-covered-files = True 30 | 31 | [coverage:run] 32 | omit = tests/*, prefect_snowflake/_version.py 33 | 34 | [coverage:report] 35 | fail_under = 80 36 | show_missing = True 37 | 38 | [tool:pytest] 39 | asyncio_mode = auto 40 | -------------------------------------------------------------------------------- /.github/workflows/windows-tests.yml: -------------------------------------------------------------------------------- 1 | 2 | name: Windows Tests 3 | 4 | on: [pull_request] 5 | 6 | jobs: 7 | run-tests: 8 | name: Run Tests 9 | runs-on: windows-latest 10 | strategy: 11 | matrix: 12 | # Prefect only tests 3.9 on Windows 13 | python-version: 14 | - "3.9" 15 | fail-fast: false 16 | steps: 17 | - uses: actions/checkout@v4 18 | 19 | - name: Set up Python ${{ matrix.python-version }} 20 | uses: actions/setup-python@v4 21 | with: 22 | python-version: ${{ matrix.python-version }} 23 | cache: pip 24 | cache-dependency-path: requirements*.txt 25 | 26 | - name: Install dependencies 27 | run: | 28 | python -m pip install --upgrade pip 29 | python -m pip install --upgrade --upgrade-strategy eager -e ".[dev]" 30 | - name: Run tests 31 | env: 32 | PREFECT_ORION_DATABASE_CONNECTION_URL: "sqlite+aiosqlite:///./orion-tests.db" 33 | run: | 34 | pytest tests -vv 35 | -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | Closes 6 | 7 | ### Example 8 | 9 | 10 | ### Screenshots 11 | 17 | 18 | ### Checklist 19 | 20 | 21 | - [ ] References any related issue by including "Closes #" or "Closes ". 22 | - If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect-snowflake/issues/new/choose) first. 23 | - [ ] Includes tests or only affects documentation. 24 | - [ ] Passes `pre-commit` checks. 25 | - Run `pre-commit install && pre-commit run --all` locally for formatting and linting. 26 | - [ ] Includes screenshots of documentation updates. 27 | - Run `mkdocs serve` view documentation locally. 28 | -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: Tests 2 | 3 | on: [pull_request] 4 | 5 | jobs: 6 | run-tests: 7 | name: Run Tests 8 | runs-on: ubuntu-latest 9 | strategy: 10 | matrix: 11 | python-version: 12 | - "3.8" 13 | - "3.9" 14 | - "3.10" 15 | fail-fast: false 16 | steps: 17 | - uses: actions/checkout@v4 18 | 19 | - name: Set up Python ${{ matrix.python-version }} 20 | uses: actions/setup-python@v4 21 | with: 22 | python-version: ${{ matrix.python-version }} 23 | cache: pip 24 | cache-dependency-path: requirements*.txt 25 | 26 | - name: Install dependencies 27 | run: | 28 | python -m pip install --upgrade pip 29 | python -m pip install --upgrade --upgrade-strategy eager -e ".[dev]" 30 | 31 | - name: Run tests 32 | env: 33 | PREFECT_ORION_DATABASE_CONNECTION_URL: "sqlite+aiosqlite:///./orion-tests.db" 34 | run: | 35 | coverage run --branch -m pytest tests -vv 36 | coverage report 37 | 38 | - name: Run mkdocs build 39 | run: | 40 | mkdocs build --verbose --clean 41 | -------------------------------------------------------------------------------- /.github/workflows/template-sync.yml: -------------------------------------------------------------------------------- 1 | name: Template Synchronization 2 | on: 3 | schedule: 4 | - cron: "0 0 * * *" 5 | workflow_dispatch: 6 | 7 | jobs: 8 | submit-update-pr: 9 | name: Submit update PR 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v4 13 | 14 | - name: Set up Python 15 | uses: actions/setup-python@v4 16 | with: 17 | python-version: 3.9 18 | 19 | - name: Install cruft 20 | run: pip install "cookiecutter>=1.7.3,<2.0.0" cruft 21 | 22 | - name: Perform updates 23 | run: cruft update -y 24 | 25 | - uses: tibdex/github-app-token@v1 26 | id: generate-token 27 | name: Generate GitHub token 28 | with: 29 | app_id: ${{ secrets.SYNC_APP_ID }} 30 | private_key: ${{ secrets.SYNC_APP_PRIVATE_KEY }} 31 | 32 | - name: Submit PR 33 | uses: peter-evans/create-pull-request@v4 34 | with: 35 | commit-message: Updating collection with changes to prefect-collection-template 36 | token: ${{ steps.generate-token.outputs.token }} 37 | branch: sync-with-template 38 | delete-branch: true 39 | title: Sync Collection with changes to prefect-collection-template 40 | body: | 41 | Automated PR created to propagate changes from prefect-collection-template to this collection 42 | 43 | Feel free to make any necessary changes to this PR before merging. 44 | labels: | 45 | template sync 46 | automated pr 47 | -------------------------------------------------------------------------------- /tests/test_data/test_cert_no_pass.p8: -------------------------------------------------------------------------------- 1 | -----BEGIN PRIVATE KEY----- 2 | MIIEvwIBADANBgkqhkiG9w0BAQEFAASCBKkwggSlAgEAAoIBAQDx+W7+mNsZ39tH 3 | 620SNB8HH498QM/4uQFqdqVS+jKKfGQeOOzsyDd4R6RI/rcvqZswAYVpzRzKs261 4 | cfWS2rSqObasPlhi/r4OuNgoz/g1DPj+s/KpB/Pbi5me4g4/b0X9e1UkDoPJUrA4 5 | H9ZMpxbFAJF7kTB27wfzd/3UIDf+u2iJymUW4mV96aNGo8NZlK7ucORvtlkV0K3M 6 | 0ocqegpXDMxB0FCWQ5EPZXkREVFNkGYhl08MyEOliI+Wv4iV1v17ONACKQErF/m1 7 | ZPieucTA3/FlEB1mZWwsS+jwjTbvhHqb/625z0/wbUJvWAqJpVFpRHwZ+4oiemUc 8 | H3HQ9Q4nAgMBAAECggEBAOMTj2QBkmu/tgrSFnNk0lFe/so+EG17QKbmXNN31II9 9 | pb7u8PhsFBIeOT//OW3THr14T/kv8XsP1C8WmLyN3cl5i8IsZ8nPAhDSypfSQz2V 10 | xP1RZcHWorQ/V5TQ00oYZHgyi33g2S/PF2w0BGm+zBFoLR7m67JRMwb6bgvpiGLV 11 | LWiRI7nN+PbCGgeHGbJMTndlQyAOepKJHjxygsgQFtWmOGge/hOonK5lRni9VcNH 12 | FlonrAlGxJ06PAo+wlFBak/1lOEH9Pc80uYXhlR2sKmymnQHxC08DrUyahiDzWdb 13 | MaVv8ClMstb00GKMYawUnhenFV4OACPrD2vOW6fgFmECgYEA+t8pGkk8MNFKnXoL 14 | pP9LoiR+JNaBWAp60R3flnUL4UW5oNPVeZFBW7tcwOU5yLGurlR9zpLFn3UyHABa 15 | Fw283oQHRU6z2HvqZNDKsZAB+OQkoPt+jNFhtTazEFdi80z3srQ5JMZX7/tLcufT 16 | QE7SoSW+AWzILlGF1oyv5r7Ywb8CgYEA9uu2S3cjRr6j5dU/WD8ryN3yI3eZ12Py 17 | gDH41EbUeFjbTHhHwM5RvCeYgyqV7lqLSkGsbYRL7RtfOqJ2zpRQ2dJBmNvyV0hE 18 | H/2PHLEkNy+GtU7T+M5bDUI7UYOanAhYPxL70PK/LUKY1cAxELCpRXIIJORH7+Rj 19 | ltkrqfrxfZkCgYBG6KTZhTG6Kq4IbOK1tGNQZTgyeV7935hvWx5DjLna5rZdOwLv 20 | 5Zqvrvm8nA4FKBPTupYEuX+aXqnXOFI+ieeEjZTwhhpXak8KR+nC3o1wKCwiRHO0 21 | ocoYSmm6iLizRGIO5NnyBw38Cu98fwI6/wyR9/UIuhKq5OgLiKB/fBoQSQKBgQCu 22 | UHUfDXpP5SHbjRHtAJuucFESjMqEbkCyE5UwdthkJYabk4ZELOMyy1k1sAisfis7 23 | PuW1YuIe/2XZBPyOldIGFBLPEOciixfc5an0fXGtq9WThZTLXxKUZ59sylWJtbJm 24 | xMYVGmmwUAWJUObSe4TkS75IHQhT6I4N3j6e/1MgcQKBgQDY7jGYzKdPqoHpPHie 25 | sthjH8YdPFwsTRDIh1qwRNQr4JBTEiX00SovtCu15H5/2BjY4DTRl0XOjNftB7ZR 26 | dv+Oze06D+7XoLv3uK7S/rqbtza5ejv7QJ9rrV2VLnJsjMafiZZfGjgN9K7PSiOE 27 | jtUynGlKVzyVvSmLdp3sbgYKCA== 28 | -----END PRIVATE KEY----- -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import find_packages, setup 2 | 3 | import versioneer 4 | 5 | with open("requirements.txt") as install_requires_file: 6 | install_requires = install_requires_file.read().strip().split("\n") 7 | 8 | with open("requirements-dev.txt") as dev_requires_file: 9 | dev_requires = dev_requires_file.read().strip().split("\n") 10 | 11 | with open("README.md") as readme_file: 12 | readme = readme_file.read() 13 | 14 | setup( 15 | name="prefect-snowflake", 16 | description="Prefect integrations for interacting with Snowflake.", 17 | license="Apache License 2.0", 18 | author="Prefect Technologies, Inc.", 19 | author_email="help@prefect.io", 20 | keywords="prefect", 21 | url="https://github.com/PrefectHQ/prefect-snowflake", 22 | long_description=readme, 23 | long_description_content_type="text/markdown", 24 | version=versioneer.get_version(), 25 | cmdclass=versioneer.get_cmdclass(), 26 | packages=find_packages(exclude=("tests", "docs")), 27 | python_requires=">=3.7", 28 | install_requires=install_requires, 29 | extras_require={"dev": dev_requires}, 30 | entry_points={ 31 | "prefect.collections": [ 32 | "prefect_snowflake = prefect_snowflake", 33 | ] 34 | }, 35 | classifiers=[ 36 | "Natural Language :: English", 37 | "Intended Audience :: Developers", 38 | "Intended Audience :: System Administrators", 39 | "License :: OSI Approved :: Apache Software License", 40 | "Programming Language :: Python :: 3 :: Only", 41 | "Programming Language :: Python :: 3.7", 42 | "Programming Language :: Python :: 3.8", 43 | "Programming Language :: Python :: 3.9", 44 | "Programming Language :: Python :: 3.10", 45 | "Topic :: Software Development :: Libraries", 46 | ], 47 | ) 48 | -------------------------------------------------------------------------------- /tests/test_data/test_cert.p8: -------------------------------------------------------------------------------- 1 | -----BEGIN ENCRYPTED PRIVATE KEY----- 2 | MIIFHDBOBgkqhkiG9w0BBQ0wQTApBgkqhkiG9w0BBQwwHAQIwbace/NmI2kCAggA 3 | MAwGCCqGSIb3DQIJBQAwFAYIKoZIhvcNAwcECFBJKD71p3udBIIEyE43T+94NaN8 4 | OOP0TWPYAfKKNimgVgKDzAn9pOuf0sppreHdYp4uTDFgK5vDq2SLUV7xIBQzF6OB 5 | TvXiOitBaFW3ENM7vaIlRgpUcRNol1VS+5NaLSGG/ZLJ8WVsPwsOby78piNx6kcp 6 | TBU0W0kdL1J3ovLXMGvq4iYp2OBXN56sRM5NF277G4bF1Hjk6Tc23BErHraiFoSE 7 | ZygoW1pc2cZ7WqGnB7v2HjVsRoAioJmy4nJnvZ/SwoD0+2EXMKSwZv3hC+76mgmE 8 | Ir/k0dH1cAe0eNaJcvK7QpCoPSWnoKzBFQYyH9a6bleMxJAj7nSSo8/SiAM3vq3z 9 | I4BhQt4UGy/nQ2LqD0QPCGLyacMPb198fD4ZvTe8N19xD6BEMHNiFeJ8JLAVQUFY 10 | XgxAcrh9wPXgwLEBcXQbZJpVLD38KLfTA+E9rJcnvhclthK1nf62ei6HcN2LgFRs 11 | m5nZTA1nxxBJ9rR4Paao3jEtSxaITSSAJRMeBMq+dfwQLgcmG8K6agx6z0FjcX97 12 | JBe1NFug5xrmuAWS39EjiZIxRHijFx9BzTmtV3W8jeBkmUeq3yFNcRBUIoPbTqPW 13 | 3MxT54AdrPmTJZ7159vzSoZIgct/gZYcw35raw99pD2RzH/KK71TfL/b4eK9Noa2 14 | RP18SHwUIILPcJ6k+qH6IGeXwQ3KyvHcBXL6Y4j1bBNIkr4UsGD0VAVl8ucjF5Pl 15 | hhNYH1KR9IAUml0LNH4SIYy9W0yS86yzrcD3Dq8szDs8WVB/1GqTTBnIYqfizhFC 16 | Mem4mGLEFQxOH8S98U1p/ruxiNAz8k/UrKVQEyTQ0WE9GSbOGZpdPDIScfREYcq4 17 | lkx7z5R6i3nzMldV7GTzcRPm44xn05/ALkGSu5G3viAezQaLL99+zLY13h7/R67J 18 | KG1YKERBsVVh6zjv58OOJonxmYaUumFmKDqe7JuIyZBNknvr3KvV+8iY0fxyaMVC 19 | X0tYWyFpdpDhtR59XRfHzwwc6S4f0H/uaXEMQQ2UxE15w0aYiVNHbI5fRiZHWRJN 20 | ppzpKCRfQPVbulI0oQraPawClQtacbYbQc7KQ23kM2WU443zmukkcE2EIbjUH67Z 21 | ESGhvIKpXZkRwRR5Wyyu1lKa5jNxJvAG6ArHskJ1mBGM2UmOGgRiDlcxMzNermpY 22 | uEwtfKNeTzqszVCKAVG8kMaLPXXCTApONnitPIL41qlC74G2YTNV4WvtZKRiC9La 23 | YsjGcwUJbX/JJBTD28aVZAgL9Lq0aQXjZPVDEhxuRLPW1ery8OuVdsUnOg9kdqpB 24 | Gn23wr4jYUX28KRW6uPXNhKsm+mOUtvkHz7+mk0SeUJeJFv+dnuRlKvZ4O30k+BJ 25 | GtXlpIeHZ6tYzcU5h1ZGoyThPXwX0qU57zhs8PLxbiQ/oUbBwFAgKUghLzYG5lXs 26 | TeoWta5H7eANd6XkIfFLV2EEcGrz1GuZWOv8iYvZDcxGLUoESmaV4e4Z56ly58Yz 27 | Qu1CJmiBESm9+qdF0wb+fNgLigHeVxHo5ebS4njNggmcPQbym0Iw1kUNnxRjRFeq 28 | 035QwksqQIuwLhQ9SvJQY2MOUyFTvADGS+VazushZ2SxW0KTtieCbVPKPM0+N4HQ 29 | sshzYu9lhehMrOo/FvFMtQ== 30 | -----END ENCRYPTED PRIVATE KEY----- -------------------------------------------------------------------------------- /tests/test_data/test_cert_malformed_format.p8: -------------------------------------------------------------------------------- 1 | -----BEGIN ENCRYPTED PRIVATE KEY----- 2 | MIIFHDBOBgkqhkiG9w0BBQ0wQTApBgkqhkiG9w0BBQwwHAQIwbace/NmI2kCAggA 3 | MAwGCCqGSIb3DQIJBQAwFAYIKoZIhvcNAwcECFBJKD71p3udBIIEyE43T+94NaN8 4 | OOP0TWPYAfKKNimgVgKDzAn9pOuf0sppreHdYp4uTDFgK5vDq2SLUV7xIBQzF6OB 5 | TvXiOitBaFW3ENM7vaIlRgpUcRNol1VS+5NaLSGG/ZLJ8WVsPwsOby78piNx6kcp 6 | TBU0W0kdL1J3ovLXMGvq4iYp2OBXN56sRM5NF277G4bF1Hjk6Tc23BErHraiFoSE 7 | ZygoW1pc2cZ7WqGnB7v2HjVsRoAioJmy4nJnvZ/SwoD0+2EXMKSwZv3hC+76mgmE 8 | Ir/k0dH1cAe0eNaJcvK7QpCoPSWnoKzBFQYyH9a6bleMxJAj7nSSo8/SiAM3vq3z 9 | I4BhQt4UGy/nQ2LqD0QPCGLyacMPb198fD4ZvTe8N19xD6BEMHNiFeJ8JLAVQUFY 10 | XgxAcrh9wPXgwLEBcXQbZJpVLD38KLfTA+E9rJcnvhclthK1nf62ei6HcN2LgFRs 11 | m5nZTA1nxxBJ9rR4Paao3jEtSxaITSSAJRMeBMq+dfwQLgcmG8K6agx6z0FjcX97 12 | JBe1NFug5xrmuAWS39EjiZIxRHijFx9BzTmtV3W8jeBkmUeq3yFNcRBUIoPbTqPW 13 | 3MxT54AdrPmTJZ7159vzSoZIgct/gZYcw35raw99pD2RzH/KK71TfL/b4eK9Noa2 14 | RP18SHwUIILPcJ6k+qH6IGeXwQ3KyvHcBXL6Y4j1bBNIkr4UsGD0VAVl8ucjF5Pl 15 | hhNYH1KR9IAUml0LNH4SIYy9W0yS86yzrcD3Dq8szDs8WVB/1GqTTBnIYqfizhFC 16 | Mem4mGLEFQxOH8S98U1p/ruxiNAz8k/UrKVQEyTQ0WE9GSbOGZpdPDIScfREYcq4 17 | lkx7z5R6i3nzMldV7GTzcRPm44xn05/ALkGSu5G3viAezQaLL99+zLY13h7/R67J 18 | KG1YKERBsVVh6zjv58OOJonxmYaUumFmKDqe7JuIyZBNknvr3KvV+8iY0fxyaMVC 19 | X0tYWyFpdpDhtR59XRfHzwwc6S4f0H/uaXEMQQ2UxE15w0aYiVNHbI5fRiZHWRJN 20 | ppzpKCRfQPVbulI0oQraPawClQtacbYbQc7KQ23kM2WU443zmukkcE2EIbjUH67Z 21 | ESGhvIKpXZkRwRR5Wyyu1lKa5jNxJvAG6ArHskJ1mBGM2UmOGgRiDlcxMzNermpY 22 | uEwtfKNeTzqszVCKAVG8kMaLPXXCTApONnitPIL41qlC74G2YTNV4WvtZKRiC9La 23 | YsjGcwUJbX/JJBTD28aVZAgL9Lq0aQXjZPVDEhxuRLPW1ery8OuVdsUnOg9kdqpB 24 | Gn23wr4jYUX28KRW6uPXNhKsm+mOUtvkHz7+mk0SeUJeJFv+dnuRlKvZ4O30k+BJ 25 | GtXlpIeHZ6tYzcU5h1ZGoyThPXwX0qU57zhs8PLxbiQ/oUbBwFAgKUghLzYG5lXs 26 | TeoWta5H7eANd6XkIfFLV2EEcGrz1GuZWOv8iYvZDcxGLUoESmaV4e4Z56ly58Yz Qu1CJmiBESm9+qdF0wb+fNgLigHeVxHo5ebS4njNggmcPQbym0Iw1kUNnxRjRFeq 27 | 035QwksqQIuwLhQ9SvJQY2MOUyFTvADGS+VazushZ2SxW0KTtieCbVPKPM0+N4HQ sshzYu9lhehMrOo/FvFMtQ== -----END ENCRYPTED PRIVATE KEY----- -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Build & Release 2 | 3 | on: 4 | push: 5 | tags: 6 | - "v*" 7 | 8 | jobs: 9 | build-release: 10 | name: Build Release 11 | runs-on: ubuntu-latest 12 | steps: 13 | - uses: actions/checkout@v4 14 | 15 | - name: Set up Python 16 | uses: actions/setup-python@v4 17 | with: 18 | python-version: 3.8 19 | 20 | - name: Install packages 21 | run: | 22 | python -m pip install --upgrade pip build 23 | python -m pip install --upgrade --upgrade-strategy eager -e .[dev] 24 | 25 | - name: Build a binary wheel and a source tarball 26 | run: | 27 | python -m build --sdist --wheel --outdir dist/ 28 | 29 | - name: Publish build artifacts 30 | uses: actions/upload-artifact@v3 31 | with: 32 | name: built-package 33 | path: "./dist" 34 | 35 | publish-release: 36 | name: Publish release to PyPI 37 | needs: [build-release] 38 | environment: "prod" 39 | runs-on: ubuntu-latest 40 | 41 | steps: 42 | - name: Download build artifacts 43 | uses: actions/download-artifact@v3 44 | with: 45 | name: built-package 46 | path: "./dist" 47 | 48 | - name: Publish distribution to PyPI 49 | uses: pypa/gh-action-pypi-publish@release/v1 50 | with: 51 | password: ${{ secrets.PYPI_API_TOKEN }} 52 | verbose: true 53 | 54 | build-and-publish-docs: 55 | name: Build and publish docs 56 | needs: [build-release, publish-release] 57 | runs-on: ubuntu-latest 58 | 59 | steps: 60 | - uses: actions/checkout@v4 61 | 62 | - name: Set up Python 3.10 63 | uses: actions/setup-python@v4 64 | with: 65 | python-version: "3.10" 66 | cache: pip 67 | cache-dependency-path: requirements*.txt 68 | 69 | - name: Build docs 70 | run: | 71 | python -m pip install --upgrade pip 72 | python -m pip install --upgrade --upgrade-strategy eager -e .[dev] 73 | mkdocs build 74 | 75 | - name: Publish docs 76 | uses: JamesIves/github-pages-deploy-action@v4.4.1 77 | with: 78 | branch: docs 79 | folder: site 80 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # OS files 132 | .DS_Store 133 | 134 | # VS Code 135 | .vscode 136 | 137 | # Jupyter notebook 138 | *.ipynb 139 | 140 | # Pycharm 141 | .idea -------------------------------------------------------------------------------- /docs/stylesheets/extra.css: -------------------------------------------------------------------------------- 1 | /* theme */ 2 | :root > * { 3 | /* theme */ 4 | --md-primary-fg-color: #26272B; 5 | --md-primary-fg-color--light: #26272B; 6 | --md-primary-fg-color--dark: #26272B; 7 | } 8 | 9 | /* Table formatting */ 10 | .md-typeset table:not([class]) td { 11 | padding: 0.5em 1.25em; 12 | } 13 | .md-typeset table:not([class]) th { 14 | padding: 0.5em 1.25em; 15 | } 16 | 17 | /* convenience class to keep lines from breaking 18 | useful for wrapping table cell text in a span 19 | to force column width */ 20 | .no-wrap { 21 | white-space: nowrap; 22 | } 23 | 24 | /* badge formatting */ 25 | .badge::before { 26 | background-color: #1860F2; 27 | color: white; 28 | font-size: 0.8rem; 29 | font-weight: normal; 30 | padding: 4px 8px; 31 | margin-left: 0.5rem; 32 | vertical-align: super; 33 | text-align: center; 34 | border-radius: 5px; 35 | } 36 | 37 | .badge-api::before { 38 | background-color: #1860F2; 39 | color: white; 40 | font-size: 0.8rem; 41 | font-weight: normal; 42 | padding: 4px 8px; 43 | text-align: center; 44 | border-radius: 5px; 45 | } 46 | 47 | .experimental::before { 48 | background-color: #FCD14E; 49 | content: "Experimental"; 50 | } 51 | 52 | .cloud::before { 53 | background-color: #799AF7; 54 | content: "Prefect Cloud"; 55 | } 56 | 57 | .deprecated::before { 58 | background-color: #FA1C2F; 59 | content: "Deprecated"; 60 | } 61 | 62 | .new::before { 63 | background-color: #2AC769; 64 | content: "New"; 65 | } 66 | 67 | .expert::before { 68 | background-color: #726576; 69 | content: "Advanced"; 70 | } 71 | 72 | /* dark mode slate theme */ 73 | /* dark mode code overrides */ 74 | [data-md-color-scheme="slate"] { 75 | --md-code-bg-color: #1c1d20; 76 | --md-code-fg-color: #eee; 77 | --md-code-hl-color: #3b3d54; 78 | --md-code-hl-name-color: #eee; 79 | } 80 | 81 | /* dark mode link overrides */ 82 | [data-md-color-scheme="slate"] .md-typeset a { 83 | color: var(--blue); 84 | } 85 | 86 | [data-md-color-scheme="slate"] .md-typeset a:hover { 87 | font-weight: bold; 88 | } 89 | 90 | /* dark mode nav overrides */ 91 | [data-md-color-scheme="slate"] .md-nav--primary .md-nav__item--active>.md-nav__link { 92 | color: var(--blue); 93 | font-weight: bold; 94 | } 95 | 96 | [data-md-color-scheme="slate"] .md-nav--primary .md-nav__link--active { 97 | color: var(--blue); 98 | font-weight: bold; 99 | } 100 | 101 | /* dark mode collection catalog overrides */ 102 | [data-md-color-scheme="slate"] .collection-item { 103 | background-color: #26272B; 104 | } 105 | 106 | /* dark mode recipe collection overrides */ 107 | [data-md-color-scheme="slate"] .recipe-item { 108 | background-color: #26272B; 109 | } 110 | 111 | /* dark mode API doc overrides */ 112 | [data-md-color-scheme="slate"] .prefect-table th { 113 | background-color: #26272B; 114 | } -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: prefect-snowflake 2 | site_url: https://PrefectHQ.github.io/prefect-snowflake 3 | repo_url: https://github.com/prefecthq/prefect-snowflake 4 | edit_uri: edit/main/docs/ 5 | theme: 6 | name: material 7 | favicon: img/favicon.ico 8 | palette: 9 | - media: "(prefers-color-scheme)" 10 | toggle: 11 | icon: material/brightness-auto 12 | name: Switch to light mode 13 | - media: "(prefers-color-scheme: light)" 14 | accent: blue 15 | primary: blue 16 | scheme: default 17 | toggle: 18 | icon: material/weather-sunny 19 | name: Switch to dark mode 20 | - media: "(prefers-color-scheme: dark)" 21 | accent: blue 22 | primary: blue 23 | scheme: slate 24 | toggle: 25 | icon: material/weather-night 26 | name: Switch to light mode 27 | icon: 28 | repo: fontawesome/brands/github 29 | logo: 30 | img/prefect-logo-mark.png 31 | font: 32 | text: Inter 33 | code: Source Code Pro 34 | features: 35 | - content.code.copy 36 | - content.code.annotate 37 | extra_css: 38 | - stylesheets/extra.css 39 | markdown_extensions: 40 | - admonition 41 | - attr_list 42 | - codehilite 43 | - md_in_html 44 | - meta 45 | - pymdownx.highlight: 46 | use_pygments: true 47 | - pymdownx.superfences 48 | - pymdownx.tabbed: 49 | alternate_style: true 50 | - pymdownx.inlinehilite 51 | - pymdownx.snippets 52 | 53 | plugins: 54 | - search 55 | - gen-files: 56 | scripts: 57 | - docs/gen_home_page.py 58 | - docs/gen_examples_catalog.py 59 | - docs/gen_blocks_catalog.py 60 | - mkdocstrings: 61 | handlers: 62 | python: 63 | options: 64 | show_root_heading: True 65 | show_object_full_path: False 66 | show_category_heading: True 67 | show_bases: True 68 | show_signature: False 69 | heading_level: 1 70 | watch: 71 | - prefect_snowflake/ 72 | - README.md 73 | 74 | nav: 75 | - Home: index.md 76 | - Blocks Catalog: blocks_catalog.md 77 | - Examples Catalog: examples_catalog.md 78 | - API Reference: 79 | - Credentials: credentials.md 80 | - Database: database.md 81 | 82 | extra: 83 | social: 84 | - icon: fontawesome/brands/slack 85 | link: https://www.prefect.io/slack/ 86 | - icon: fontawesome/brands/discourse 87 | link: https://discourse.prefect.io/ 88 | - icon: fontawesome/brands/youtube 89 | link: https://www.youtube.com/c/PrefectIO/videos 90 | - icon: fontawesome/regular/newspaper 91 | link: https://prefect.io/guide/ 92 | - icon: fontawesome/brands/twitter 93 | link: https://twitter.com/PrefectIO 94 | - icon: fontawesome/brands/linkedin 95 | link: https://www.linkedin.com/company/prefect/ 96 | - icon: fontawesome/brands/github 97 | link: https://github.com/PrefectHQ/prefect 98 | - icon: fontawesome/brands/docker 99 | link: https://hub.docker.com/r/prefecthq/prefect/ 100 | - icon: fontawesome/brands/python 101 | link: https://pypi.org/project/prefect/ 102 | analytics: 103 | provider: google 104 | property: G-8CSMBCQDKN 105 | -------------------------------------------------------------------------------- /docs/gen_blocks_catalog.py: -------------------------------------------------------------------------------- 1 | """ 2 | Discovers all blocks and generates a list of them in the docs 3 | under the Blocks Catalog heading. 4 | """ 5 | 6 | from pathlib import Path 7 | from textwrap import dedent 8 | 9 | import mkdocs_gen_files 10 | from prefect.blocks.core import Block 11 | from prefect.utilities.dispatch import get_registry_for_type 12 | from prefect.utilities.importtools import from_qualified_name, to_qualified_name 13 | 14 | COLLECTION_SLUG = "prefect_snowflake" 15 | 16 | 17 | def find_module_blocks(): 18 | blocks = get_registry_for_type(Block) 19 | collection_blocks = [ 20 | block 21 | for block in blocks.values() 22 | if to_qualified_name(block).startswith(COLLECTION_SLUG) 23 | ] 24 | module_blocks = {} 25 | for block in collection_blocks: 26 | block_name = block.__name__ 27 | module_nesting = tuple(to_qualified_name(block).split(".")[1:-1]) 28 | if module_nesting not in module_blocks: 29 | module_blocks[module_nesting] = [] 30 | module_blocks[module_nesting].append(block_name) 31 | return module_blocks 32 | 33 | 34 | def insert_blocks_catalog(generated_file): 35 | module_blocks = find_module_blocks() 36 | if len(module_blocks) == 0: 37 | return 38 | generated_file.write( 39 | dedent( 40 | f""" 41 | Below is a list of Blocks available for registration in 42 | `prefect-snowflake`. 43 | 44 | To register blocks in this module to 45 | [view and edit them](https://orion-docs.prefect.io/ui/blocks/) 46 | on Prefect Cloud, first [install the required packages]( 47 | https://PrefectHQ.github.io/prefect-snowflake/#installation), 48 | then 49 | ```bash 50 | prefect block register -m {COLLECTION_SLUG} 51 | ``` 52 | """ # noqa 53 | ) 54 | ) 55 | generated_file.write( 56 | "Note, to use the `load` method on Blocks, you must already have a block document " # noqa 57 | "[saved through code](https://orion-docs.prefect.io/concepts/blocks/#saving-blocks) " # noqa 58 | "or [saved through the UI](https://orion-docs.prefect.io/ui/blocks/).\n" 59 | ) 60 | for module_nesting, block_names in module_blocks.items(): 61 | module_path = f"{COLLECTION_SLUG}." + " ".join(module_nesting) 62 | module_title = ( 63 | module_path.replace(COLLECTION_SLUG, "") 64 | .lstrip(".") 65 | .replace("_", " ") 66 | .title() 67 | ) 68 | generated_file.write(f"## [{module_title} Module][{module_path}]\n") 69 | for block_name in block_names: 70 | block_obj = from_qualified_name(f"{module_path}.{block_name}") 71 | block_description = block_obj.get_description() 72 | if not block_description.endswith("."): 73 | block_description += "." 74 | generated_file.write( 75 | f"[{block_name}][{module_path}.{block_name}]\n\n{block_description}\n\n" 76 | ) 77 | generated_file.write( 78 | dedent( 79 | f""" 80 | To load the {block_name}: 81 | ```python 82 | from prefect import flow 83 | from {module_path} import {block_name} 84 | 85 | @flow 86 | def my_flow(): 87 | my_block = {block_name}.load("MY_BLOCK_NAME") 88 | 89 | my_flow() 90 | ``` 91 | """ 92 | ) 93 | ) 94 | generated_file.write( 95 | f"For additional examples, check out the [{module_title} Module]" 96 | f"(../examples_catalog/#{module_nesting[-1]}-module) " 97 | f"under Examples Catalog.\n" 98 | ) 99 | 100 | 101 | blocks_catalog_path = Path("blocks_catalog.md") 102 | with mkdocs_gen_files.open(blocks_catalog_path, "w") as generated_file: 103 | insert_blocks_catalog(generated_file) 104 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import os 2 | from unittest.mock import MagicMock 3 | 4 | import pytest 5 | from prefect.testing.utilities import prefect_test_harness 6 | 7 | from prefect_snowflake.credentials import SnowflakeCredentials 8 | 9 | 10 | def _read_test_file(name: str) -> bytes: 11 | """ 12 | Args: 13 | name: File to load from test_data folder. 14 | 15 | Returns: 16 | File content as binary. 17 | """ 18 | full_name = os.path.join(os.path.split(__file__)[0], "test_data", name) 19 | with open(full_name, "rb") as fd: 20 | return fd.read() 21 | 22 | 23 | @pytest.fixture(scope="session", autouse=True) 24 | def prefect_db(): 25 | """ 26 | Sets up test harness for temporary DB during test runs. 27 | """ 28 | with prefect_test_harness(): 29 | yield 30 | 31 | 32 | @pytest.fixture(autouse=True) 33 | def reset_object_registry(): 34 | """ 35 | Ensures each test has a clean object registry. 36 | """ 37 | from prefect.context import PrefectObjectRegistry 38 | 39 | with PrefectObjectRegistry(): 40 | yield 41 | 42 | 43 | @pytest.fixture() 44 | def credentials_params(): 45 | return { 46 | "account": "account", 47 | "user": "user", 48 | "password": "password", 49 | } 50 | 51 | 52 | @pytest.fixture() 53 | def connector_params(credentials_params): 54 | snowflake_credentials = SnowflakeCredentials(**credentials_params) 55 | _connector_params = { 56 | "schema": "schema_input", 57 | "database": "database", 58 | "warehouse": "warehouse", 59 | "credentials": snowflake_credentials, 60 | } 61 | return _connector_params 62 | 63 | 64 | @pytest.fixture() 65 | def private_credentials_params(): 66 | return { 67 | "account": "account", 68 | "user": "user", 69 | "password": "letmein", 70 | "private_key": _read_test_file("test_cert.p8"), 71 | } 72 | 73 | 74 | @pytest.fixture() 75 | def private_key_path_credentials_params(): 76 | return { 77 | "account": "account", 78 | "user": "user", 79 | "private_key_path": "path/to/private/key", 80 | "private_key_passphrase": "letmein", 81 | } 82 | 83 | 84 | @pytest.fixture() 85 | def private_connector_params(private_credentials_params): 86 | 87 | snowflake_credentials = SnowflakeCredentials(**private_credentials_params) 88 | _connector_params = { 89 | "schema": "schema_input", 90 | "database": "database", 91 | "warehouse": "warehouse", 92 | "credentials": snowflake_credentials, 93 | } 94 | return _connector_params 95 | 96 | 97 | @pytest.fixture() 98 | def private_no_pass_credentials_params(): 99 | return { 100 | "account": "account", 101 | "user": "user", 102 | "password": "letmein", 103 | "private_key": _read_test_file("test_cert_no_pass.p8"), 104 | } 105 | 106 | 107 | @pytest.fixture() 108 | def private_no_pass_connector_params(private_no_pass_credentials_params): 109 | 110 | snowflake_credentials = SnowflakeCredentials(**private_no_pass_credentials_params) 111 | _connector_params = { 112 | "schema": "schema_input", 113 | "database": "database", 114 | "warehouse": "warehouse", 115 | "credentials": snowflake_credentials, 116 | } 117 | return _connector_params 118 | 119 | 120 | @pytest.fixture() 121 | def private_malformed_credentials_params(): 122 | return { 123 | "account": "account", 124 | "user": "user", 125 | "password": "letmein", 126 | "private_key": _read_test_file("test_cert_malformed_format.p8"), 127 | } 128 | 129 | 130 | @pytest.fixture(autouse=True) 131 | def snowflake_connect_mock(monkeypatch): 132 | mock_cursor = MagicMock(name="cursor mock") 133 | results = iter([0, 1, 2, 3, 4]) 134 | mock_cursor.return_value.fetchone.side_effect = lambda: (next(results),) 135 | mock_cursor.return_value.fetchmany.side_effect = lambda size: list( 136 | (next(results),) for i in range(size) 137 | ) 138 | mock_cursor.return_value.fetchall.side_effect = lambda: [ 139 | (result,) for result in results 140 | ] 141 | 142 | mock_connection = MagicMock(name="connection mock") 143 | mock_connection.return_value.is_still_running.return_value = False 144 | mock_connection.return_value.cursor = mock_cursor 145 | 146 | monkeypatch.setattr("snowflake.connector.connect", mock_connection) 147 | return mock_connection 148 | -------------------------------------------------------------------------------- /docs/gen_examples_catalog.py: -------------------------------------------------------------------------------- 1 | """ 2 | Locates all the examples in the Collection and puts them in a single page. 3 | """ 4 | 5 | import re 6 | from collections import defaultdict 7 | from inspect import getmembers, isclass, isfunction, ismodule 8 | from pathlib import Path 9 | from textwrap import dedent 10 | from types import ModuleType 11 | from typing import Callable, Set, Union 12 | 13 | import mkdocs_gen_files 14 | from griffe.dataclasses import Docstring 15 | from griffe.docstrings.dataclasses import DocstringSectionKind 16 | from griffe.docstrings.parsers import Parser, parse 17 | from prefect.logging.loggers import disable_logger 18 | from prefect.utilities.importtools import to_qualified_name 19 | 20 | import prefect_snowflake 21 | 22 | COLLECTION_SLUG = "prefect_snowflake" 23 | 24 | 25 | def skip_parsing(name: str, obj: Union[ModuleType, Callable], module_nesting: str): 26 | """ 27 | Skips parsing the object if it's a private object or if it's not in the 28 | module nesting, preventing imports from other libraries from being added to the 29 | examples catalog. 30 | """ 31 | try: 32 | wrong_module = not to_qualified_name(obj).startswith(module_nesting) 33 | except AttributeError: 34 | wrong_module = False 35 | return obj.__doc__ is None or name.startswith("_") or wrong_module 36 | 37 | 38 | def skip_block_load_code_example(code_example: str) -> bool: 39 | """ 40 | Skips the code example if it's just showing how to load a Block. 41 | """ 42 | return re.search(r'\.load\("BLOCK_NAME"\)\s*$', code_example.rstrip("`")) 43 | 44 | 45 | def get_code_examples(obj: Union[ModuleType, Callable]) -> Set[str]: 46 | """ 47 | Gathers all the code examples within an object. 48 | """ 49 | code_examples = set() 50 | with disable_logger("griffe.docstrings.google"): 51 | with disable_logger("griffe.agents.nodes"): 52 | docstring = Docstring(obj.__doc__) 53 | parsed_sections = parse(docstring, Parser.google) 54 | 55 | for section in parsed_sections: 56 | if section.kind == DocstringSectionKind.examples: 57 | code_example = "\n".join( 58 | (part[1] for part in section.as_dict().get("value", [])) 59 | ) 60 | if not skip_block_load_code_example(code_example): 61 | code_examples.add(code_example) 62 | if section.kind == DocstringSectionKind.admonition: 63 | value = section.as_dict().get("value", {}) 64 | if value.get("annotation") == "example": 65 | code_example = value.get("description") 66 | if not skip_block_load_code_example(code_example): 67 | code_examples.add(code_example) 68 | 69 | return code_examples 70 | 71 | 72 | code_examples_grouping = defaultdict(set) 73 | for module_name, module_obj in getmembers(prefect_snowflake, ismodule): 74 | 75 | module_nesting = f"{COLLECTION_SLUG}.{module_name}" 76 | # find all module examples 77 | if skip_parsing(module_name, module_obj, module_nesting): 78 | continue 79 | code_examples_grouping[module_name] |= get_code_examples(module_obj) 80 | 81 | # find all class and method examples 82 | for class_name, class_obj in getmembers(module_obj, isclass): 83 | if skip_parsing(class_name, class_obj, module_nesting): 84 | continue 85 | code_examples_grouping[module_name] |= get_code_examples(class_obj) 86 | for method_name, method_obj in getmembers(class_obj, isfunction): 87 | if skip_parsing(method_name, method_obj, module_nesting): 88 | continue 89 | code_examples_grouping[module_name] |= get_code_examples(method_obj) 90 | 91 | # find all function examples 92 | for function_name, function_obj in getmembers(module_obj, isfunction): 93 | if skip_parsing(function_name, function_obj, module_nesting): 94 | continue 95 | code_examples_grouping[module_name] |= get_code_examples(function_obj) 96 | 97 | 98 | examples_catalog_path = Path("examples_catalog.md") 99 | with mkdocs_gen_files.open(examples_catalog_path, "w") as generated_file: 100 | generated_file.write( 101 | dedent( 102 | """ 103 | # Examples Catalog 104 | 105 | Below is a list of examples for `prefect-snowflake`. 106 | """ 107 | ) 108 | ) 109 | for module_name, code_examples in code_examples_grouping.items(): 110 | if len(code_examples) == 0: 111 | continue 112 | module_title = module_name.replace("_", " ").title() 113 | generated_file.write( 114 | f"## [{module_title} Module][{COLLECTION_SLUG}.{module_name}]\n" 115 | ) 116 | for code_example in code_examples: 117 | generated_file.write(code_example + "\n") 118 | -------------------------------------------------------------------------------- /MAINTAINERS.md: -------------------------------------------------------------------------------- 1 | # prefect-snowflake 2 | 3 | ## Getting Started 4 | 5 | Now that you've bootstrapped a project, follow the steps below to get started developing your Prefect Collection! 6 | 7 | ### Python setup 8 | 9 | Requires an installation of Python 3.7+ 10 | 11 | We recommend using a Python virtual environment manager such as pipenv, conda or virtualenv. 12 | 13 | ### GitHub setup 14 | 15 | Create a Git respoitory for the newly generated collection and create the first commit: 16 | 17 | ```bash 18 | git init 19 | git add . 20 | git commit -m "Initial commit: project generated by prefect-collection-template" 21 | ``` 22 | 23 | Then, create a new repo following the prompts at: 24 | https://github.com/organizations/PrefectHQ/repositories/new 25 | 26 | Upon creation, push the repository to GitHub: 27 | ```bash 28 | git remote add origin https://github.com/PrefectHQ/prefect-snowflake.git 29 | git branch -M main 30 | git push -u origin main 31 | ``` 32 | 33 | It's recommended to setup some protection rules for main at: 34 | https://github.com/PrefectHQ/prefect-snowflake/settings/branches 35 | 36 | - Require a pull request before merging 37 | - Require approvals 38 | 39 | Lastly, [code owners](https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners) for the repository can be set, like this [example here](https://github.com/PrefectHQ/prefect/blob/master/.github/CODEOWNERS). 40 | 41 | ### Project setup 42 | 43 | To setup your project run the following: 44 | 45 | ```bash 46 | # Create an editable install of your project 47 | pip install -e ".[dev]" 48 | 49 | # Configure pre-commit hooks 50 | pre-commit install 51 | ``` 52 | 53 | To verify the setup was successful you can run the following: 54 | 55 | - Run the tests for tasks and flows in the collection: 56 | ```bash 57 | pytest tests 58 | ``` 59 | - Serve the docs with `mkdocs`: 60 | ```bash 61 | mkdocs serve 62 | ``` 63 | 64 | ## Developing tasks and flows 65 | 66 | For information about the use and development of tasks and flow, check out the [flows](https://orion-docs.prefect.io/concepts/flows/) and [tasks](https://orion-docs.prefect.io/concepts/tasks/) concepts docs in the Prefect docs. 67 | 68 | ## Writing documentation 69 | 70 | This collection has been setup to with [mkdocs](https://www.mkdocs.org/) for automatically generated documentation. The signatures and docstrings of your tasks and flow will be used to generate documentation for the users of this collection. You can make changes to the structure of the generated documentation by editing the `mkdocs.yml` file in this project. 71 | 72 | To add a new page for a module in your collection, create a new markdown file in the `docs` directory and add that file to the `nav` section of `mkdocs.yml`. If you want to automatically generate documentation based on the docstrings and signatures of the contents of the module with `mkdocstrings`, add a line to the new markdown file in the following format: 73 | 74 | ```markdown 75 | ::: prefect_snowflake.{module_name} 76 | ``` 77 | 78 | You can also refer to the `flows.md` and `tasks.md` files included in your generated project as examples. 79 | 80 | ## Development lifecycle 81 | 82 | ### CI Pipeline 83 | 84 | This collection comes with [GitHub Actions](https://docs.github.com/en/actions) for testing and linting. To add additional actions, you can add jobs in the `.github/workflows` folder. Upon a pull request, the pipeline will run linting via [`black`](https://black.readthedocs.io/en/stable/), [`flake8`](https://flake8.pycqa.org/en/latest/), [`interrogate`](https://interrogate.readthedocs.io/en/latest/), and unit tests via `pytest` alongside `coverage`. 85 | 86 | `interrogate` will tell you which methods, functions, classes, and modules have docstrings, and which do not--the job has a fail threshold of 95%, meaning that it will fail if more than 5% of the codebase is undocumented. We recommend following the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) for docstring format. 87 | 88 | Simiarly, `coverage` ensures that the codebase includes tests--the job has a fail threshold of 80%, meaning that it will fail if more than 20% of the codebase is missing tests. 89 | 90 | ### Track Issues on Project Board 91 | 92 | To automatically add issues to a GitHub Project Board, you'll need a [secret added](https://docs.github.com/en/actions/security-guides/encrypted-secrets#creating-encrypted-secrets-for-an-environment) to the repository. Specifically, a secret named `ADD_TO_PROJECT_URL`, formatted like `https://github.com/orgs//projects/`. 93 | 94 | ### Package and Publish 95 | 96 | GitHub actions will handle packaging and publishing of your collection to [PyPI](https://pypi.org/) so other Prefect users can your collection in their flows. 97 | 98 | To publish to PyPI, you'll need a PyPI account and to generate an API token to authenticate with PyPI when publishing new versions of your collection. The [PyPI documentation](https://pypi.org/help/#apitoken) outlines the steps needed to get an API token. 99 | 100 | Once you've obtained a PyPI API token, [create a GitHub secret](https://docs.github.com/en/actions/security-guides/encrypted-secrets#creating-encrypted-secrets-for-a-repository) named `PYPI_API_TOKEN`. 101 | 102 | To publish a new version of your collection, [create a new GitHub release](https://docs.github.com/en/repositories/releasing-projects-on-github/managing-releases-in-a-repository#creating-a-release) and tag it with the version that you want to deploy (e.g. v0.3.2). This will trigger a workflow to publish the new version on PyPI and deploy the updated docs to GitHub pages. 103 | 104 | Upon publishing, a `docs` branch is automatically created. To hook this up to GitHub Pages, simply head over to https://github.com/PrefectHQ/prefect-snowflake/settings/pages, select `docs` under the dropdown menu, keep the default `/root` folder, `Save`, and upon refresh, you should see a prompt stating "Your site is published at https://PrefectHQ.github.io/prefect-snowflake". Don't forget to add this link to the repo's "About" section, under "Website" so users can access the docs easily. 105 | 106 | Feel free to [submit your collection](https://orion-docs.prefect.io/collections/overview/#listing-in-the-collections-catalog) to the Prefect [Collections Catalog](https://orion-docs.prefect.io/collections/catalog/)! 107 | 108 | ## Further guidance 109 | 110 | If you run into any issues during the bootstrapping process, feel free to open an issue in the [prefect-collection-template](https://github.com/PrefectHQ/prefect-collection-template) repository. 111 | 112 | If you have any questions or issues while developing your collection, you can find help in either the [Prefect Discourse forum](https://discourse.prefect.io/) or the [Prefect Slack community](https://prefect.io/slack). 113 | -------------------------------------------------------------------------------- /tests/test_database.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import MagicMock 2 | 3 | import pytest 4 | from prefect import flow 5 | from pydantic import VERSION as PYDANTIC_VERSION 6 | 7 | if PYDANTIC_VERSION.startswith("2."): 8 | from pydantic.v1 import SecretBytes, SecretStr 9 | else: 10 | from pydantic import SecretBytes, SecretStr 11 | 12 | from snowflake.connector import DictCursor 13 | from snowflake.connector.cursor import SnowflakeCursor as OriginalSnowflakeCursorClass 14 | 15 | from prefect_snowflake.database import ( 16 | BEGIN_TRANSACTION_STATEMENT, 17 | END_TRANSACTION_STATEMENT, 18 | SnowflakeConnector, 19 | snowflake_multiquery, 20 | snowflake_query, 21 | snowflake_query_sync, 22 | ) 23 | 24 | 25 | def test_snowflake_connector_init(connector_params): 26 | snowflake_connector = SnowflakeConnector(**connector_params) 27 | actual_connector_params = snowflake_connector.dict() 28 | for param in connector_params: 29 | expected = connector_params[param] 30 | if param == "schema": 31 | param = "schema_" 32 | actual = actual_connector_params[param] 33 | if isinstance(actual, SecretStr): 34 | actual = actual.get_secret_value() 35 | assert actual == expected 36 | 37 | 38 | def test_snowflake_connector_password_is_secret_str(connector_params): 39 | snowflake_connector = SnowflakeConnector(**connector_params) 40 | password = snowflake_connector.credentials.password 41 | assert isinstance(password, SecretStr) 42 | assert password.get_secret_value() == "password" 43 | 44 | 45 | def test_snowflake_connector_private_key_is_secret(private_connector_params): 46 | snowflake_connector = SnowflakeConnector(**private_connector_params) 47 | private_key = snowflake_connector.credentials.private_key 48 | assert isinstance(private_key, (SecretStr, SecretBytes)) 49 | 50 | 51 | class SnowflakeCursor: 52 | def __enter__(self): 53 | return self 54 | 55 | def __exit__(self, *exc): 56 | return False 57 | 58 | def execute_async(self, query, params): 59 | query_id = "1234" 60 | self.result = {query_id: [(query, params)]} 61 | return {"queryId": query_id} 62 | 63 | def get_results_from_sfqid(self, query_id): 64 | self.query_result = self.result[query_id] 65 | 66 | def fetchall(self): 67 | return self.query_result 68 | 69 | def execute(self, query, params=None): 70 | self.query_result = [(query, params, "sync")] 71 | return self 72 | 73 | 74 | class SnowflakeConnection: 75 | def __enter__(self): 76 | return self 77 | 78 | def __exit__(self, *exc): 79 | return False 80 | 81 | def cursor(self, cursor_type): 82 | return SnowflakeCursor() 83 | 84 | def is_still_running(self, state): 85 | return state 86 | 87 | def get_query_status_throw_if_error(self, query_id): 88 | return False 89 | 90 | 91 | @pytest.fixture() 92 | def snowflake_connector(snowflake_connect_mock): 93 | snowflake_connector_mock = MagicMock() 94 | snowflake_connector_mock.get_connection.return_value = SnowflakeConnection() 95 | return snowflake_connector_mock 96 | 97 | 98 | def test_snowflake_query(snowflake_connector): 99 | @flow 100 | def test_flow(): 101 | result = snowflake_query( 102 | "query", 103 | snowflake_connector, 104 | params=("param",), 105 | ) 106 | return result 107 | 108 | result = test_flow() 109 | assert result[0][0] == "query" 110 | assert result[0][1] == ("param",) 111 | 112 | 113 | def test_snowflake_multiquery(snowflake_connector): 114 | @flow 115 | def test_flow(): 116 | result = snowflake_multiquery( 117 | ["query1", "query2"], 118 | snowflake_connector, 119 | params=("param",), 120 | ) 121 | return result 122 | 123 | result = test_flow() 124 | assert result[0][0][0] == "query1" 125 | assert result[0][0][1] == ("param",) 126 | assert result[1][0][0] == "query2" 127 | assert result[1][0][1] == ("param",) 128 | 129 | 130 | def test_snowflake_multiquery_transaction(snowflake_connector): 131 | @flow 132 | def test_flow(): 133 | result = snowflake_multiquery( 134 | ["query1", "query2"], 135 | snowflake_connector, 136 | params=("param",), 137 | as_transaction=True, 138 | ) 139 | return result 140 | 141 | result = test_flow() 142 | assert result[0][0][0] == "query1" 143 | assert result[0][0][1] == ("param",) 144 | assert result[1][0][0] == "query2" 145 | assert result[1][0][1] == ("param",) 146 | 147 | 148 | def test_snowflake_multiquery_transaction_with_transaction_control_results( 149 | snowflake_connector, 150 | ): 151 | @flow 152 | def test_flow(): 153 | result = snowflake_multiquery( 154 | ["query1", "query2"], 155 | snowflake_connector, 156 | params=("param",), 157 | as_transaction=True, 158 | return_transaction_control_results=True, 159 | ) 160 | return result 161 | 162 | result = test_flow() 163 | assert result[0][0][0] == BEGIN_TRANSACTION_STATEMENT 164 | assert result[1][0][0] == "query1" 165 | assert result[1][0][1] == ("param",) 166 | assert result[2][0][0] == "query2" 167 | assert result[2][0][1] == ("param",) 168 | assert result[3][0][0] == END_TRANSACTION_STATEMENT 169 | 170 | 171 | def test_snowflake_query_sync(snowflake_connector): 172 | @flow() 173 | def test_snowflake_query_sync_flow(): 174 | result = snowflake_query_sync("query", snowflake_connector, params=("param",)) 175 | return result 176 | 177 | result = test_snowflake_query_sync_flow() 178 | assert result[0][0] == "query" 179 | assert result[0][1] == ("param",) 180 | assert result[0][2] == "sync" 181 | 182 | 183 | def test_snowflake_private_connector_init(private_connector_params): 184 | snowflake_connector = SnowflakeConnector(**private_connector_params) 185 | actual_connector_params = snowflake_connector.dict() 186 | for param in private_connector_params: 187 | expected = private_connector_params[param] 188 | if param == "schema": 189 | param = "schema_" 190 | actual = actual_connector_params[param] 191 | if isinstance(actual, (SecretStr, SecretBytes)): 192 | actual = actual.get_secret_value() 193 | assert actual == expected 194 | 195 | 196 | class TestSnowflakeConnector: 197 | @pytest.fixture 198 | def snowflake_connector(self, connector_params, snowflake_connect_mock): 199 | connector = SnowflakeConnector(**connector_params) 200 | return connector 201 | 202 | def test_block_initialization(self, snowflake_connector): 203 | assert snowflake_connector._connection is None 204 | assert snowflake_connector._unique_cursors is None 205 | 206 | def test_get_connection(self, snowflake_connector: SnowflakeConnector, caplog): 207 | connection = snowflake_connector.get_connection() 208 | assert snowflake_connector._connection is connection 209 | assert caplog.records[0].msg == "Started a new connection to Snowflake." 210 | 211 | def test_reset_cursors(self, snowflake_connector: SnowflakeConnector, caplog): 212 | mock_cursor = MagicMock() 213 | snowflake_connector.reset_cursors() 214 | assert caplog.records[0].msg == "There were no cursors to reset." 215 | 216 | snowflake_connector._start_connection() 217 | snowflake_connector._unique_cursors["12345"] = mock_cursor 218 | snowflake_connector.reset_cursors() 219 | assert len(snowflake_connector._unique_cursors) == 0 220 | mock_cursor.close.assert_called_once() 221 | 222 | def test_fetch_one(self, snowflake_connector: SnowflakeConnector): 223 | result = snowflake_connector.fetch_one("query", parameters=("param",)) 224 | assert result == (0,) 225 | result = snowflake_connector.fetch_one("query", parameters=("param",)) 226 | assert result == (1,) 227 | 228 | def test_fetch_one_cursor_set_to_dict_cursor( 229 | self, snowflake_connector: SnowflakeConnector 230 | ): 231 | _ = snowflake_connector.fetch_one( 232 | "query", parameters=("param",), cursor_type=DictCursor 233 | ) 234 | args, _ = snowflake_connector._connection.cursor.call_args 235 | 236 | assert args[0] == DictCursor 237 | 238 | def test_fetch_one_cursor_default(self, snowflake_connector: SnowflakeConnector): 239 | _ = snowflake_connector.fetch_one("query", parameters=("param",)) 240 | args, kwargs = snowflake_connector._connection.cursor.call_args 241 | 242 | assert args[0] == OriginalSnowflakeCursorClass 243 | 244 | def test_fetch_all_cursor_set_to_dict_cursor( 245 | self, snowflake_connector: SnowflakeConnector 246 | ): 247 | _ = snowflake_connector.fetch_all( 248 | "query", parameters=("param",), cursor_type=DictCursor 249 | ) 250 | args, _ = snowflake_connector._connection.cursor.call_args 251 | 252 | assert args[0] == DictCursor 253 | 254 | def test_fetch_all_cursor_default(self, snowflake_connector: SnowflakeConnector): 255 | _ = snowflake_connector.fetch_all("query", parameters=("param",)) 256 | args, _ = snowflake_connector._connection.cursor.call_args 257 | 258 | assert args[0] == OriginalSnowflakeCursorClass 259 | 260 | def test_fetch_many(self, snowflake_connector: SnowflakeConnector): 261 | result = snowflake_connector.fetch_many("query", parameters=("param",), size=2) 262 | assert result == [(0,), (1,)] 263 | result = snowflake_connector.fetch_many("query", parameters=("param",)) 264 | assert result == [(2,)] 265 | 266 | def test_fetch_all(self, snowflake_connector: SnowflakeConnector): 267 | result = snowflake_connector.fetch_all("query", parameters=("param",)) 268 | assert result == [(0,), (1,), (2,), (3,), (4,)] 269 | 270 | def test_execute(self, snowflake_connector: SnowflakeConnector): 271 | assert snowflake_connector.execute("query", parameters=("param",)) is None 272 | 273 | def test_execute_Many(self, snowflake_connector: SnowflakeConnector): 274 | assert ( 275 | snowflake_connector.execute_many("query", seq_of_parameters=[("param",)]) 276 | is None 277 | ) 278 | 279 | def test_close(self, snowflake_connector: SnowflakeConnector, caplog): 280 | assert snowflake_connector.close() is None 281 | assert caplog.records[0].msg == "There were no cursors to reset." 282 | assert caplog.records[1].msg == "There was no connection open to be closed." 283 | 284 | snowflake_connector._start_connection() 285 | assert snowflake_connector.close() is None 286 | assert snowflake_connector._connection is None 287 | assert snowflake_connector._unique_cursors == {} 288 | 289 | def test_context_management(self, snowflake_connector): 290 | with snowflake_connector: 291 | assert snowflake_connector._connection is None 292 | assert snowflake_connector._unique_cursors is None 293 | 294 | assert snowflake_connector._connection is None 295 | assert snowflake_connector._unique_cursors is None 296 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | > [!NOTE] 2 | > Active development of this project has moved within PrefectHQ/prefect. The code can be found [here](https://github.com/PrefectHQ/prefect/tree/main/src/integrations/prefect-snowflake) and documentation [here](https://docs.prefect.io/latest/integrations/prefect-snowflake). 3 | > Please open issues and PRs against PrefectHQ/prefect instead of this repository. 4 | 5 | 6 | # prefect-snowflake 7 | 8 |

9 | 10 | PyPI 11 | 12 | 13 | 14 | 15 | 16 | 17 |
18 | 19 | 20 |

21 | 22 | Visit the full docs [here](https://PrefectHQ.github.io/prefect-snowflake) to see additional examples and the API reference. 23 | 24 | ## Welcome! 25 | 26 | The prefect-snowflake collection makes it easy to connect to a Snowflake database in your Prefect flows. Check out the examples below to get started! 27 | 28 | ## Getting Started 29 | 30 | ### Integrate with Prefect flows 31 | 32 | Prefect works with Snowflake by providing dataflow automation for faster, more efficient data pipeline creation, execution, and monitoring. 33 | 34 | This results in reduced errors, increased confidence in your data, and ultimately, faster insights. 35 | 36 | To set up a table, use the `execute` and `execute_many` methods. Then, use the `fetch_many` method to retrieve data in a stream until there's no more data. 37 | 38 | By using the `SnowflakeConnector` as a context manager, you can make sure that the Snowflake connection and cursors are closed properly after you're done with them. 39 | 40 | Be sure to install [prefect-snowflake](#installation) and [save to block](#saving-credentials-to-block) to run the examples below! 41 | 42 | === "Sync" 43 | 44 | ```python 45 | from prefect import flow, task 46 | from prefect_snowflake import SnowflakeConnector 47 | 48 | 49 | @task 50 | def setup_table(block_name: str) -> None: 51 | with SnowflakeConnector.load(block_name) as connector: 52 | connector.execute( 53 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 54 | ) 55 | connector.execute_many( 56 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 57 | seq_of_parameters=[ 58 | {"name": "Ford", "address": "Highway 42"}, 59 | {"name": "Unknown", "address": "Space"}, 60 | {"name": "Me", "address": "Myway 88"}, 61 | ], 62 | ) 63 | 64 | @task 65 | def fetch_data(block_name: str) -> list: 66 | all_rows = [] 67 | with SnowflakeConnector.load(block_name) as connector: 68 | while True: 69 | # Repeated fetch* calls using the same operation will 70 | # skip re-executing and instead return the next set of results 71 | new_rows = connector.fetch_many("SELECT * FROM customers", size=2) 72 | if len(new_rows) == 0: 73 | break 74 | all_rows.append(new_rows) 75 | return all_rows 76 | 77 | @flow 78 | def snowflake_flow(block_name: str) -> list: 79 | setup_table(block_name) 80 | all_rows = fetch_data(block_name) 81 | return all_rows 82 | 83 | snowflake_flow() 84 | ``` 85 | 86 | === "Async" 87 | 88 | ```python 89 | from prefect import flow, task 90 | from prefect_snowflake import SnowflakeConnector 91 | import asyncio 92 | 93 | @task 94 | async def setup_table(block_name: str) -> None: 95 | with await SnowflakeConnector.load(block_name) as connector: 96 | await connector.execute( 97 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 98 | ) 99 | await connector.execute_many( 100 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 101 | seq_of_parameters=[ 102 | {"name": "Ford", "address": "Highway 42"}, 103 | {"name": "Unknown", "address": "Space"}, 104 | {"name": "Me", "address": "Myway 88"}, 105 | ], 106 | ) 107 | 108 | @task 109 | async def fetch_data(block_name: str) -> list: 110 | all_rows = [] 111 | with await SnowflakeConnector.load(block_name) as connector: 112 | while True: 113 | # Repeated fetch* calls using the same operation will 114 | # skip re-executing and instead return the next set of results 115 | new_rows = await connector.fetch_many("SELECT * FROM customers", size=2) 116 | if len(new_rows) == 0: 117 | break 118 | all_rows.append(new_rows) 119 | return all_rows 120 | 121 | @flow 122 | async def snowflake_flow(block_name: str) -> list: 123 | await setup_table(block_name) 124 | all_rows = await fetch_data(block_name) 125 | return all_rows 126 | 127 | asyncio.run(snowflake_flow("example")) 128 | ``` 129 | 130 | ### Access underlying Snowflake connection 131 | 132 | If the native methods of the block don't meet your requirements, don't worry. 133 | 134 | You have the option to access the underlying Snowflake connection and utilize its built-in methods as well. 135 | 136 | ```python 137 | import pandas as pd 138 | from prefect import flow 139 | from prefect_snowflake.database import SnowflakeConnector 140 | from snowflake.connector.pandas_tools import write_pandas 141 | 142 | @flow 143 | def snowflake_write_pandas_flow(): 144 | connector = SnowflakeConnector.load("my-block") 145 | with connector.get_connection() as connection: 146 | table_name = "TABLE_NAME" 147 | ddl = "NAME STRING, NUMBER INT" 148 | statement = f'CREATE TABLE IF NOT EXISTS {table_name} ({ddl})' 149 | with connection.cursor() as cursor: 150 | cursor.execute(statement) 151 | 152 | # case sensitivity matters here! 153 | df = pd.DataFrame([('Marvin', 42), ('Ford', 88)], columns=['NAME', 'NUMBER']) 154 | success, num_chunks, num_rows, _ = write_pandas( 155 | conn=connection, 156 | df=df, 157 | table_name=table_name, 158 | database=snowflake_connector.database, 159 | schema=snowflake_connector.schema_ # note the "_" suffix 160 | ) 161 | ``` 162 | 163 | ## Resources 164 | 165 | For more tips on how to use tasks and flows in an integration, check out [Using Collections](https://docs.prefect.io/integrations/usage/)! 166 | 167 | ### Installation 168 | 169 | Install `prefect-snowflake` with `pip`: 170 | 171 | ```bash 172 | pip install prefect-snowflake 173 | ``` 174 | 175 | A list of available blocks in `prefect-snowflake` and their setup instructions can be found [here](https://PrefectHQ.github.io/prefect-snowflake/blocks_catalog). 176 | 177 | Requires an installation of Python 3.8+. 178 | 179 | We recommend using a Python virtual environment manager such as pipenv, conda or virtualenv. 180 | 181 | These tasks are designed to work with Prefect 2. For more information about how to use Prefect, please refer to the [Prefect documentation](https://docs.prefect.io/). 182 | 183 | ### Saving credentials to block 184 | 185 | Note, to use the `load` method on Blocks, you must already have a block document [saved through code](https://docs.prefect.io/concepts/blocks/#saving-blocks) or saved through the UI. 186 | 187 | Below is a walkthrough on saving a `SnowflakeCredentials` block through code. 188 | 189 | 1. Head over to https://app.snowflake.com/. 190 | 2. Login to your Snowflake account, e.g. nh12345.us-east-2.aws, with your username and password. 191 | 3. Use those credentials to fill replace the placeholders below. 192 | 193 | ```python 194 | from prefect_snowflake import SnowflakeCredentials 195 | 196 | credentials = SnowflakeCredentials( 197 | account="ACCOUNT-PLACEHOLDER", # resembles nh12345.us-east-2.aws 198 | user="USER-PLACEHOLDER", 199 | password="PASSWORD-PLACEHOLDER" 200 | ) 201 | credentials.save("CREDENTIALS-BLOCK-NAME-PLACEHOLDER") 202 | ``` 203 | 204 | Then, to create a `SnowflakeConnector` block: 205 | 206 | 1. After logging in, click on any worksheet. 207 | 2. On the left side, select a database and schema. 208 | 3. On the top right, select a warehouse. 209 | 3. Create a short script, replacing the placeholders below. 210 | 211 | ```python 212 | from prefect_snowflake import SnowflakeCredentials, SnowflakeConnector 213 | 214 | credentials = SnowflakeCredentials.load("CREDENTIALS-BLOCK-NAME-PLACEHOLDER") 215 | 216 | connector = SnowflakeConnector( 217 | credentials=credentials, 218 | database="DATABASE-PLACEHOLDER", 219 | schema="SCHEMA-PLACEHOLDER", 220 | warehouse="COMPUTE_WH", 221 | ) 222 | connector.save("CONNECTOR-BLOCK-NAME-PLACEHOLDER") 223 | ``` 224 | 225 | Congrats! You can now easily load the saved block, which holds your credentials and connection info: 226 | 227 | ```python 228 | from prefect_snowflake import SnowflakeCredentials, SnowflakeConnector 229 | 230 | SnowflakeCredentials.load("CREDENTIALS-BLOCK-NAME-PLACEHOLDER") 231 | SnowflakeConnector.load("CONNECTOR-BLOCK-NAME-PLACEHOLDER") 232 | ``` 233 | 234 | !!! info "Registering blocks" 235 | 236 | Register blocks in this module to 237 | [view and edit them](https://docs.prefect.io/ui/blocks/) 238 | on Prefect Cloud: 239 | 240 | ```bash 241 | prefect block register -m prefect_snowflake 242 | ``` 243 | 244 | A list of available blocks in `prefect-snowflake` and their setup instructions can be found [here](https://PrefectHQ.github.io/prefect-snowflake/blocks_catalog). 245 | 246 | ### Feedback 247 | 248 | If you encounter any bugs while using `prefect-snowflake`, feel free to open an issue in the [prefect-snowflake](https://github.com/PrefectHQ/prefect-snowflake) repository. 249 | 250 | If you have any questions or issues while using `prefect-snowflake`, you can find help in either the [Prefect Discourse forum](https://discourse.prefect.io/) or the [Prefect Slack community](https://prefect.io/slack). 251 | 252 | Feel free to star or watch [`prefect-snowflake`](https://github.com/PrefectHQ/prefect-snowflake) for updates too! 253 | 254 | ### Contributing 255 | 256 | If you'd like to help contribute to fix an issue or add a feature to `prefect-snowflake`, please [propose changes through a pull request from a fork of the repository](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork). 257 | 258 | Here are the steps: 259 | 260 | 1. [Fork the repository](https://docs.github.com/en/get-started/quickstart/fork-a-repo#forking-a-repository) 261 | 2. [Clone the forked repository](https://docs.github.com/en/get-started/quickstart/fork-a-repo#cloning-your-forked-repository) 262 | 3. Install the repository and its dependencies: 263 | ``` 264 | pip install -e ".[dev]" 265 | ``` 266 | 4. Make desired changes 267 | 5. Add tests 268 | 6. Insert an entry to [CHANGELOG.md](https://github.com/PrefectHQ/prefect-snowflake/blob/main/CHANGELOG.md) 269 | 7. Install `pre-commit` to perform quality checks prior to commit: 270 | ``` 271 | pre-commit install 272 | ``` 273 | 8. `git commit`, `git push`, and create a pull request 274 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright 2021 Prefect Technologies, Inc. 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. -------------------------------------------------------------------------------- /prefect_snowflake/credentials.py: -------------------------------------------------------------------------------- 1 | """Credentials block for authenticating with Snowflake.""" 2 | 3 | import re 4 | import warnings 5 | from pathlib import Path 6 | from typing import Any, Optional, Union 7 | 8 | from cryptography.hazmat.backends import default_backend 9 | from cryptography.hazmat.primitives.serialization import ( 10 | Encoding, 11 | NoEncryption, 12 | PrivateFormat, 13 | load_pem_private_key, 14 | ) 15 | 16 | try: 17 | from typing import Literal 18 | except ImportError: 19 | from typing_extensions import Literal 20 | 21 | import snowflake.connector 22 | from prefect.blocks.abstract import CredentialsBlock 23 | from pydantic import VERSION as PYDANTIC_VERSION 24 | 25 | if PYDANTIC_VERSION.startswith("2."): 26 | from pydantic.v1 import Field, SecretBytes, SecretField, SecretStr, root_validator 27 | else: 28 | from pydantic import Field, SecretBytes, SecretField, SecretStr, root_validator 29 | 30 | # PEM certificates have the pattern: 31 | # -----BEGIN PRIVATE KEY----- 32 | # <- multiple lines of encoded data-> 33 | # -----END PRIVATE KEY----- 34 | # 35 | # The regex captures the header and footer into groups 1 and 3, the body into group 2 36 | # group 1: "header" captures series of hyphens followed by anything that is 37 | # not a hyphen followed by another string of hyphens 38 | # group 2: "body" capture everything upto the next hyphen 39 | # group 3: "footer" duplicates group 1 40 | _SIMPLE_PEM_CERTIFICATE_REGEX = "^(-+[^-]+-+)([^-]+)(-+[^-]+-+)" 41 | 42 | 43 | class InvalidPemFormat(Exception): 44 | """Invalid PEM Format Certificate""" 45 | 46 | 47 | class SnowflakeCredentials(CredentialsBlock): 48 | """ 49 | Block used to manage authentication with Snowflake. 50 | 51 | Args: 52 | account (str): The snowflake account name. 53 | user (str): The user name used to authenticate. 54 | password (SecretStr): The password used to authenticate. 55 | private_key (SecretStr): The PEM used to authenticate. 56 | authenticator (str): The type of authenticator to use for initializing 57 | connection (oauth, externalbrowser, etc); refer to 58 | [Snowflake documentation](https://docs.snowflake.com/en/user-guide/python-connector-api.html#connect) 59 | for details, and note that `externalbrowser` will only 60 | work in an environment where a browser is available. 61 | token (SecretStr): The OAuth or JWT Token to provide when 62 | authenticator is set to OAuth. 63 | endpoint (str): The Okta endpoint to use when authenticator is 64 | set to `okta_endpoint`, e.g. `https://.okta.com`. 65 | role (str): The name of the default role to use. 66 | autocommit (bool): Whether to automatically commit. 67 | 68 | Example: 69 | Load stored Snowflake credentials: 70 | ```python 71 | from prefect_snowflake import SnowflakeCredentials 72 | 73 | snowflake_credentials_block = SnowflakeCredentials.load("BLOCK_NAME") 74 | ``` 75 | """ # noqa E501 76 | 77 | _block_type_name = "Snowflake Credentials" 78 | _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/bd359de0b4be76c2254bd329fe3a267a1a3879c2-250x250.png" # noqa 79 | _documentation_url = "https://prefecthq.github.io/prefect-snowflake/credentials/#prefect_snowflake.credentials.SnowflakeCredentials" # noqa 80 | 81 | account: str = Field( 82 | ..., description="The snowflake account name.", example="nh12345.us-east-2.aws" 83 | ) 84 | user: str = Field(..., description="The user name used to authenticate.") 85 | password: Optional[SecretStr] = Field( 86 | default=None, description="The password used to authenticate." 87 | ) 88 | private_key: Optional[SecretBytes] = Field( 89 | default=None, description="The PEM used to authenticate." 90 | ) 91 | private_key_path: Optional[Path] = Field( 92 | default=None, description="The path to the private key." 93 | ) 94 | private_key_passphrase: Optional[SecretStr] = Field( 95 | default=None, description="The password to use for the private key." 96 | ) 97 | authenticator: Literal[ 98 | "snowflake", 99 | "snowflake_jwt", 100 | "externalbrowser", 101 | "okta_endpoint", 102 | "oauth", 103 | "username_password_mfa", 104 | ] = Field( # noqa 105 | default="snowflake", 106 | description=("The type of authenticator to use for initializing connection."), 107 | ) 108 | token: Optional[SecretStr] = Field( 109 | default=None, 110 | description=( 111 | "The OAuth or JWT Token to provide when authenticator is set to `oauth`." 112 | ), 113 | ) 114 | endpoint: Optional[str] = Field( 115 | default=None, 116 | description=( 117 | "The Okta endpoint to use when authenticator is set to `okta_endpoint`." 118 | ), 119 | ) 120 | role: Optional[str] = Field( 121 | default=None, description="The name of the default role to use." 122 | ) 123 | autocommit: Optional[bool] = Field( 124 | default=None, description="Whether to automatically commit." 125 | ) 126 | 127 | @root_validator(pre=True) 128 | def _validate_auth_kwargs(cls, values): 129 | """ 130 | Ensure an authorization value has been provided by the user. 131 | """ 132 | auth_params = ( 133 | "password", 134 | "private_key", 135 | "private_key_path", 136 | "authenticator", 137 | "token", 138 | ) 139 | if not any(values.get(param) for param in auth_params): 140 | auth_str = ", ".join(auth_params) 141 | raise ValueError( 142 | f"One of the authentication keys must be provided: {auth_str}\n" 143 | ) 144 | elif values.get("private_key") and values.get("private_key_path"): 145 | raise ValueError( 146 | "Do not provide both private_key and private_key_path; select one." 147 | ) 148 | elif values.get("password") and values.get("private_key_passphrase"): 149 | raise ValueError( 150 | "Do not provide both password and private_key_passphrase; " 151 | "specify private_key_passphrase only instead." 152 | ) 153 | return values 154 | 155 | @root_validator(pre=True) 156 | def _validate_token_kwargs(cls, values): 157 | """ 158 | Ensure an authorization value has been provided by the user. 159 | """ 160 | authenticator = values.get("authenticator") 161 | token = values.get("token") 162 | if authenticator == "oauth" and not token: 163 | raise ValueError( 164 | "If authenticator is set to `oauth`, `token` must be provided" 165 | ) 166 | return values 167 | 168 | @root_validator(pre=True) 169 | def _validate_okta_kwargs(cls, values): 170 | """ 171 | Ensure an authorization value has been provided by the user. 172 | """ 173 | authenticator = values.get("authenticator") 174 | 175 | # did not want to make a breaking change so we will allow both 176 | # see https://github.com/PrefectHQ/prefect-snowflake/issues/44 177 | if "okta_endpoint" in values.keys(): 178 | warnings.warn( 179 | "Please specify `endpoint` instead of `okta_endpoint`; " 180 | "`okta_endpoint` will be removed March 31, 2023.", 181 | DeprecationWarning, 182 | stacklevel=2, 183 | ) 184 | # remove okta endpoint from fields 185 | okta_endpoint = values.pop("okta_endpoint") 186 | if "endpoint" not in values.keys(): 187 | values["endpoint"] = okta_endpoint 188 | 189 | endpoint = values.get("endpoint") 190 | if authenticator == "okta_endpoint" and not endpoint: 191 | raise ValueError( 192 | "If authenticator is set to `okta_endpoint`, " 193 | "`endpoint` must be provided" 194 | ) 195 | return values 196 | 197 | def resolve_private_key(self) -> Optional[bytes]: 198 | """ 199 | Converts a PEM encoded private key into a DER binary key. 200 | 201 | Returns: 202 | DER encoded key if private_key has been provided otherwise returns None. 203 | 204 | Raises: 205 | InvalidPemFormat: If private key is not in PEM format. 206 | """ 207 | if self.private_key_path is None and self.private_key is None: 208 | return None 209 | elif self.private_key_path: 210 | private_key = self.private_key_path.read_bytes() 211 | else: 212 | private_key = self._decode_secret(self.private_key) 213 | 214 | if self.private_key_passphrase is not None: 215 | password = self._decode_secret(self.private_key_passphrase) 216 | elif self.password is not None: 217 | warnings.warn( 218 | "Using the password field for private_key is deprecated " 219 | "and will not work after March 31, 2023; please use " 220 | "private_key_passphrase instead", 221 | DeprecationWarning, 222 | stacklevel=2, 223 | ) 224 | password = self._decode_secret(self.password) 225 | else: 226 | password = None 227 | 228 | composed_private_key = self._compose_pem(private_key) 229 | return load_pem_private_key( 230 | data=composed_private_key, 231 | password=password, 232 | backend=default_backend(), 233 | ).private_bytes( 234 | encoding=Encoding.DER, 235 | format=PrivateFormat.PKCS8, 236 | encryption_algorithm=NoEncryption(), 237 | ) 238 | 239 | @staticmethod 240 | def _decode_secret(secret: Union[SecretStr, SecretBytes]) -> Optional[bytes]: 241 | """ 242 | Decode the provided secret into bytes. If the secret is not a 243 | string or bytes, or it is whitespace, then return None. 244 | 245 | Args: 246 | secret: The value to decode. 247 | 248 | Returns: 249 | The decoded secret as bytes. 250 | 251 | """ 252 | if isinstance(secret, (SecretBytes, SecretStr)): 253 | secret = secret.get_secret_value() 254 | 255 | if not isinstance(secret, (bytes, str)) or len(secret) == 0 or secret.isspace(): 256 | return None 257 | 258 | return secret if isinstance(secret, bytes) else secret.encode() 259 | 260 | @staticmethod 261 | def _compose_pem(private_key: bytes) -> bytes: 262 | """Validate structure of PEM certificate. 263 | 264 | The original key passed from Prefect is sometimes malformed. 265 | This function recomposes the key into a valid key that will 266 | pass the serialization step when resolving the key to a DER. 267 | 268 | Args: 269 | private_key: A valid PEM format byte encoded string. 270 | 271 | Returns: 272 | byte encoded certificate. 273 | 274 | Raises: 275 | InvalidPemFormat: if private key is an invalid format. 276 | """ 277 | pem_parts = re.match(_SIMPLE_PEM_CERTIFICATE_REGEX, private_key.decode()) 278 | if pem_parts is None: 279 | raise InvalidPemFormat() 280 | 281 | body = "\n".join(re.split(r"\s+", pem_parts[2].strip())) 282 | # reassemble header+body+footer 283 | return f"{pem_parts[1]}\n{body}\n{pem_parts[3]}".encode() 284 | 285 | def get_client( 286 | self, **connect_kwargs: Any 287 | ) -> snowflake.connector.SnowflakeConnection: 288 | """ 289 | Returns an authenticated connection that can be used to query 290 | Snowflake databases. 291 | 292 | Any additional arguments passed to this method will be used to configure 293 | the SnowflakeConnection. For available parameters, please refer to the 294 | [Snowflake Python connector documentation](https://docs.snowflake.com/en/user-guide/python-connector-api.html#connect). 295 | 296 | Args: 297 | **connect_kwargs: Additional arguments to pass to 298 | `snowflake.connector.connect`. 299 | 300 | Returns: 301 | An authenticated Snowflake connection. 302 | 303 | Example: 304 | Get Snowflake connection with only block configuration: 305 | ```python 306 | from prefect_snowflake import SnowflakeCredentials 307 | 308 | snowflake_credentials_block = SnowflakeCredentials.load("BLOCK_NAME") 309 | 310 | connection = snowflake_credentials_block.get_client() 311 | ``` 312 | 313 | Get Snowflake connector scoped to a specified database: 314 | ```python 315 | from prefect_snowflake import SnowflakeCredentials 316 | 317 | snowflake_credentials_block = SnowflakeCredentials.load("BLOCK_NAME") 318 | 319 | connection = snowflake_credentials_block.get_client(database="my_database") 320 | ``` 321 | """ # noqa 322 | connect_params = { 323 | # required to track task's usage in the Snowflake Partner Network Portal 324 | "application": "Prefect_Snowflake_Collection", 325 | **self.dict(exclude_unset=True, exclude={"block_type_slug"}), 326 | **connect_kwargs, 327 | } 328 | 329 | for key, value in connect_params.items(): 330 | if isinstance(value, SecretField): 331 | connect_params[key] = connect_params[key].get_secret_value() 332 | 333 | # set authenticator to the actual okta_endpoint 334 | if connect_params.get("authenticator") == "okta_endpoint": 335 | endpoint = connect_params.pop("endpoint", None) or connect_params.pop( 336 | "okta_endpoint", None 337 | ) # okta_endpoint is deprecated 338 | connect_params["authenticator"] = endpoint 339 | 340 | private_der_key = self.resolve_private_key() 341 | if private_der_key is not None: 342 | connect_params["private_key"] = private_der_key 343 | connect_params.pop("password", None) 344 | connect_params.pop("private_key_passphrase", None) 345 | 346 | return snowflake.connector.connect(**connect_params) 347 | -------------------------------------------------------------------------------- /tests/test_credentials.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pathlib import Path 3 | from unittest.mock import MagicMock 4 | 5 | import pytest 6 | from prefect import flow 7 | from prefect.utilities.filesystem import relative_path_to_current_platform 8 | from pydantic import VERSION as PYDANTIC_VERSION 9 | 10 | if PYDANTIC_VERSION.startswith("2."): 11 | from pydantic.v1 import SecretBytes, SecretStr 12 | else: 13 | from pydantic import SecretBytes, SecretStr 14 | 15 | from prefect_snowflake.credentials import InvalidPemFormat, SnowflakeCredentials 16 | from prefect_snowflake.database import SnowflakeConnector 17 | 18 | 19 | def test_snowflake_credentials_init(credentials_params): 20 | snowflake_credentials = SnowflakeCredentials(**credentials_params) 21 | actual_credentials_params = snowflake_credentials.dict() 22 | for param in credentials_params: 23 | actual = actual_credentials_params[param] 24 | expected = credentials_params[param] 25 | if isinstance(actual, SecretStr): 26 | actual = actual.get_secret_value() 27 | assert actual == expected 28 | 29 | 30 | def test_snowflake_credentials_validate_auth_kwargs(credentials_params): 31 | credentials_params_missing = credentials_params.copy() 32 | credentials_params_missing.pop("password") 33 | with pytest.raises(ValueError, match="One of the authentication keys"): 34 | SnowflakeCredentials(**credentials_params_missing) 35 | 36 | 37 | def test_snowflake_credentials_validate_auth_kwargs_private_keys(credentials_params): 38 | credentials_params["private_key"] = "key" 39 | credentials_params["private_key_path"] = "keypath" 40 | with pytest.raises(ValueError, match="Do not provide both private_key and private"): 41 | SnowflakeCredentials(**credentials_params) 42 | 43 | 44 | def test_snowflake_credentials_validate_auth_kwargs_private_key_password( 45 | credentials_params, 46 | ): 47 | credentials_params["private_key_passphrase"] = "key" 48 | with pytest.raises(ValueError, match="Do not provide both password and private_"): 49 | SnowflakeCredentials(**credentials_params) 50 | 51 | 52 | def test_snowflake_credentials_validate_token_kwargs(credentials_params): 53 | credentials_params_missing = credentials_params.copy() 54 | credentials_params_missing.pop("password") 55 | credentials_params_missing["authenticator"] = "oauth" 56 | with pytest.raises(ValueError, match="If authenticator is set to `oauth`"): 57 | SnowflakeCredentials(**credentials_params_missing) 58 | 59 | # now test if passing both works 60 | credentials_params_missing["token"] = "some_token" 61 | assert SnowflakeCredentials(**credentials_params_missing) 62 | 63 | 64 | def test_snowflake_credentials_validate_okta_endpoint_kwargs(credentials_params): 65 | credentials_params_missing = credentials_params.copy() 66 | credentials_params_missing.pop("password") 67 | credentials_params_missing["authenticator"] = "okta_endpoint" 68 | with pytest.raises(ValueError, match="If authenticator is set to `okta_endpoint`"): 69 | SnowflakeCredentials(**credentials_params_missing) 70 | 71 | # now test if passing both works 72 | credentials_params_missing["endpoint"] = "https://account_name.okta.com" 73 | snowflake_credentials = SnowflakeCredentials(**credentials_params_missing) 74 | assert snowflake_credentials.endpoint == "https://account_name.okta.com" 75 | 76 | 77 | def test_snowflake_credentials_support_deprecated_okta_endpoint(credentials_params): 78 | credentials_params_missing = credentials_params.copy() 79 | credentials_params_missing.pop("password") 80 | credentials_params_missing["authenticator"] = "okta_endpoint" 81 | credentials_params_missing["okta_endpoint"] = "deprecated.com" 82 | snowflake_credentials = SnowflakeCredentials(**credentials_params_missing) 83 | assert snowflake_credentials.endpoint == "deprecated.com" 84 | 85 | 86 | def test_snowflake_credentials_support_endpoint_overrides_okta_endpoint( 87 | credentials_params, 88 | ): 89 | credentials_params_missing = credentials_params.copy() 90 | credentials_params_missing.pop("password") 91 | credentials_params_missing["authenticator"] = "okta_endpoint" 92 | credentials_params_missing["okta_endpoint"] = "deprecated.com" 93 | credentials_params_missing["endpoint"] = "new.com" 94 | snowflake_credentials = SnowflakeCredentials(**credentials_params_missing) 95 | assert snowflake_credentials.endpoint == "new.com" 96 | 97 | 98 | def test_snowflake_private_credentials_init(private_credentials_params): 99 | snowflake_credentials = SnowflakeCredentials(**private_credentials_params) 100 | actual_credentials_params = snowflake_credentials.dict() 101 | for param in private_credentials_params: 102 | actual = actual_credentials_params[param] 103 | expected = private_credentials_params[param] 104 | if isinstance(actual, (SecretStr, SecretBytes)): 105 | actual = actual.get_secret_value() 106 | if sys.platform != "win32": 107 | assert actual == expected 108 | 109 | 110 | def test_snowflake_private_credentials_malformed_certificate( 111 | private_credentials_params, private_malformed_credentials_params 112 | ): 113 | correct = SnowflakeCredentials(**private_credentials_params) 114 | malformed = SnowflakeCredentials(**private_malformed_credentials_params) 115 | c1 = correct.resolve_private_key() 116 | c2 = malformed.resolve_private_key() 117 | assert isinstance(c1, bytes) 118 | assert isinstance(c2, bytes) 119 | assert c1 == c2 120 | 121 | 122 | def test_snowflake_private_credentials_invalid_certificate(private_credentials_params): 123 | private_credentials_params["private_key"] = "---- INVALID CERTIFICATE ----" 124 | with pytest.raises(InvalidPemFormat): 125 | SnowflakeCredentials(**private_credentials_params).resolve_private_key() 126 | 127 | 128 | def test_snowflake_credentials_validate_private_key_password( 129 | private_credentials_params, 130 | ): 131 | credentials_params_missing = private_credentials_params.copy() 132 | password = credentials_params_missing.pop("password") 133 | private_key = credentials_params_missing.pop("private_key") 134 | assert password == "letmein" 135 | assert isinstance(private_key, bytes) 136 | # Test cert as string 137 | credentials = SnowflakeCredentials(**private_credentials_params) 138 | assert credentials.resolve_private_key() is not None 139 | 140 | 141 | def test_snowflake_credentials_validate_private_key_passphrase( 142 | private_credentials_params, 143 | ): 144 | private_credentials_params[ 145 | "private_key_passphrase" 146 | ] = private_credentials_params.pop("password") 147 | credentials_params_missing = private_credentials_params.copy() 148 | password = credentials_params_missing.pop("private_key_passphrase") 149 | private_key = credentials_params_missing.pop("private_key") 150 | assert password == "letmein" 151 | assert isinstance(private_key, bytes) 152 | # Test cert as string 153 | credentials = SnowflakeCredentials(**private_credentials_params) 154 | assert credentials.resolve_private_key() is not None 155 | 156 | 157 | def test_snowflake_credentials_validate_private_key_path( 158 | private_credentials_params, tmp_path 159 | ): 160 | private_key_path = tmp_path / "private_key.pem" 161 | private_key_path.write_bytes(private_credentials_params.pop("private_key")) 162 | private_credentials_params["private_key_path"] = private_key_path 163 | private_credentials_params[ 164 | "private_key_passphrase" 165 | ] = private_credentials_params.pop("password") 166 | credentials = SnowflakeCredentials(**private_credentials_params) 167 | assert credentials.resolve_private_key() is not None 168 | 169 | 170 | def test_snowflake_credentials_validate_private_key_invalid(private_credentials_params): 171 | credentials_params_missing = private_credentials_params.copy() 172 | private_key = credentials_params_missing.pop("private_key") 173 | assert isinstance(private_key, bytes) 174 | assert private_key.startswith(b"----") 175 | with pytest.raises(ValueError, match="Bad decrypt. Incorrect password?"): 176 | credentials = SnowflakeCredentials(**private_credentials_params) 177 | credentials.password = "_wrong_password" 178 | assert credentials.resolve_private_key() is not None 179 | 180 | 181 | def test_snowflake_credentials_validate_private_key_unexpected_password( 182 | private_credentials_params, 183 | ): 184 | credentials_params_missing = private_credentials_params.copy() 185 | private_key = credentials_params_missing.pop("private_key") 186 | assert isinstance(private_key, bytes) 187 | assert private_key.startswith(b"----") 188 | with pytest.raises( 189 | TypeError, match="Password was not given but private key is encrypted" 190 | ): 191 | credentials = SnowflakeCredentials(**private_credentials_params) 192 | credentials.password = None 193 | assert credentials.resolve_private_key() is not None 194 | 195 | 196 | def test_snowflake_credentials_validate_private_key_no_pass_password( 197 | private_no_pass_credentials_params, 198 | ): 199 | credentials_params_missing = private_no_pass_credentials_params.copy() 200 | password = credentials_params_missing.pop("password") 201 | private_key = credentials_params_missing.pop("private_key") 202 | assert password == "letmein" 203 | assert isinstance(private_key, bytes) 204 | assert private_key.startswith(b"----") 205 | 206 | with pytest.raises( 207 | TypeError, match="Password was given but private key is not encrypted" 208 | ): 209 | credentials = SnowflakeCredentials(**private_no_pass_credentials_params) 210 | assert credentials.resolve_private_key() is not None 211 | 212 | 213 | def test_snowflake_credentials_validate_private_key_is_pem( 214 | private_no_pass_credentials_params, 215 | ): 216 | private_no_pass_credentials_params["private_key"] = "_invalid_key_" 217 | with pytest.raises(InvalidPemFormat): 218 | credentials = SnowflakeCredentials(**private_no_pass_credentials_params) 219 | assert credentials.resolve_private_key() is not None 220 | 221 | 222 | def test_snowflake_credentials_validate_private_key_is_pem_bytes( 223 | private_no_pass_credentials_params, 224 | ): 225 | private_no_pass_credentials_params["private_key"] = "_invalid_key_" 226 | with pytest.raises(InvalidPemFormat): 227 | credentials = SnowflakeCredentials(**private_no_pass_credentials_params) 228 | assert credentials.resolve_private_key() is not None 229 | 230 | 231 | def test_snowflake_credentials_validate_private_key_path_init( 232 | private_key_path_credentials_params, 233 | ): 234 | snowflake_credentials = SnowflakeCredentials(**private_key_path_credentials_params) 235 | actual_credentials_params = snowflake_credentials.dict() 236 | for param in private_key_path_credentials_params: 237 | actual = actual_credentials_params[param] 238 | expected = private_key_path_credentials_params[param] 239 | if isinstance(actual, (SecretStr, SecretBytes)): 240 | actual = actual.get_secret_value() 241 | elif isinstance(actual, Path): 242 | actual = relative_path_to_current_platform(actual) 243 | expected = relative_path_to_current_platform(expected) 244 | assert actual == expected 245 | 246 | 247 | def test_get_client(credentials_params, snowflake_connect_mock: MagicMock): 248 | snowflake_credentials = SnowflakeCredentials(**credentials_params) 249 | snowflake_credentials.get_client() 250 | snowflake_connect_mock.assert_called_with( 251 | application="Prefect_Snowflake_Collection", 252 | account="account", 253 | user="user", 254 | password="password", 255 | ) 256 | 257 | 258 | def test_get_client_okta_endpoint( 259 | credentials_params, snowflake_connect_mock: MagicMock 260 | ): 261 | okta_endpoint = "https://account_name.okta.com" 262 | credentials_params_okta_endpoint = credentials_params.copy() 263 | del credentials_params_okta_endpoint["password"] 264 | credentials_params_okta_endpoint["authenticator"] = "okta_endpoint" 265 | credentials_params_okta_endpoint["endpoint"] = okta_endpoint 266 | snowflake_credentials = SnowflakeCredentials(**credentials_params_okta_endpoint) 267 | snowflake_credentials.get_client() 268 | snowflake_connect_mock.assert_called_with( 269 | application="Prefect_Snowflake_Collection", 270 | account="account", 271 | user="user", 272 | authenticator="https://account_name.okta.com", 273 | ) 274 | 275 | 276 | def test_snowflake_credentials_deprecated_okta_endpoint( 277 | credentials_params, snowflake_connect_mock: MagicMock 278 | ): 279 | okta_endpoint = "https://account_name.okta.com" 280 | credentials_params_okta_endpoint = credentials_params.copy() 281 | del credentials_params_okta_endpoint["password"] 282 | credentials_params_okta_endpoint["authenticator"] = "okta_endpoint" 283 | credentials_params_okta_endpoint["endpoint"] = okta_endpoint 284 | snowflake_credentials = SnowflakeCredentials(**credentials_params_okta_endpoint) 285 | snowflake_credentials.get_client() 286 | snowflake_connect_mock.assert_called_with( 287 | application="Prefect_Snowflake_Collection", 288 | account="account", 289 | user="user", 290 | authenticator="https://account_name.okta.com", 291 | ) 292 | 293 | 294 | def test_snowflake_credentials_unencrypted_private_key_password( 295 | private_no_pass_credentials_params, 296 | ): 297 | snowflake_credentials = SnowflakeCredentials(**private_no_pass_credentials_params) 298 | assert snowflake_credentials.private_key is not None 299 | assert snowflake_credentials.password is not None 300 | # Raises error if invalid 301 | with pytest.raises( 302 | TypeError, match="Password was given but private key is not encrypted" 303 | ): 304 | snowflake_credentials.get_client() 305 | 306 | 307 | def test_snowflake_credentials_unencrypted_private_key_no_password( 308 | private_no_pass_credentials_params, snowflake_connect_mock: MagicMock 309 | ): 310 | snowflake_credentials = SnowflakeCredentials(**private_no_pass_credentials_params) 311 | snowflake_credentials.password = None 312 | assert snowflake_credentials.private_key is not None 313 | # Raises error if invalid 314 | snowflake_credentials.get_client() 315 | 316 | 317 | def test_snowflake_credentials_unencrypted_private_key_empty_password( 318 | private_no_pass_credentials_params, snowflake_connect_mock: MagicMock 319 | ): 320 | snowflake_credentials = SnowflakeCredentials(**private_no_pass_credentials_params) 321 | assert snowflake_credentials.private_key is not None 322 | 323 | snowflake_credentials.password = SecretBytes(b" ") 324 | snowflake_credentials.get_client() 325 | snowflake_credentials.password = SecretBytes(b"") 326 | snowflake_credentials.get_client() 327 | snowflake_credentials.password = SecretStr("") 328 | snowflake_credentials.get_client() 329 | snowflake_credentials.password = SecretStr(" ") 330 | snowflake_credentials.get_client() 331 | 332 | 333 | def test_snowflake_credentials_encrypted_private_key_is_valid( 334 | private_credentials_params, snowflake_connect_mock: MagicMock 335 | ): 336 | snowflake_credentials = SnowflakeCredentials(**private_credentials_params) 337 | assert snowflake_credentials.private_key is not None 338 | assert snowflake_credentials.password is not None 339 | # Raises error if invalid 340 | snowflake_credentials.get_client() 341 | 342 | 343 | def test_snowflake_with_no_private_key_flow(): 344 | """ 345 | https://github.com/PrefectHQ/prefect-snowflake/issues/62 346 | """ 347 | 348 | @flow 349 | def snowflake_query_flow(): 350 | snowflake_credentials = SnowflakeCredentials( 351 | account="account", user="user", password="password", role="MY_ROLE" 352 | ) 353 | snowflake_connector = SnowflakeConnector( 354 | database="database", 355 | warehouse="warehouse", 356 | schema="schema", 357 | credentials=snowflake_credentials, 358 | ) 359 | return snowflake_connector 360 | 361 | snowflake_connector = snowflake_query_flow() 362 | assert isinstance(snowflake_connector, SnowflakeConnector) 363 | assert isinstance(snowflake_connector.credentials, SnowflakeCredentials) 364 | assert snowflake_connector.credentials.private_key is None 365 | assert snowflake_connector.credentials.private_key_path is None 366 | assert snowflake_connector.credentials.private_key_passphrase is None 367 | assert snowflake_connector.credentials.password.get_secret_value() == "password" 368 | 369 | 370 | def test_snowflake_with_private_key_path_flow(): 371 | """ 372 | https://github.com/PrefectHQ/prefect-snowflake/issues/62 373 | """ 374 | 375 | @flow 376 | def snowflake_query_flow(): 377 | snowflake_credentials = SnowflakeCredentials( 378 | account="account", 379 | user="user", 380 | private_key_path="private_key_path", 381 | private_key_passphrase="passphrase", 382 | ) 383 | snowflake_connector = SnowflakeConnector( 384 | database="database", 385 | warehouse="warehouse", 386 | schema="schema", 387 | credentials=snowflake_credentials, 388 | ) 389 | return snowflake_connector 390 | 391 | snowflake_connector = snowflake_query_flow() 392 | assert isinstance(snowflake_connector, SnowflakeConnector) 393 | assert isinstance(snowflake_connector.credentials, SnowflakeCredentials) 394 | assert snowflake_connector.credentials.private_key is None 395 | assert snowflake_connector.credentials.private_key_path == Path("private_key_path") 396 | assert ( 397 | snowflake_connector.credentials.private_key_passphrase.get_secret_value() 398 | == "passphrase" 399 | ) 400 | assert snowflake_connector.credentials.password is None 401 | 402 | 403 | def test_snowflake_connect_params(credentials_params, snowflake_connect_mock): 404 | snowflake_credentials = SnowflakeCredentials(**credentials_params) 405 | 406 | # The returned client is mocked, use the fixture instead 407 | _ = snowflake_credentials.get_client(autocommit=False) 408 | 409 | assert snowflake_connect_mock.call_args_list[0][1]["autocommit"] is False 410 | -------------------------------------------------------------------------------- /prefect_snowflake/_version.py: -------------------------------------------------------------------------------- 1 | # This file helps to compute a version number in source trees obtained from 2 | # git-archive tarball (such as those provided by githubs download-from-tag 3 | # feature). Distribution tarballs (built by setup.py sdist) and build 4 | # directories (produced by setup.py build) will contain a much shorter file 5 | # that just contains the computed version number. 6 | 7 | # This file is released into the public domain. Generated by 8 | # versioneer-0.21 (https://github.com/python-versioneer/python-versioneer) 9 | 10 | """Git implementation of _version.py.""" 11 | 12 | import errno 13 | import os 14 | import re 15 | import subprocess 16 | import sys 17 | from typing import Callable, Dict 18 | 19 | 20 | def get_keywords(): 21 | """Get the keywords needed to look up the version information.""" 22 | # these strings will be replaced by git during git-archive. 23 | # setup.py/versioneer.py will grep for the variable names, so they must 24 | # each be defined on a line of their own. _version.py will just call 25 | # get_keywords(). 26 | git_refnames = " (HEAD -> main)" 27 | git_full = "4ca016910935059fe6fb77c964a954930384a564" 28 | git_date = "2024-04-26 11:34:14 -0500" 29 | keywords = {"refnames": git_refnames, "full": git_full, "date": git_date} 30 | return keywords 31 | 32 | 33 | class VersioneerConfig: 34 | """Container for Versioneer configuration parameters.""" 35 | 36 | 37 | def get_config(): 38 | """Create, populate and return the VersioneerConfig() object.""" 39 | # these strings are filled in when 'setup.py versioneer' creates 40 | # _version.py 41 | cfg = VersioneerConfig() 42 | cfg.VCS = "git" 43 | cfg.style = "pep440" 44 | cfg.tag_prefix = "" 45 | cfg.parentdir_prefix = "" 46 | cfg.versionfile_source = "prefect_snowflake/_version.py" 47 | cfg.verbose = False 48 | return cfg 49 | 50 | 51 | class NotThisMethod(Exception): 52 | """Exception raised if a method is not valid for the current scenario.""" 53 | 54 | 55 | LONG_VERSION_PY: Dict[str, str] = {} 56 | HANDLERS: Dict[str, Dict[str, Callable]] = {} 57 | 58 | 59 | def register_vcs_handler(vcs, method): # decorator 60 | """Create decorator to mark a method as the handler of a VCS.""" 61 | 62 | def decorate(f): 63 | """Store f in HANDLERS[vcs][method].""" 64 | if vcs not in HANDLERS: 65 | HANDLERS[vcs] = {} 66 | HANDLERS[vcs][method] = f 67 | return f 68 | 69 | return decorate 70 | 71 | 72 | def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=None): 73 | """Call the given command(s).""" 74 | assert isinstance(commands, list) 75 | process = None 76 | for command in commands: 77 | try: 78 | dispcmd = str([command] + args) 79 | # remember shell=False, so use git.cmd on windows, not just git 80 | process = subprocess.Popen( 81 | [command] + args, 82 | cwd=cwd, 83 | env=env, 84 | stdout=subprocess.PIPE, 85 | stderr=(subprocess.PIPE if hide_stderr else None), 86 | ) 87 | break 88 | except OSError: 89 | e = sys.exc_info()[1] 90 | if e.errno == errno.ENOENT: 91 | continue 92 | if verbose: 93 | print("unable to run %s" % dispcmd) 94 | print(e) 95 | return None, None 96 | else: 97 | if verbose: 98 | print("unable to find command, tried %s" % (commands,)) 99 | return None, None 100 | stdout = process.communicate()[0].strip().decode() 101 | if process.returncode != 0: 102 | if verbose: 103 | print("unable to run %s (error)" % dispcmd) 104 | print("stdout was %s" % stdout) 105 | return None, process.returncode 106 | return stdout, process.returncode 107 | 108 | 109 | def versions_from_parentdir(parentdir_prefix, root, verbose): 110 | """Try to determine the version from the parent directory name. 111 | 112 | Source tarballs conventionally unpack into a directory that includes both 113 | the project name and a version string. We will also support searching up 114 | two directory levels for an appropriately named parent directory 115 | """ 116 | rootdirs = [] 117 | 118 | for _ in range(3): 119 | dirname = os.path.basename(root) 120 | if dirname.startswith(parentdir_prefix): 121 | return { 122 | "version": dirname[len(parentdir_prefix) :], 123 | "full-revisionid": None, 124 | "dirty": False, 125 | "error": None, 126 | "date": None, 127 | } 128 | rootdirs.append(root) 129 | root = os.path.dirname(root) # up a level 130 | 131 | if verbose: 132 | print( 133 | "Tried directories %s but none started with prefix %s" 134 | % (str(rootdirs), parentdir_prefix) 135 | ) 136 | raise NotThisMethod("rootdir doesn't start with parentdir_prefix") 137 | 138 | 139 | @register_vcs_handler("git", "get_keywords") 140 | def git_get_keywords(versionfile_abs): 141 | """Extract version information from the given file.""" 142 | # the code embedded in _version.py can just fetch the value of these 143 | # keywords. When used from setup.py, we don't want to import _version.py, 144 | # so we do it with a regexp instead. This function is not used from 145 | # _version.py. 146 | keywords = {} 147 | try: 148 | with open(versionfile_abs, "r") as fobj: 149 | for line in fobj: 150 | if line.strip().startswith("git_refnames ="): 151 | mo = re.search(r'=\s*"(.*)"', line) 152 | if mo: 153 | keywords["refnames"] = mo.group(1) 154 | if line.strip().startswith("git_full ="): 155 | mo = re.search(r'=\s*"(.*)"', line) 156 | if mo: 157 | keywords["full"] = mo.group(1) 158 | if line.strip().startswith("git_date ="): 159 | mo = re.search(r'=\s*"(.*)"', line) 160 | if mo: 161 | keywords["date"] = mo.group(1) 162 | except OSError: 163 | pass 164 | return keywords 165 | 166 | 167 | @register_vcs_handler("git", "keywords") 168 | def git_versions_from_keywords(keywords, tag_prefix, verbose): 169 | """Get version information from git keywords.""" 170 | if "refnames" not in keywords: 171 | raise NotThisMethod("Short version file found") 172 | date = keywords.get("date") 173 | if date is not None: 174 | # Use only the last line. Previous lines may contain GPG signature 175 | # information. 176 | date = date.splitlines()[-1] 177 | 178 | # git-2.2.0 added "%cI", which expands to an ISO-8601 -compliant 179 | # datestamp. However we prefer "%ci" (which expands to an "ISO-8601 180 | # -like" string, which we must then edit to make compliant), because 181 | # it's been around since git-1.5.3, and it's too difficult to 182 | # discover which version we're using, or to work around using an 183 | # older one. 184 | date = date.strip().replace(" ", "T", 1).replace(" ", "", 1) 185 | refnames = keywords["refnames"].strip() 186 | if refnames.startswith("$Format"): 187 | if verbose: 188 | print("keywords are unexpanded, not using") 189 | raise NotThisMethod("unexpanded keywords, not a git-archive tarball") 190 | refs = {r.strip() for r in refnames.strip("()").split(",")} 191 | # starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of 192 | # just "foo-1.0". If we see a "tag: " prefix, prefer those. 193 | TAG = "tag: " 194 | tags = {r[len(TAG) :] for r in refs if r.startswith(TAG)} 195 | if not tags: 196 | # Either we're using git < 1.8.3, or there really are no tags. We use 197 | # a heuristic: assume all version tags have a digit. The old git %d 198 | # expansion behaves like git log --decorate=short and strips out the 199 | # refs/heads/ and refs/tags/ prefixes that would let us distinguish 200 | # between branches and tags. By ignoring refnames without digits, we 201 | # filter out many common branch names like "release" and 202 | # "stabilization", as well as "HEAD" and "master". 203 | tags = {r for r in refs if re.search(r"\d", r)} 204 | if verbose: 205 | print("discarding '%s', no digits" % ",".join(refs - tags)) 206 | if verbose: 207 | print("likely tags: %s" % ",".join(sorted(tags))) 208 | for ref in sorted(tags): 209 | # sorting will prefer e.g. "2.0" over "2.0rc1" 210 | if ref.startswith(tag_prefix): 211 | r = ref[len(tag_prefix) :] 212 | # Filter out refs that exactly match prefix or that don't start 213 | # with a number once the prefix is stripped (mostly a concern 214 | # when prefix is '') 215 | if not re.match(r"\d", r): 216 | continue 217 | if verbose: 218 | print("picking %s" % r) 219 | return { 220 | "version": r, 221 | "full-revisionid": keywords["full"].strip(), 222 | "dirty": False, 223 | "error": None, 224 | "date": date, 225 | } 226 | # no suitable tags, so version is "0+unknown", but full hex is still there 227 | if verbose: 228 | print("no suitable tags, using unknown + full revision id") 229 | return { 230 | "version": "0+unknown", 231 | "full-revisionid": keywords["full"].strip(), 232 | "dirty": False, 233 | "error": "no suitable tags", 234 | "date": None, 235 | } 236 | 237 | 238 | @register_vcs_handler("git", "pieces_from_vcs") 239 | def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): 240 | """Get version from 'git describe' in the root of the source tree. 241 | 242 | This only gets called if the git-archive 'subst' keywords were *not* 243 | expanded, and _version.py hasn't already been rewritten with a short 244 | version string, meaning we're inside a checked out source tree. 245 | """ 246 | GITS = ["git"] 247 | TAG_PREFIX_REGEX = "*" 248 | if sys.platform == "win32": 249 | GITS = ["git.cmd", "git.exe"] 250 | TAG_PREFIX_REGEX = r"\*" 251 | 252 | _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=True) 253 | if rc != 0: 254 | if verbose: 255 | print("Directory %s not under git control" % root) 256 | raise NotThisMethod("'git rev-parse --git-dir' returned error") 257 | 258 | # if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty] 259 | # if there isn't one, this yields HEX[-dirty] (no NUM) 260 | describe_out, rc = runner( 261 | GITS, 262 | [ 263 | "describe", 264 | "--tags", 265 | "--dirty", 266 | "--always", 267 | "--long", 268 | "--match", 269 | "%s%s" % (tag_prefix, TAG_PREFIX_REGEX), 270 | ], 271 | cwd=root, 272 | ) 273 | # --long was added in git-1.5.5 274 | if describe_out is None: 275 | raise NotThisMethod("'git describe' failed") 276 | describe_out = describe_out.strip() 277 | full_out, rc = runner(GITS, ["rev-parse", "HEAD"], cwd=root) 278 | if full_out is None: 279 | raise NotThisMethod("'git rev-parse' failed") 280 | full_out = full_out.strip() 281 | 282 | pieces = {} 283 | pieces["long"] = full_out 284 | pieces["short"] = full_out[:7] # maybe improved later 285 | pieces["error"] = None 286 | 287 | branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"], cwd=root) 288 | # --abbrev-ref was added in git-1.6.3 289 | if rc != 0 or branch_name is None: 290 | raise NotThisMethod("'git rev-parse --abbrev-ref' returned error") 291 | branch_name = branch_name.strip() 292 | 293 | if branch_name == "HEAD": 294 | # If we aren't exactly on a branch, pick a branch which represents 295 | # the current commit. If all else fails, we are on a branchless 296 | # commit. 297 | branches, rc = runner(GITS, ["branch", "--contains"], cwd=root) 298 | # --contains was added in git-1.5.4 299 | if rc != 0 or branches is None: 300 | raise NotThisMethod("'git branch --contains' returned error") 301 | branches = branches.split("\n") 302 | 303 | # Remove the first line if we're running detached 304 | if "(" in branches[0]: 305 | branches.pop(0) 306 | 307 | # Strip off the leading "* " from the list of branches. 308 | branches = [branch[2:] for branch in branches] 309 | if "master" in branches: 310 | branch_name = "master" 311 | elif not branches: 312 | branch_name = None 313 | else: 314 | # Pick the first branch that is returned. Good or bad. 315 | branch_name = branches[0] 316 | 317 | pieces["branch"] = branch_name 318 | 319 | # parse describe_out. It will be like TAG-NUM-gHEX[-dirty] or HEX[-dirty] 320 | # TAG might have hyphens. 321 | git_describe = describe_out 322 | 323 | # look for -dirty suffix 324 | dirty = git_describe.endswith("-dirty") 325 | pieces["dirty"] = dirty 326 | if dirty: 327 | git_describe = git_describe[: git_describe.rindex("-dirty")] 328 | 329 | # now we have TAG-NUM-gHEX or HEX 330 | 331 | if "-" in git_describe: 332 | # TAG-NUM-gHEX 333 | mo = re.search(r"^(.+)-(\d+)-g([0-9a-f]+)$", git_describe) 334 | if not mo: 335 | # unparsable. Maybe git-describe is misbehaving? 336 | pieces["error"] = "unable to parse git-describe output: '%s'" % describe_out 337 | return pieces 338 | 339 | # tag 340 | full_tag = mo.group(1) 341 | if not full_tag.startswith(tag_prefix): 342 | if verbose: 343 | fmt = "tag '%s' doesn't start with prefix '%s'" 344 | print(fmt % (full_tag, tag_prefix)) 345 | pieces["error"] = "tag '%s' doesn't start with prefix '%s'" % ( 346 | full_tag, 347 | tag_prefix, 348 | ) 349 | return pieces 350 | pieces["closest-tag"] = full_tag[len(tag_prefix) :] 351 | 352 | # distance: number of commits since tag 353 | pieces["distance"] = int(mo.group(2)) 354 | 355 | # commit: short hex revision ID 356 | pieces["short"] = mo.group(3) 357 | 358 | else: 359 | # HEX: no tags 360 | pieces["closest-tag"] = None 361 | count_out, rc = runner(GITS, ["rev-list", "HEAD", "--count"], cwd=root) 362 | pieces["distance"] = int(count_out) # total number of commits 363 | 364 | # commit date: see ISO-8601 comment in git_versions_from_keywords() 365 | date = runner(GITS, ["show", "-s", "--format=%ci", "HEAD"], cwd=root)[0].strip() 366 | # Use only the last line. Previous lines may contain GPG signature 367 | # information. 368 | date = date.splitlines()[-1] 369 | pieces["date"] = date.strip().replace(" ", "T", 1).replace(" ", "", 1) 370 | 371 | return pieces 372 | 373 | 374 | def plus_or_dot(pieces): 375 | """Return a + if we don't already have one, else return a .""" 376 | if "+" in pieces.get("closest-tag", ""): 377 | return "." 378 | return "+" 379 | 380 | 381 | def render_pep440(pieces): 382 | """Build up version string, with post-release "local version identifier". 383 | 384 | Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you 385 | get a tagged build and then dirty it, you'll get TAG+0.gHEX.dirty 386 | 387 | Exceptions: 388 | 1: no tags. git_describe was just HEX. 0+untagged.DISTANCE.gHEX[.dirty] 389 | """ 390 | if pieces["closest-tag"]: 391 | rendered = pieces["closest-tag"] 392 | if pieces["distance"] or pieces["dirty"]: 393 | rendered += plus_or_dot(pieces) 394 | rendered += "%d.g%s" % (pieces["distance"], pieces["short"]) 395 | if pieces["dirty"]: 396 | rendered += ".dirty" 397 | else: 398 | # exception #1 399 | rendered = "0+untagged.%d.g%s" % (pieces["distance"], pieces["short"]) 400 | if pieces["dirty"]: 401 | rendered += ".dirty" 402 | return rendered 403 | 404 | 405 | def render_pep440_branch(pieces): 406 | """TAG[[.dev0]+DISTANCE.gHEX[.dirty]] . 407 | 408 | The ".dev0" means not master branch. Note that .dev0 sorts backwards 409 | (a feature branch will appear "older" than the master branch). 410 | 411 | Exceptions: 412 | 1: no tags. 0[.dev0]+untagged.DISTANCE.gHEX[.dirty] 413 | """ 414 | if pieces["closest-tag"]: 415 | rendered = pieces["closest-tag"] 416 | if pieces["distance"] or pieces["dirty"]: 417 | if pieces["branch"] != "master": 418 | rendered += ".dev0" 419 | rendered += plus_or_dot(pieces) 420 | rendered += "%d.g%s" % (pieces["distance"], pieces["short"]) 421 | if pieces["dirty"]: 422 | rendered += ".dirty" 423 | else: 424 | # exception #1 425 | rendered = "0" 426 | if pieces["branch"] != "master": 427 | rendered += ".dev0" 428 | rendered += "+untagged.%d.g%s" % (pieces["distance"], pieces["short"]) 429 | if pieces["dirty"]: 430 | rendered += ".dirty" 431 | return rendered 432 | 433 | 434 | def pep440_split_post(ver): 435 | """Split pep440 version string at the post-release segment. 436 | 437 | Returns the release segments before the post-release and the 438 | post-release version number (or -1 if no post-release segment is present). 439 | """ 440 | vc = str.split(ver, ".post") 441 | return vc[0], int(vc[1] or 0) if len(vc) == 2 else None 442 | 443 | 444 | def render_pep440_pre(pieces): 445 | """TAG[.postN.devDISTANCE] -- No -dirty. 446 | 447 | Exceptions: 448 | 1: no tags. 0.post0.devDISTANCE 449 | """ 450 | if pieces["closest-tag"]: 451 | if pieces["distance"]: 452 | # update the post release segment 453 | tag_version, post_version = pep440_split_post(pieces["closest-tag"]) 454 | rendered = tag_version 455 | if post_version is not None: 456 | rendered += ".post%d.dev%d" % (post_version + 1, pieces["distance"]) 457 | else: 458 | rendered += ".post0.dev%d" % (pieces["distance"]) 459 | else: 460 | # no commits, use the tag as the version 461 | rendered = pieces["closest-tag"] 462 | else: 463 | # exception #1 464 | rendered = "0.post0.dev%d" % pieces["distance"] 465 | return rendered 466 | 467 | 468 | def render_pep440_post(pieces): 469 | """TAG[.postDISTANCE[.dev0]+gHEX] . 470 | 471 | The ".dev0" means dirty. Note that .dev0 sorts backwards 472 | (a dirty tree will appear "older" than the corresponding clean one), 473 | but you shouldn't be releasing software with -dirty anyways. 474 | 475 | Exceptions: 476 | 1: no tags. 0.postDISTANCE[.dev0] 477 | """ 478 | if pieces["closest-tag"]: 479 | rendered = pieces["closest-tag"] 480 | if pieces["distance"] or pieces["dirty"]: 481 | rendered += ".post%d" % pieces["distance"] 482 | if pieces["dirty"]: 483 | rendered += ".dev0" 484 | rendered += plus_or_dot(pieces) 485 | rendered += "g%s" % pieces["short"] 486 | else: 487 | # exception #1 488 | rendered = "0.post%d" % pieces["distance"] 489 | if pieces["dirty"]: 490 | rendered += ".dev0" 491 | rendered += "+g%s" % pieces["short"] 492 | return rendered 493 | 494 | 495 | def render_pep440_post_branch(pieces): 496 | """TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] . 497 | 498 | The ".dev0" means not master branch. 499 | 500 | Exceptions: 501 | 1: no tags. 0.postDISTANCE[.dev0]+gHEX[.dirty] 502 | """ 503 | if pieces["closest-tag"]: 504 | rendered = pieces["closest-tag"] 505 | if pieces["distance"] or pieces["dirty"]: 506 | rendered += ".post%d" % pieces["distance"] 507 | if pieces["branch"] != "master": 508 | rendered += ".dev0" 509 | rendered += plus_or_dot(pieces) 510 | rendered += "g%s" % pieces["short"] 511 | if pieces["dirty"]: 512 | rendered += ".dirty" 513 | else: 514 | # exception #1 515 | rendered = "0.post%d" % pieces["distance"] 516 | if pieces["branch"] != "master": 517 | rendered += ".dev0" 518 | rendered += "+g%s" % pieces["short"] 519 | if pieces["dirty"]: 520 | rendered += ".dirty" 521 | return rendered 522 | 523 | 524 | def render_pep440_old(pieces): 525 | """TAG[.postDISTANCE[.dev0]] . 526 | 527 | The ".dev0" means dirty. 528 | 529 | Exceptions: 530 | 1: no tags. 0.postDISTANCE[.dev0] 531 | """ 532 | if pieces["closest-tag"]: 533 | rendered = pieces["closest-tag"] 534 | if pieces["distance"] or pieces["dirty"]: 535 | rendered += ".post%d" % pieces["distance"] 536 | if pieces["dirty"]: 537 | rendered += ".dev0" 538 | else: 539 | # exception #1 540 | rendered = "0.post%d" % pieces["distance"] 541 | if pieces["dirty"]: 542 | rendered += ".dev0" 543 | return rendered 544 | 545 | 546 | def render_git_describe(pieces): 547 | """TAG[-DISTANCE-gHEX][-dirty]. 548 | 549 | Like 'git describe --tags --dirty --always'. 550 | 551 | Exceptions: 552 | 1: no tags. HEX[-dirty] (note: no 'g' prefix) 553 | """ 554 | if pieces["closest-tag"]: 555 | rendered = pieces["closest-tag"] 556 | if pieces["distance"]: 557 | rendered += "-%d-g%s" % (pieces["distance"], pieces["short"]) 558 | else: 559 | # exception #1 560 | rendered = pieces["short"] 561 | if pieces["dirty"]: 562 | rendered += "-dirty" 563 | return rendered 564 | 565 | 566 | def render_git_describe_long(pieces): 567 | """TAG-DISTANCE-gHEX[-dirty]. 568 | 569 | Like 'git describe --tags --dirty --always -long'. 570 | The distance/hash is unconditional. 571 | 572 | Exceptions: 573 | 1: no tags. HEX[-dirty] (note: no 'g' prefix) 574 | """ 575 | if pieces["closest-tag"]: 576 | rendered = pieces["closest-tag"] 577 | rendered += "-%d-g%s" % (pieces["distance"], pieces["short"]) 578 | else: 579 | # exception #1 580 | rendered = pieces["short"] 581 | if pieces["dirty"]: 582 | rendered += "-dirty" 583 | return rendered 584 | 585 | 586 | def render(pieces, style): 587 | """Render the given version pieces into the requested style.""" 588 | if pieces["error"]: 589 | return { 590 | "version": "unknown", 591 | "full-revisionid": pieces.get("long"), 592 | "dirty": None, 593 | "error": pieces["error"], 594 | "date": None, 595 | } 596 | 597 | if not style or style == "default": 598 | style = "pep440" # the default 599 | 600 | if style == "pep440": 601 | rendered = render_pep440(pieces) 602 | elif style == "pep440-branch": 603 | rendered = render_pep440_branch(pieces) 604 | elif style == "pep440-pre": 605 | rendered = render_pep440_pre(pieces) 606 | elif style == "pep440-post": 607 | rendered = render_pep440_post(pieces) 608 | elif style == "pep440-post-branch": 609 | rendered = render_pep440_post_branch(pieces) 610 | elif style == "pep440-old": 611 | rendered = render_pep440_old(pieces) 612 | elif style == "git-describe": 613 | rendered = render_git_describe(pieces) 614 | elif style == "git-describe-long": 615 | rendered = render_git_describe_long(pieces) 616 | else: 617 | raise ValueError("unknown style '%s'" % style) 618 | 619 | return { 620 | "version": rendered, 621 | "full-revisionid": pieces["long"], 622 | "dirty": pieces["dirty"], 623 | "error": None, 624 | "date": pieces.get("date"), 625 | } 626 | 627 | 628 | def get_versions(): 629 | """Get version information or return default if unable to do so.""" 630 | # I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have 631 | # __file__, we can work backwards from there to the root. Some 632 | # py2exe/bbfreeze/non-CPython implementations don't do __file__, in which 633 | # case we can only use expanded keywords. 634 | 635 | cfg = get_config() 636 | verbose = cfg.verbose 637 | 638 | try: 639 | return git_versions_from_keywords(get_keywords(), cfg.tag_prefix, verbose) 640 | except NotThisMethod: 641 | pass 642 | 643 | try: 644 | root = os.path.realpath(__file__) 645 | # versionfile_source is the relative path from the top of the source 646 | # tree (where the .git directory might live) to this file. Invert 647 | # this to find the root from __file__. 648 | for _ in cfg.versionfile_source.split("/"): 649 | root = os.path.dirname(root) 650 | except NameError: 651 | return { 652 | "version": "0+unknown", 653 | "full-revisionid": None, 654 | "dirty": None, 655 | "error": "unable to find root of source tree", 656 | "date": None, 657 | } 658 | 659 | try: 660 | pieces = git_pieces_from_vcs(cfg.tag_prefix, root, verbose) 661 | return render(pieces, cfg.style) 662 | except NotThisMethod: 663 | pass 664 | 665 | try: 666 | if cfg.parentdir_prefix: 667 | return versions_from_parentdir(cfg.parentdir_prefix, root, verbose) 668 | except NotThisMethod: 669 | pass 670 | 671 | return { 672 | "version": "0+unknown", 673 | "full-revisionid": None, 674 | "dirty": None, 675 | "error": "unable to compute version", 676 | "date": None, 677 | } 678 | -------------------------------------------------------------------------------- /prefect_snowflake/database.py: -------------------------------------------------------------------------------- 1 | """Module for querying against Snowflake databases.""" 2 | 3 | import asyncio 4 | from typing import Any, Dict, List, Optional, Tuple, Type, Union 5 | 6 | from prefect import task 7 | from prefect.blocks.abstract import DatabaseBlock 8 | from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible 9 | from prefect.utilities.hashing import hash_objects 10 | from pydantic import VERSION as PYDANTIC_VERSION 11 | 12 | if PYDANTIC_VERSION.startswith("2."): 13 | from pydantic.v1 import Field 14 | else: 15 | from pydantic import Field 16 | 17 | from snowflake.connector.connection import SnowflakeConnection 18 | from snowflake.connector.cursor import SnowflakeCursor 19 | 20 | from prefect_snowflake import SnowflakeCredentials 21 | 22 | BEGIN_TRANSACTION_STATEMENT = "BEGIN TRANSACTION" 23 | END_TRANSACTION_STATEMENT = "COMMIT" 24 | 25 | 26 | class SnowflakeConnector(DatabaseBlock): 27 | 28 | """ 29 | Block used to manage connections with Snowflake. 30 | 31 | Upon instantiating, a connection is created and maintained for the life of 32 | the object until the close method is called. 33 | 34 | It is recommended to use this block as a context manager, which will automatically 35 | close the engine and its connections when the context is exited. 36 | 37 | It is also recommended that this block is loaded and consumed within a single task 38 | or flow because if the block is passed across separate tasks and flows, 39 | the state of the block's connection and cursor will be lost. 40 | 41 | Args: 42 | credentials: The credentials to authenticate with Snowflake. 43 | database: The name of the default database to use. 44 | warehouse: The name of the default warehouse to use. 45 | schema: The name of the default schema to use; 46 | this attribute is accessible through `SnowflakeConnector(...).schema_`. 47 | fetch_size: The number of rows to fetch at a time. 48 | poll_frequency_s: The number of seconds before checking query. 49 | 50 | Examples: 51 | Load stored Snowflake connector as a context manager: 52 | ```python 53 | from prefect_snowflake.database import SnowflakeConnector 54 | 55 | snowflake_connector = SnowflakeConnector.load("BLOCK_NAME"): 56 | ``` 57 | 58 | Insert data into database and fetch results. 59 | ```python 60 | from prefect_snowflake.database import SnowflakeConnector 61 | 62 | with SnowflakeConnector.load("BLOCK_NAME") as conn: 63 | conn.execute( 64 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 65 | ) 66 | conn.execute_many( 67 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 68 | seq_of_parameters=[ 69 | {"name": "Ford", "address": "Highway 42"}, 70 | {"name": "Unknown", "address": "Space"}, 71 | {"name": "Me", "address": "Myway 88"}, 72 | ], 73 | ) 74 | results = conn.fetch_all( 75 | "SELECT * FROM customers WHERE address = %(address)s", 76 | parameters={"address": "Space"} 77 | ) 78 | print(results) 79 | ``` 80 | """ # noqa 81 | 82 | _block_type_name = "Snowflake Connector" 83 | _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/bd359de0b4be76c2254bd329fe3a267a1a3879c2-250x250.png" # noqa 84 | _documentation_url = "https://prefecthq.github.io/prefect-snowflake/database/#prefect_snowflake.database.SnowflakeConnector" # noqa 85 | _description = "Perform data operations against a Snowflake database." 86 | 87 | credentials: SnowflakeCredentials = Field( 88 | default=..., description="The credentials to authenticate with Snowflake." 89 | ) 90 | database: str = Field( 91 | default=..., description="The name of the default database to use." 92 | ) 93 | warehouse: str = Field( 94 | default=..., description="The name of the default warehouse to use." 95 | ) 96 | schema_: str = Field( 97 | default=..., 98 | alias="schema", 99 | description="The name of the default schema to use.", 100 | ) 101 | fetch_size: int = Field( 102 | default=1, description="The default number of rows to fetch at a time." 103 | ) 104 | poll_frequency_s: int = Field( 105 | default=1, 106 | title="Poll Frequency [seconds]", 107 | description=( 108 | "The number of seconds between checking query " 109 | "status for long running queries." 110 | ), 111 | ) 112 | 113 | _connection: Optional[SnowflakeConnection] = None 114 | _unique_cursors: Dict[str, SnowflakeCursor] = None 115 | 116 | def get_connection(self, **connect_kwargs: Any) -> SnowflakeConnection: 117 | """ 118 | Returns an authenticated connection that can be 119 | used to query from Snowflake databases. 120 | 121 | Args: 122 | **connect_kwargs: Additional arguments to pass to 123 | `snowflake.connector.connect`. 124 | 125 | Returns: 126 | The authenticated SnowflakeConnection. 127 | 128 | Examples: 129 | ```python 130 | from prefect_snowflake.credentials import SnowflakeCredentials 131 | from prefect_snowflake.database import SnowflakeConnector 132 | 133 | snowflake_credentials = SnowflakeCredentials( 134 | account="account", 135 | user="user", 136 | password="password", 137 | ) 138 | snowflake_connector = SnowflakeConnector( 139 | database="database", 140 | warehouse="warehouse", 141 | schema="schema", 142 | credentials=snowflake_credentials 143 | ) 144 | with snowflake_connector.get_connection() as connection: 145 | ... 146 | ``` 147 | """ 148 | if self._connection is not None: 149 | return self._connection 150 | 151 | connect_params = { 152 | "database": self.database, 153 | "warehouse": self.warehouse, 154 | "schema": self.schema_, 155 | } 156 | connection = self.credentials.get_client(**connect_kwargs, **connect_params) 157 | self._connection = connection 158 | self.logger.info("Started a new connection to Snowflake.") 159 | return connection 160 | 161 | def _start_connection(self): 162 | """ 163 | Starts Snowflake database connection. 164 | """ 165 | self.get_connection() 166 | if self._unique_cursors is None: 167 | self._unique_cursors = {} 168 | 169 | def _get_cursor( 170 | self, 171 | inputs: Dict[str, Any], 172 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 173 | ) -> Tuple[bool, SnowflakeCursor]: 174 | """ 175 | Get a Snowflake cursor. 176 | 177 | Args: 178 | inputs: The inputs to generate a unique hash, used to decide 179 | whether a new cursor should be used. 180 | cursor_type: The class of the cursor to use when creating a 181 | Snowflake cursor. 182 | 183 | Returns: 184 | Whether a cursor is new and a Snowflake cursor. 185 | """ 186 | self._start_connection() 187 | 188 | input_hash = hash_objects(inputs) 189 | if input_hash is None: 190 | raise RuntimeError( 191 | "We were not able to hash your inputs, " 192 | "which resulted in an unexpected data return; " 193 | "please open an issue with a reproducible example." 194 | ) 195 | if input_hash not in self._unique_cursors.keys(): 196 | new_cursor = self._connection.cursor(cursor_type) 197 | self._unique_cursors[input_hash] = new_cursor 198 | return True, new_cursor 199 | else: 200 | existing_cursor = self._unique_cursors[input_hash] 201 | return False, existing_cursor 202 | 203 | async def _execute_async(self, cursor: SnowflakeCursor, inputs: Dict[str, Any]): 204 | """Helper method to execute operations asynchronously.""" 205 | response = await run_sync_in_worker_thread(cursor.execute_async, **inputs) 206 | self.logger.info( 207 | f"Executing the operation, {inputs['command']!r}, asynchronously; " 208 | f"polling for the result every {self.poll_frequency_s} seconds." 209 | ) 210 | 211 | query_id = response["queryId"] 212 | while self._connection.is_still_running( 213 | await run_sync_in_worker_thread( 214 | self._connection.get_query_status_throw_if_error, query_id 215 | ) 216 | ): 217 | await asyncio.sleep(self.poll_frequency_s) 218 | await run_sync_in_worker_thread(cursor.get_results_from_sfqid, query_id) 219 | 220 | def reset_cursors(self) -> None: 221 | """ 222 | Tries to close all opened cursors. 223 | 224 | Examples: 225 | Reset the cursors to refresh cursor position. 226 | ```python 227 | from prefect_snowflake.database import SnowflakeConnector 228 | 229 | with SnowflakeConnector.load("BLOCK_NAME") as conn: 230 | conn.execute( 231 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 232 | ) 233 | conn.execute_many( 234 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 235 | seq_of_parameters=[ 236 | {"name": "Ford", "address": "Highway 42"}, 237 | {"name": "Unknown", "address": "Space"}, 238 | {"name": "Me", "address": "Myway 88"}, 239 | ], 240 | ) 241 | print(conn.fetch_one("SELECT * FROM customers")) # Ford 242 | conn.reset_cursors() 243 | print(conn.fetch_one("SELECT * FROM customers")) # should be Ford again 244 | ``` 245 | """ # noqa 246 | if not self._unique_cursors: 247 | self.logger.info("There were no cursors to reset.") 248 | return 249 | 250 | input_hashes = tuple(self._unique_cursors.keys()) 251 | for input_hash in input_hashes: 252 | cursor = self._unique_cursors.pop(input_hash) 253 | try: 254 | cursor.close() 255 | except Exception as exc: 256 | self.logger.warning( 257 | f"Failed to close cursor for input hash {input_hash!r}: {exc}" 258 | ) 259 | self.logger.info("Successfully reset the cursors.") 260 | 261 | @sync_compatible 262 | async def fetch_one( 263 | self, 264 | operation: str, 265 | parameters: Optional[Dict[str, Any]] = None, 266 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 267 | **execute_kwargs: Any, 268 | ) -> Tuple[Any]: 269 | """ 270 | Fetch a single result from the database. 271 | Repeated calls using the same inputs to *any* of the fetch methods of this 272 | block will skip executing the operation again, and instead, 273 | return the next set of results from the previous execution, 274 | until the reset_cursors method is called. 275 | 276 | Args: 277 | operation: The SQL query or other operation to be executed. 278 | parameters: The parameters for the operation. 279 | cursor_type: The class of the cursor to use when creating a Snowflake cursor. 280 | **execute_kwargs: Additional options to pass to `cursor.execute_async`. 281 | 282 | Returns: 283 | A tuple containing the data returned by the database, 284 | where each row is a tuple and each column is a value in the tuple. 285 | 286 | Examples: 287 | Fetch one row from the database where address is Space. 288 | ```python 289 | from prefect_snowflake.database import SnowflakeConnector 290 | 291 | with SnowflakeConnector.load("BLOCK_NAME") as conn: 292 | conn.execute( 293 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 294 | ) 295 | conn.execute_many( 296 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 297 | seq_of_parameters=[ 298 | {"name": "Ford", "address": "Highway 42"}, 299 | {"name": "Unknown", "address": "Space"}, 300 | {"name": "Me", "address": "Myway 88"}, 301 | ], 302 | ) 303 | result = conn.fetch_one( 304 | "SELECT * FROM customers WHERE address = %(address)s", 305 | parameters={"address": "Space"} 306 | ) 307 | print(result) 308 | ``` 309 | """ # noqa 310 | inputs = dict( 311 | command=operation, 312 | params=parameters, 313 | **execute_kwargs, 314 | ) 315 | new, cursor = self._get_cursor(inputs, cursor_type=cursor_type) 316 | if new: 317 | await self._execute_async(cursor, inputs) 318 | self.logger.debug("Preparing to fetch a row.") 319 | result = await run_sync_in_worker_thread(cursor.fetchone) 320 | return result 321 | 322 | @sync_compatible 323 | async def fetch_many( 324 | self, 325 | operation: str, 326 | parameters: Optional[Dict[str, Any]] = None, 327 | size: Optional[int] = None, 328 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 329 | **execute_kwargs: Any, 330 | ) -> List[Tuple[Any]]: 331 | """ 332 | Fetch a limited number of results from the database. 333 | Repeated calls using the same inputs to *any* of the fetch methods of this 334 | block will skip executing the operation again, and instead, 335 | return the next set of results from the previous execution, 336 | until the reset_cursors method is called. 337 | 338 | Args: 339 | operation: The SQL query or other operation to be executed. 340 | parameters: The parameters for the operation. 341 | size: The number of results to return; if None or 0, uses the value of 342 | `fetch_size` configured on the block. 343 | cursor_type: The class of the cursor to use when creating a Snowflake cursor. 344 | **execute_kwargs: Additional options to pass to `cursor.execute_async`. 345 | 346 | Returns: 347 | A list of tuples containing the data returned by the database, 348 | where each row is a tuple and each column is a value in the tuple. 349 | 350 | Examples: 351 | Repeatedly fetch two rows from the database where address is Highway 42. 352 | ```python 353 | from prefect_snowflake.database import SnowflakeConnector 354 | 355 | with SnowflakeConnector.load("BLOCK_NAME") as conn: 356 | conn.execute( 357 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 358 | ) 359 | conn.execute_many( 360 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 361 | seq_of_parameters=[ 362 | {"name": "Marvin", "address": "Highway 42"}, 363 | {"name": "Ford", "address": "Highway 42"}, 364 | {"name": "Unknown", "address": "Highway 42"}, 365 | {"name": "Me", "address": "Highway 42"}, 366 | ], 367 | ) 368 | result = conn.fetch_many( 369 | "SELECT * FROM customers WHERE address = %(address)s", 370 | parameters={"address": "Highway 42"}, 371 | size=2 372 | ) 373 | print(result) # Marvin, Ford 374 | result = conn.fetch_many( 375 | "SELECT * FROM customers WHERE address = %(address)s", 376 | parameters={"address": "Highway 42"}, 377 | size=2 378 | ) 379 | print(result) # Unknown, Me 380 | ``` 381 | """ # noqa 382 | inputs = dict( 383 | command=operation, 384 | params=parameters, 385 | **execute_kwargs, 386 | ) 387 | new, cursor = self._get_cursor(inputs, cursor_type) 388 | if new: 389 | await self._execute_async(cursor, inputs) 390 | size = size or self.fetch_size 391 | self.logger.debug(f"Preparing to fetch {size} rows.") 392 | result = await run_sync_in_worker_thread(cursor.fetchmany, size=size) 393 | return result 394 | 395 | @sync_compatible 396 | async def fetch_all( 397 | self, 398 | operation: str, 399 | parameters: Optional[Dict[str, Any]] = None, 400 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 401 | **execute_kwargs: Any, 402 | ) -> List[Tuple[Any]]: 403 | """ 404 | Fetch all results from the database. 405 | Repeated calls using the same inputs to *any* of the fetch methods of this 406 | block will skip executing the operation again, and instead, 407 | return the next set of results from the previous execution, 408 | until the reset_cursors method is called. 409 | 410 | Args: 411 | operation: The SQL query or other operation to be executed. 412 | parameters: The parameters for the operation. 413 | cursor_type: The class of the cursor to use when creating a Snowflake cursor. 414 | **execute_kwargs: Additional options to pass to `cursor.execute_async`. 415 | 416 | Returns: 417 | A list of tuples containing the data returned by the database, 418 | where each row is a tuple and each column is a value in the tuple. 419 | 420 | Examples: 421 | Fetch all rows from the database where address is Highway 42. 422 | ```python 423 | from prefect_snowflake.database import SnowflakeConnector 424 | 425 | with SnowflakeConnector.load("BLOCK_NAME") as conn: 426 | conn.execute( 427 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 428 | ) 429 | conn.execute_many( 430 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 431 | seq_of_parameters=[ 432 | {"name": "Marvin", "address": "Highway 42"}, 433 | {"name": "Ford", "address": "Highway 42"}, 434 | {"name": "Unknown", "address": "Highway 42"}, 435 | {"name": "Me", "address": "Myway 88"}, 436 | ], 437 | ) 438 | result = conn.fetch_all( 439 | "SELECT * FROM customers WHERE address = %(address)s", 440 | parameters={"address": "Highway 42"}, 441 | ) 442 | print(result) # Marvin, Ford, Unknown 443 | ``` 444 | """ # noqa 445 | inputs = dict( 446 | command=operation, 447 | params=parameters, 448 | **execute_kwargs, 449 | ) 450 | new, cursor = self._get_cursor(inputs, cursor_type) 451 | if new: 452 | await self._execute_async(cursor, inputs) 453 | self.logger.debug("Preparing to fetch all rows.") 454 | result = await run_sync_in_worker_thread(cursor.fetchall) 455 | return result 456 | 457 | @sync_compatible 458 | async def execute( 459 | self, 460 | operation: str, 461 | parameters: Optional[Dict[str, Any]] = None, 462 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 463 | **execute_kwargs: Any, 464 | ) -> None: 465 | """ 466 | Executes an operation on the database. This method is intended to be used 467 | for operations that do not return data, such as INSERT, UPDATE, or DELETE. 468 | Unlike the fetch methods, this method will always execute the operation 469 | upon calling. 470 | 471 | Args: 472 | operation: The SQL query or other operation to be executed. 473 | parameters: The parameters for the operation. 474 | cursor_type: The class of the cursor to use when creating a Snowflake cursor. 475 | **execute_kwargs: Additional options to pass to `cursor.execute_async`. 476 | 477 | Examples: 478 | Create table named customers with two columns, name and address. 479 | ```python 480 | from prefect_snowflake.database import SnowflakeConnector 481 | 482 | with SnowflakeConnector.load("BLOCK_NAME") as conn: 483 | conn.execute( 484 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 485 | ) 486 | ``` 487 | """ # noqa 488 | self._start_connection() 489 | 490 | inputs = dict( 491 | command=operation, 492 | params=parameters, 493 | **execute_kwargs, 494 | ) 495 | with self._connection.cursor(cursor_type) as cursor: 496 | await run_sync_in_worker_thread(cursor.execute, **inputs) 497 | self.logger.info(f"Executed the operation, {operation!r}.") 498 | 499 | @sync_compatible 500 | async def execute_many( 501 | self, 502 | operation: str, 503 | seq_of_parameters: List[Dict[str, Any]], 504 | ) -> None: 505 | """ 506 | Executes many operations on the database. This method is intended to be used 507 | for operations that do not return data, such as INSERT, UPDATE, or DELETE. 508 | Unlike the fetch methods, this method will always execute the operations 509 | upon calling. 510 | 511 | Args: 512 | operation: The SQL query or other operation to be executed. 513 | seq_of_parameters: The sequence of parameters for the operation. 514 | 515 | Examples: 516 | Create table and insert three rows into it. 517 | ```python 518 | from prefect_snowflake.database import SnowflakeConnector 519 | 520 | with SnowflakeConnector.load("BLOCK_NAME") as conn: 521 | conn.execute( 522 | "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);" 523 | ) 524 | conn.execute_many( 525 | "INSERT INTO customers (name, address) VALUES (%(name)s, %(address)s);", 526 | seq_of_parameters=[ 527 | {"name": "Marvin", "address": "Highway 42"}, 528 | {"name": "Ford", "address": "Highway 42"}, 529 | {"name": "Unknown", "address": "Space"}, 530 | ], 531 | ) 532 | ``` 533 | """ # noqa 534 | self._start_connection() 535 | 536 | inputs = dict( 537 | command=operation, 538 | seqparams=seq_of_parameters, 539 | ) 540 | with self._connection.cursor() as cursor: 541 | await run_sync_in_worker_thread(cursor.executemany, **inputs) 542 | self.logger.info( 543 | f"Executed {len(seq_of_parameters)} operations off {operation!r}." 544 | ) 545 | 546 | def close(self): 547 | """ 548 | Closes connection and its cursors. 549 | """ 550 | try: 551 | self.reset_cursors() 552 | finally: 553 | if self._connection is None: 554 | self.logger.info("There was no connection open to be closed.") 555 | return 556 | self._connection.close() 557 | self._connection = None 558 | self.logger.info("Successfully closed the Snowflake connection.") 559 | 560 | def __enter__(self): 561 | """ 562 | Start a connection upon entry. 563 | """ 564 | return self 565 | 566 | def __exit__(self, *args): 567 | """ 568 | Closes connection and its cursors upon exit. 569 | """ 570 | self.close() 571 | 572 | def __getstate__(self): 573 | """Allows block to be pickled and dumped.""" 574 | data = self.__dict__.copy() 575 | data.update({k: None for k in {"_connection", "_unique_cursors"}}) 576 | return data 577 | 578 | def __setstate__(self, data: dict): 579 | """Reset connection and cursors upon loading.""" 580 | self.__dict__.update(data) 581 | self._start_connection() 582 | 583 | 584 | @task 585 | async def snowflake_query( 586 | query: str, 587 | snowflake_connector: SnowflakeConnector, 588 | params: Union[Tuple[Any], Dict[str, Any]] = None, 589 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 590 | poll_frequency_seconds: int = 1, 591 | ) -> List[Tuple[Any]]: 592 | """ 593 | Executes a query against a Snowflake database. 594 | 595 | Args: 596 | query: The query to execute against the database. 597 | params: The params to replace the placeholders in the query. 598 | snowflake_connector: The credentials to use to authenticate. 599 | cursor_type: The type of database cursor to use for the query. 600 | poll_frequency_seconds: Number of seconds to wait in between checks for 601 | run completion. 602 | 603 | Returns: 604 | The output of `response.fetchall()`. 605 | 606 | Examples: 607 | Query Snowflake table with the ID value parameterized. 608 | ```python 609 | from prefect import flow 610 | from prefect_snowflake.credentials import SnowflakeCredentials 611 | from prefect_snowflake.database import SnowflakeConnector, snowflake_query 612 | 613 | 614 | @flow 615 | def snowflake_query_flow(): 616 | snowflake_credentials = SnowflakeCredentials( 617 | account="account", 618 | user="user", 619 | password="password", 620 | ) 621 | snowflake_connector = SnowflakeConnector( 622 | database="database", 623 | warehouse="warehouse", 624 | schema="schema", 625 | credentials=snowflake_credentials 626 | ) 627 | result = snowflake_query( 628 | "SELECT * FROM table WHERE id=%{id_param}s LIMIT 8;", 629 | snowflake_connector, 630 | params={"id_param": 1} 631 | ) 632 | return result 633 | 634 | snowflake_query_flow() 635 | ``` 636 | """ 637 | # context manager automatically rolls back failed transactions and closes 638 | with snowflake_connector.get_connection() as connection: 639 | with connection.cursor(cursor_type) as cursor: 640 | response = cursor.execute_async(query, params=params) 641 | query_id = response["queryId"] 642 | while connection.is_still_running( 643 | connection.get_query_status_throw_if_error(query_id) 644 | ): 645 | await asyncio.sleep(poll_frequency_seconds) 646 | cursor.get_results_from_sfqid(query_id) 647 | result = cursor.fetchall() 648 | return result 649 | 650 | 651 | @task 652 | async def snowflake_multiquery( 653 | queries: List[str], 654 | snowflake_connector: SnowflakeConnector, 655 | params: Union[Tuple[Any], Dict[str, Any]] = None, 656 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 657 | as_transaction: bool = False, 658 | return_transaction_control_results: bool = False, 659 | poll_frequency_seconds: int = 1, 660 | ) -> List[List[Tuple[Any]]]: 661 | """ 662 | Executes multiple queries against a Snowflake database in a shared session. 663 | Allows execution in a transaction. 664 | 665 | Args: 666 | queries: The list of queries to execute against the database. 667 | params: The params to replace the placeholders in the query. 668 | snowflake_connector: The credentials to use to authenticate. 669 | cursor_type: The type of database cursor to use for the query. 670 | as_transaction: If True, queries are executed in a transaction. 671 | return_transaction_control_results: Determines if the results of queries 672 | controlling the transaction (BEGIN/COMMIT) should be returned. 673 | poll_frequency_seconds: Number of seconds to wait in between checks for 674 | run completion. 675 | 676 | Returns: 677 | List of the outputs of `response.fetchall()` for each query. 678 | 679 | Examples: 680 | Query Snowflake table with the ID value parameterized. 681 | ```python 682 | from prefect import flow 683 | from prefect_snowflake.credentials import SnowflakeCredentials 684 | from prefect_snowflake.database import SnowflakeConnector, snowflake_multiquery 685 | 686 | 687 | @flow 688 | def snowflake_multiquery_flow(): 689 | snowflake_credentials = SnowflakeCredentials( 690 | account="account", 691 | user="user", 692 | password="password", 693 | ) 694 | snowflake_connector = SnowflakeConnector( 695 | database="database", 696 | warehouse="warehouse", 697 | schema="schema", 698 | credentials=snowflake_credentials 699 | ) 700 | result = snowflake_multiquery( 701 | ["SELECT * FROM table WHERE id=%{id_param}s LIMIT 8;", "SELECT 1,2"], 702 | snowflake_connector, 703 | params={"id_param": 1}, 704 | as_transaction=True 705 | ) 706 | return result 707 | 708 | snowflake_multiquery_flow() 709 | ``` 710 | """ 711 | with snowflake_connector.get_connection() as connection: 712 | if as_transaction: 713 | queries.insert(0, BEGIN_TRANSACTION_STATEMENT) 714 | queries.append(END_TRANSACTION_STATEMENT) 715 | 716 | with connection.cursor(cursor_type) as cursor: 717 | results = [] 718 | for query in queries: 719 | response = cursor.execute_async(query, params=params) 720 | query_id = response["queryId"] 721 | while connection.is_still_running( 722 | connection.get_query_status_throw_if_error(query_id) 723 | ): 724 | await asyncio.sleep(poll_frequency_seconds) 725 | cursor.get_results_from_sfqid(query_id) 726 | result = cursor.fetchall() 727 | results.append(result) 728 | 729 | # cut off results from BEGIN/COMMIT queries 730 | if as_transaction and not return_transaction_control_results: 731 | return results[1:-1] 732 | else: 733 | return results 734 | 735 | 736 | @task 737 | async def snowflake_query_sync( 738 | query: str, 739 | snowflake_connector: SnowflakeConnector, 740 | params: Union[Tuple[Any], Dict[str, Any]] = None, 741 | cursor_type: Type[SnowflakeCursor] = SnowflakeCursor, 742 | ) -> List[Tuple[Any]]: 743 | """ 744 | Executes a query in sync mode against a Snowflake database. 745 | 746 | Args: 747 | query: The query to execute against the database. 748 | params: The params to replace the placeholders in the query. 749 | snowflake_connector: The credentials to use to authenticate. 750 | cursor_type: The type of database cursor to use for the query. 751 | 752 | Returns: 753 | The output of `response.fetchall()`. 754 | 755 | Examples: 756 | Execute a put statement. 757 | ```python 758 | from prefect import flow 759 | from prefect_snowflake.credentials import SnowflakeCredentials 760 | from prefect_snowflake.database import SnowflakeConnector, snowflake_query 761 | 762 | 763 | @flow 764 | def snowflake_query_sync_flow(): 765 | snowflake_credentials = SnowflakeCredentials( 766 | account="account", 767 | user="user", 768 | password="password", 769 | ) 770 | snowflake_connector = SnowflakeConnector( 771 | database="database", 772 | warehouse="warehouse", 773 | schema="schema", 774 | credentials=snowflake_credentials 775 | ) 776 | result = snowflake_query_sync( 777 | "put file://afile.csv @mystage;", 778 | snowflake_connector, 779 | ) 780 | return result 781 | 782 | snowflake_query_sync_flow() 783 | ``` 784 | """ 785 | # context manager automatically rolls back failed transactions and closes 786 | with snowflake_connector.get_connection() as connection: 787 | with connection.cursor(cursor_type) as cursor: 788 | cursor.execute(query, params=params) 789 | result = cursor.fetchall() 790 | return result 791 | --------------------------------------------------------------------------------