Skip to content

Commit

Permalink
Add 'Firebolt' to Prefect task library. (PrefectHQ#5265)
Browse files Browse the repository at this point in the history
* Firebolt prefect tasks (PrefectHQ#1)

* Allow `Azure` storage to overwrite existing blobs (PrefectHQ#5103)

* Allow `Azure` storage to overwrite existing blobs

* Move flaky test mark to class level

* Add test covering `overwrite`

* Spell correctly

* Fixup assertion

* Use mock.ANY

* Today is not a day for spelling correctly

* Add changes entry

* Glob Task (PrefectHQ#5077)

* Adds glob test with requested recursion changes

* Fixes typo in function name

Co-authored-by: Michael Adkins <[email protected]>

* Fixes typo in doc string

Co-authored-by: Michael Adkins <[email protected]>

* Adds more details to doc string for path arg

Co-authored-by: Michael Adkins <[email protected]>

* Adds more detail to the doc string for pattern arg

Co-authored-by: Michael Adkins <[email protected]>

* Swaps out tmpdir for tmp_path

* Fixes formatting issues

* Fixes flake8 issues

Co-authored-by: Michael Adkins <[email protected]>

* PagerDuty integration review and revisions (PrefectHQ#5102)

* Revisions to integration setup after confirming steps, added image.

* Fix typos

* Minor revisions in response to review feedback.

* Clarify which Cloud plans support automations.

* Fixed ReadTimeoutErrors for kubernetes client after connection reset by keep-alive connections (PrefectHQ#5066)

* get batch + core k8s client from utilities and modified tests accordingly

* fixed test_generate_job_spec_prefect_logging_level_environment_variable test

* added keep-alive functionality & added test

* added change document

* black reformatting

* fixed test_kube_client_with_keep_alive

* added return type annotation

* fixed flake8 problems

Co-authored-by: Jonas Miederer <[email protected]>

* Update click requirement from <8.0,>=7.0 to >=7.0,<9.0 (PrefectHQ#4534)

Updates the requirements on [click](https://github.com/pallets/click) to permit the latest version.
- [Release notes](https://github.com/pallets/click/releases)
- [Changelog](https://github.com/pallets/click/blob/main/CHANGES.rst)
- [Commits](pallets/click@7.0...8.0.0)

Signed-off-by: dependabot[bot] <[email protected]>

Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Add padding to code blocks to remove scroll bar (PrefectHQ#5107)

* Docs: Orchestration orchestration cloud examples (PrefectHQ#5086)

* Reworded for clarify recommended use cases for Cloud and server.

* Minor change for consistency between installation pages.

* Edits to focus tutorial examples on Prefect Cloud.

* Revisions from feedback.

* Update docs/orchestration/README.md

Co-authored-by: Mariia Kerimova <[email protected]>

* Update docs/orchestration/getting-started/set-up.md

Co-authored-by: Mariia Kerimova <[email protected]>

* Update docs/orchestration/README.md

Co-authored-by: Michael Adkins <[email protected]>

* Update docs/orchestration/README.md

Co-authored-by: Michael Adkins <[email protected]>

* Update docs/orchestration/README.md

Co-authored-by: Michael Adkins <[email protected]>

Co-authored-by: Terry Dorsey <[email protected]>
Co-authored-by: Mariia Kerimova <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>

* update bigquery dependency to support version 2.0 (PrefectHQ#5084)

* update bigquery dependency to support version 2.0

* add change file update

* Retain upper pin < 3.0

Co-authored-by: Michael Adkins <[email protected]>

* Fix formatting (PrefectHQ#5119)

* Add clarification about default secrets (PrefectHQ#5097)

* Add clarification about default secrets

* Update docs/core/concepts/secrets.md

Co-authored-by: Michael Adkins <[email protected]>

Co-authored-by: Michael Adkins <[email protected]>

* Update the `FlowView` to be more robust to serialized flow changes in the backend (PrefectHQ#5116)

* Update the `FlowView` to be more robust to serialized flow changes in the backend

* Add changes/ entry

* Drop from args docstring because the docs tooling complains

* Drop `flow` from test view

* fix: typo in architecture docs (PrefectHQ#5121)

inteface -> interface

* Move artifacts functions to `prefect.backend.artifacts` (PrefectHQ#5117)

* Move artifacts functions to `prefect.backend.artifacts`

* Update API reference

* Remove artifacts from sidebar

* Update internal uses of artifacts

* Add changes/ entry

* Actually commit the new files :D

* Ignore line to long flake8 in old file

* Fixup tests

* Update src/prefect/backend/artifacts.py

Co-authored-by: Zach Angell <[email protected]>

* Update src/prefect/artifacts.py

Co-authored-by: Dustin Ngo <[email protected]>

* Rename throughout

Co-authored-by: Zach Angell <[email protected]>
Co-authored-by: Dustin Ngo <[email protected]>

* Specify dockerignore for docker storage (PrefectHQ#4980)

* Provide option to specify `.dockerignore` when using Docker storage

* Move copying `dockerignore` close to file copy code

* Test `dockerignore` is copied

* Remove needless asserts, update changelog

* Clean up changelog

* Update DockerStorage docstring

* Add idempotency_key to create_flow_run task (PrefectHQ#5125)

* Add idempotency_key to create_flow_run task

* Fix flow run tests

* Black

* Add `raise_final_state` bool to `wait_for_flow_run` (PrefectHQ#5129)

* Remove artifact creation from `create_flow_run`

* Add `raise_final_state` bool to `wait_for_flow_run`

* Release 0.15.8 (PrefectHQ#5123)

* Generate changelog for 0.15.8

* Cleanup changelog

* Bump latest link in docs

* Add PrefectHQ#5117 changelog entry

* Add PrefectHQ#4980 changelog entry

* Add changelog entries for PrefectHQ#5125 and PrefectHQ#5129

* Add server security note

* Add bnaul to contributors

* Fix missing secret while using kubernetes agent (PrefectHQ#5131)

* Fix missing secret while using kubernetes agent

* Use `client_type` instead of renaming variable

* Fix tests and add coverage for in/out cluster

* Add 0.15.9 changelog entry

* Bump timeout on failing 3.6 test?

* Bump latest link in docs (PrefectHQ#5132)

* Documentation fixes: ECSRun ambiguities in IAM roles in the ECSRun docstring (PrefectHQ#5124)

* fix local agent docstring + remove misleading versioning info

https://prefect-community.slack.com/archives/CL09KU1K7/p1635231234368400

* providing examples to clarify the docstring

* Create issue5110.yml

* Update ecs.py

* Update ecs.py

* adding a full flow example with S3 and ECR image since it's a common pattern

* Update ecs.py

* fix W291 trailing whitespace

* Update issue5110.yml

* Update docs/core/about_prefect/why-not-airflow.md

Co-authored-by: Terrence Dorsey <[email protected]>

* Update src/prefect/run_configs/ecs.py

Co-authored-by: Terrence Dorsey <[email protected]>

* Update src/prefect/run_configs/ecs.py

Co-authored-by: Terrence Dorsey <[email protected]>

* Update src/prefect/run_configs/ecs.py

Co-authored-by: Terrence Dorsey <[email protected]>

Co-authored-by: Anna Geller <[email protected]>
Co-authored-by: Terrence Dorsey <[email protected]>

* Added Firebolt to task library

* Document connecting to local server using `--network` (PrefectHQ#5156)

* Add `packaging` to requirements (PrefectHQ#5157)

* Add `packaging` to requirements

* Spell right, raise upper bound to cover most recent version

* Add `end_time` to `FlowRunView.get_logs` (PrefectHQ#5138)

* Add `end_time` to `FlowRunView.get_logs`

* Add tests for `get_logs`

* Add changes/ entry

* Refactored code for Firebolt Task Library

* Docs: Add automations UI docs and revisions to automations concepts (PrefectHQ#5137)

* Automation concepts: edits and inter-page navigation

* Automations UI - add page elements

* Add page for automations UI to nav, minor edits.

* Automations UI documentation with screenshots

* Automations - document UI for viewing team automation actions.

* Refactored code for Firebolt Task Library

* Expose ports in docker run (PrefectHQ#5159)

* expose container ports via run_config

* changes file

* remove type as default is None

* Refactored code for Firebolt Task file

* Remove outdated cloud tier information. (PrefectHQ#5168)

* Add Prefect Server troubleshooting docs (PrefectHQ#5120)

* Add local deploy troubleshooting tips

* Link to the troubleshooting guide from the README

* Use separate file for troubleshooting tip

* Wording

* Update README.md

Co-authored-by: Terrence Dorsey <[email protected]>

* Apply suggestions from code review

Co-authored-by: Terrence Dorsey <[email protected]>

* Add Troubleshooting page to docs sidebar

* Fix filename of could-not-connect.png

Co-authored-by: Terrence Dorsey <[email protected]>

* Adds dbt artifact URLs to DbtCloudRunJob (PrefectHQ#5135)

* Adds dbt artifact URLs to DbtCloudRunJob

* Adds tests and changes file

* Fixes formatting

* Switches from link artifact to markdown artifact

* Removes unused import

* Updates docstring with details about dbt artifacts

* Adds functionality to include all dbt artifacts in markdown document

* Moves retrieval of artifacts after job completion

* Adds missing create_markdown_artifact call

* Updates change log entry

Co-authored-by: Michael Adkins <[email protected]>

* Clarifies docstring for DbtCloudBaseException

Co-authored-by: Michael Adkins <[email protected]>

Co-authored-by: Michael Adkins <[email protected]>

* document environment var overwrite in ECSRun (PrefectHQ#5169)

* document environment var overwrite in ECSRun

See discussion https://prefect-community.slack.com/archives/CL09KU1K7/p1637703686216800?thread_ts=1637627432.120600&cid=CL09KU1K7

* Update src/prefect/run_configs/ecs.py

Co-authored-by: Michael Adkins <[email protected]>

Co-authored-by: Michael Adkins <[email protected]>

* Use allowed_methods in Retry as method_whitelist is depreciated  (PrefectHQ#5167)

* Use allowed_methods in Retry

* type ignore allowed_methods + use urllib3 Retry

* Update comment re: typeshed

* Resolve marshmallow deprecation warning (PrefectHQ#5175)

This change fixes a deprecation warning that had been previously taken
care of by PR PrefectHQ#4903. It seems after that PR was merged some code using
marshmallow's fields.Dict class was added that used the deprecated
keyword argument `key` again. This just switches that over to use `keys`
instead. This deprecation was [made in marshmallow 3.10](https://github.com/marshmallow-code/marshmallow/blob/dev/CHANGELOG.rst#3100-2020-12-19) and you can
see the issue in marshmallow-code/marshmallow#1350. There are no new tests with this because the code change is so small. Nothing has changed which existing tests would not cover.

For more information see PrefectHQ#4903

This should resolve PrefectHQ#4540

* AirbyteConnectionTask (PrefectHQ#5078)

* Initial airbyte commit.

* v1 of AirbyteConnectionTask completed.

* Removed superfluous logging.

* Added responses for Mocking requests / Skip integration tests / Raise Signal.

* Added check for existing schedule.

* Fixed schedule checking...

* Formatted using black

* Added changes/pr5078.yaml & docs/outline.toml changes.

* Added missing keyword parameters to run()

* back --check . now passes!

* flake8 now passes!

* Moved responses to test-requirements.txt

* Fixes test failures

* Fixes static analysis errors

* Addresses review comments

Co-authored-by: Alexander Streed <[email protected]>
Co-authored-by: Alexander Streed <[email protected]>

* Feature/add soda spark task (PrefectHQ#5144)

* Added SodaSparkScan task

* Added changes file

* Added system deps CircleCI config

* Skipping soda spark tests involving PySpark

* Removed soda-sql dep and changed extra name for sodasql and sodaspark tasks

* Added OpenJDK to sys deps and re-add Soda Spark tests with PySpark

* Revert back to sodasql and sodaspark extras, enhance docstrings and tests

* Fixed SodaSpark tests

Co-authored-by: Alessandro Lollo <[email protected]>

* Update CODEOWNERS (PrefectHQ#5178)

* Remove Zach from global; Add Alex as task library owneer

* Add Terrence to docs

* Fixup indenting [ci skip]

* Docs: Add Concepts page for Artifacts API (PrefectHQ#5180)

* Artifacts - new concepts page for artifacts API

* Artifacts - minor edits.

* Artifacts - add to TOC, minor edits

* Artifacts - edits in response to feedback.

* Artifacts - update links and import path.

* Artifacts - correct API imports.

* Artifacts - revise examples, add update examples.

* Artifacts - add delete_artifact example.

* Artifacts - minor edits.

* Artifacts - add to TOC, minor edits

* Artifacts - edits in response to feedback.

* Fixup

Co-authored-by: Terrence Dorsey <[email protected]>

* Bump minimum urllib3 version to include `allowed_methods` (PrefectHQ#5179)

* Bump minimum urllib3 version to include `allowed_methods`

* Add upper pin to urllib3

* Bump requests version to match urllib3 requirements

* Fix task library CI?

* Ignore spark tests is there is an install issue

* Fix/flow to flow docs (PrefectHQ#5095)

* Updated flow to flow docs replacing StartFlowRun with create_flow_run

* Refactored docs to include wait_for_flow_run

* Fixed description for scheduling flow of flows

* Updated docs to reflect new raise_final_state

* Moved all create_flow_run and wait_for_flow_run call inside flow block

Co-authored-by: Alessandro Lollo <[email protected]>

* Documented how to configure prefect server on a VM  (PrefectHQ#3830)

* Documented how to configure prefect server
to run on a virtual machine (same as single node deploy)

* Apply suggestions from code review

Co-authored-by: Chris White <[email protected]>

* Review from @tpdorsey

Co-authored-by: Terrence Dorsey <[email protected]>

Co-authored-by: Jack Desert <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>
Co-authored-by: Chris White <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>
Co-authored-by: Terrence Dorsey <[email protected]>

* Release 0.15.10 (PrefectHQ#5181)

* Cut changelog

* Cleanup changelog

* Fixup skip

* Allow override of Prefect API url in docker runs and improve inference (PrefectHQ#5182)

* Allow the user to override the API; auto-infer shared server prefix

* Add tests covering fix

* Fixup empty networks check

* Add changes entry

* Improve airbyte task entry

Co-authored-by: Bill Palombi <[email protected]>

Co-authored-by: Bill Palombi <[email protected]>

* Add Prefect Partnership Integration annotation to docs (PrefectHQ#5160)

* Add Prefect Partnership Integration annotation to docs

* Adds badges to task library table

* Adds Airbyte to the ptask library overview page

Co-authored-by: Alexander Streed <[email protected]>

* Use `allowed_methods` instead of `methods_whitelist` which is deprecated by urllib3 (PrefectHQ#5191)

* Use `allowed_methods` instead of `methods_whitelist` which is depreacted by urllib3

* Add type-ignore

* Add regression test for client retries

* Docs: clarify steps for setting server as backend (PrefectHQ#5187)

* Clarify backend configuration in readme

* Update tasks table.

* Update tasks table, fixing links.

* Update deployment table markup for consistency.

* Clarify setup of backend.

* Update docs/orchestration/README.md

Co-authored-by: Michael Adkins <[email protected]>

Co-authored-by: Michael Adkins <[email protected]>

* Added pytests for Firebolt task library

* Refactored code for Firebolt task library

* Run black (PrefectHQ#5202)

* Refactored code for Firebolt task library

* beta1 is deprecated in Kubernetes v1.17+ (PrefectHQ#5194)

* beta1 is no longer needed

more here: https://kubernetes.io/docs/reference/access-authn-authz/rbac/#role-and-clusterrole

* Update src/prefect/agent/kubernetes/agent.py

Co-authored-by: Michael Adkins <[email protected]>

* Create pr5194.yaml

* avoid breaking tests

* new black version

* Revert "new black version"

This reverts commit 762431b.

Co-authored-by: Anna Geller <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>

* Changes partnership integration callout to verified task callout (PrefectHQ#5203)

* Pins version of mistune to prevent version conflict (PrefectHQ#5206)

* Updates RunGreatExpectationsValidation task to work with latest version of great_expectations (PrefectHQ#5185)

* Updates RunGreatExpectationsValidation task to work with latest version of Great Expectations

* Adds changes file

* Updates return line in docstring

* Updates example in docstring and adds run_id and evaluation_parameters to checkpoint run

* specifying that the checkpoint example can be used with both APIs

* black

* Fixes formatting

* Updates description in changes

Co-authored-by: Michael Adkins <[email protected]>

* Improves warning for old great_expectations versions

Co-authored-by: Michael Adkins <[email protected]>

* Fixes formatting

Co-authored-by: Anna Geller <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>

* Refactored code for Firebolt task library

* Refactored code for Firebolt task library

* Documented Firebolt in outline.toml as per Prefect documentation

* Fix error while managing K8s jobs when the flow run is deleted (PrefectHQ#5190)

* Ignore jobs for deleted flow runs

* Reorganize to avoid failing brittle test

* Always check if the flow run id is valid

* Add test for fix

* Make `ObjectNotFoundError` a `ClientError` subtype

* Add changes entry

* Add test for `ObjectNotFoundError` in `get_flow_run_state`

* docs(gql): use flow_runs instead of flow_run (PrefectHQ#5212)

Co-authored-by: Soham Parekh <[email protected]>

* Refactored code for Firebolt task library

* Refactored code for Firebolt task library

* Docs replace missing Census task (PrefectHQ#5219)

* Corrected missing items in task list.

* Resize logos for correct aspect.

* Update tasks.md (PrefectHQ#5218)

propogate -> propagate

* Fixing Serialized Flow Examples in Docs (PrefectHQ#5220)

* Fixing examples

* Changelog and linter

* Adding pygithub installation in docs

* Improve multiprocess based timeout handler (PrefectHQ#5213)

* Use `queue.get()` with timeout instead of `queue.empty`

* Add debug logs for pickling

* Add timeout log

* Force logs to flush on process exit

* Restructure queue retrieval; add log for unpickle

* Cleanup debugging code

* Remove breakpoint

* Kubernetes agent job YAML allowing imagePullSecrets to be unset as code intended (PrefectHQ#5001)

* Fixing issue

* reformat check

* Expand test coverage of KubernetesAgent

Add tests to verify behavior with different sources of
image pull secrets -- and the absence of such secrets.

In particular, this change tests what happens if the
IMAGE_PULL_SECRETS environment variable is set to an
empty string. This situation appears to trigger
issue PrefectHQ#5001.

* Reformat with black

* Catch an issue with KubernetesRun, add more tests

* Reformat with black

* Run black again with 21.12b0

* Add changelog entry

* Add contributor

* renamed to my name

Co-authored-by: Andrew Brookins <[email protected]>

* Fixing flow of flows code (PrefectHQ#5228)

* Fixing flow of flows code

* changelog

* slight change to param handling

* docs: update idempotency docs to indicate no expiry (PrefectHQ#5221)

* docs: update idempotency docs to indicate no expiry

* Update docs/orchestration/flow-runs/creation.md

Co-authored-by: Soham Parekh <[email protected]>
Co-authored-by: Terrence Dorsey <[email protected]>

* Allow kwargs to be passed to Merge (PrefectHQ#5233)

* Allow kwargs to be passed to Merge

Signed-off-by: Connor Martin <[email protected]>

* Add docs for merge kwargs

Signed-off-by: Connor Martin <[email protected]>

* Remove unneeded tests

* Add changes file

Signed-off-by: Connor Martin <[email protected]>

* Refactored code for Firebolt task library

* Refactored code for Firebolt task library

* Refactored code for Firebolt task library

* Refactored code for Firebolt task library

* Typo fix (PrefectHQ#5244)

`flow_run.get_flow_metdata()` -> `flow_run.get_flow_metadata()`

* add flow_run.get_logs() example (PrefectHQ#5245)

* add flow_run.get_logs() example

* Apply suggestions from code review

Added punctuation.

Co-authored-by: Terrence Dorsey <[email protected]>

* Add proxies argument to slack_notifier (PrefectHQ#5237)

* Add proxies argument to slack_notifier

* Add change file

* Correct docstring

* Multiple local agents with the same label (PrefectHQ#5232)

* Multiple local agents with the same label

* Create pr5232.yaml

* make it clearer with an example

Co-authored-by: Anna Geller <[email protected]>

* The AWS provided policy for ECS is `AmazonECSTaskExecutionRolePolicy` not `AmazonECSTaskExecutionPolicy` (PrefectHQ#5243)

see https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_execution_IAM_role.html

* Add instructions to save tutorial flow to a file for following steps. (PrefectHQ#5247)

* 5239 docs fixes (PrefectHQ#5248)

* Fix some invalid markup

* Fix some broken window/document references, add quotes around width

* Wrap table in clientonly tags

* Add other routes back in

* Update Server Getting Started docs (PrefectHQ#5254)

* Added 'firebolt' to setup.py file as an 'extra'

* Added changelog file for PR PrefectHQ#5263

* Updated docstrings for PR checklist.

* Updated docstrings of arguments for PR checklist.

* Refactored code for Firebolt Task Library

Co-authored-by: Michael Adkins <[email protected]>
Co-authored-by: Alexander Streed <[email protected]>
Co-authored-by: Terrence Dorsey <[email protected]>
Co-authored-by: Jonas Miederer <[email protected]>
Co-authored-by: Jonas Miederer <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Chris Arderne <[email protected]>
Co-authored-by: Terry Dorsey <[email protected]>
Co-authored-by: Mariia Kerimova <[email protected]>
Co-authored-by: Josh Wang <[email protected]>
Co-authored-by: Bruno P. Kinoshita <[email protected]>
Co-authored-by: Ari Pollak <[email protected]>
Co-authored-by: M. Alex Boyd <[email protected]>
Co-authored-by: Zach Angell <[email protected]>
Co-authored-by: Dustin Ngo <[email protected]>
Co-authored-by: Brett Naul <[email protected]>
Co-authored-by: Anna Geller <[email protected]>
Co-authored-by: Anna Geller <[email protected]>
Co-authored-by: Shahil Mawjee <[email protected]>
Co-authored-by: Andrew Brookins <[email protected]>
Co-authored-by: Constantino Schillebeeckx <[email protected]>
Co-authored-by: Kipling <[email protected]>
Co-authored-by: Michael Peteuil <[email protected]>
Co-authored-by: Nicholas Hemley <[email protected]>
Co-authored-by: Alexander Streed <[email protected]>
Co-authored-by: AlessandroLollo <[email protected]>
Co-authored-by: Alessandro Lollo <[email protected]>
Co-authored-by: Terrence Dorsey <[email protected]>
Co-authored-by: Jack Desert <[email protected]>
Co-authored-by: Jack Desert <[email protected]>
Co-authored-by: Chris White <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>
Co-authored-by: Bill Palombi <[email protected]>
Co-authored-by: meehawk <[email protected]>
Co-authored-by: Soham Parekh <[email protected]>
Co-authored-by: bmizhen <[email protected]>
Co-authored-by: Kevin Kho <[email protected]>
Co-authored-by: Farley Farley (yes, really) <[email protected]>
Co-authored-by: Andrew Brookins <[email protected]>
Co-authored-by: Connor Martin <[email protected]>
Co-authored-by: John Shearer <[email protected]>
Co-authored-by: Vincent Chéry <[email protected]>
Co-authored-by: Nicholas Brown <[email protected]>
Co-authored-by: Jenny G <[email protected]>

* Added changelog for Firebolt

* Refactored code and updated docstrings for Firebolt task library as per PR PrefectHQ#5265 comments.

* Updated docstrings for Firebolt task library as per flake8 guidelines.

* Refactored code as per PR review comments.

Co-authored-by: Michael Adkins <[email protected]>
Co-authored-by: Alexander Streed <[email protected]>
Co-authored-by: Terrence Dorsey <[email protected]>
Co-authored-by: Jonas Miederer <[email protected]>
Co-authored-by: Jonas Miederer <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Chris Arderne <[email protected]>
Co-authored-by: Terry Dorsey <[email protected]>
Co-authored-by: Mariia Kerimova <[email protected]>
Co-authored-by: Josh Wang <[email protected]>
Co-authored-by: Bruno P. Kinoshita <[email protected]>
Co-authored-by: Ari Pollak <[email protected]>
Co-authored-by: M. Alex Boyd <[email protected]>
Co-authored-by: Zach Angell <[email protected]>
Co-authored-by: Dustin Ngo <[email protected]>
Co-authored-by: Brett Naul <[email protected]>
Co-authored-by: Anna Geller <[email protected]>
Co-authored-by: Anna Geller <[email protected]>
Co-authored-by: Shahil Mawjee <[email protected]>
Co-authored-by: Andrew Brookins <[email protected]>
Co-authored-by: Constantino Schillebeeckx <[email protected]>
Co-authored-by: Kipling <[email protected]>
Co-authored-by: Michael Peteuil <[email protected]>
Co-authored-by: Nicholas Hemley <[email protected]>
Co-authored-by: Alexander Streed <[email protected]>
Co-authored-by: AlessandroLollo <[email protected]>
Co-authored-by: Alessandro Lollo <[email protected]>
Co-authored-by: Terrence Dorsey <[email protected]>
Co-authored-by: Jack Desert <[email protected]>
Co-authored-by: Jack Desert <[email protected]>
Co-authored-by: Chris White <[email protected]>
Co-authored-by: Michael Adkins <[email protected]>
Co-authored-by: Bill Palombi <[email protected]>
Co-authored-by: meehawk <[email protected]>
Co-authored-by: Soham Parekh <[email protected]>
Co-authored-by: bmizhen <[email protected]>
Co-authored-by: Kevin Kho <[email protected]>
Co-authored-by: Farley Farley (yes, really) <[email protected]>
Co-authored-by: Andrew Brookins <[email protected]>
Co-authored-by: Connor Martin <[email protected]>
Co-authored-by: John Shearer <[email protected]>
Co-authored-by: Vincent Chéry <[email protected]>
Co-authored-by: Nicholas Brown <[email protected]>
Co-authored-by: Jenny G <[email protected]>
  • Loading branch information
Show file tree
Hide file tree
Showing 5 changed files with 264 additions and 0 deletions.
5 changes: 5 additions & 0 deletions changes/pr5265.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
task:
- "Add 'Firebolt' to Prefect Task Library - [#5265](https://github.com/PrefectHQ/prefect/pull/5265)"

contributor:
- "[Raghav Sharma](https://github.com/raghavSharmaSigmoid)"
11 changes: 11 additions & 0 deletions src/prefect/tasks/firebolt/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
"""
This module contains a collection of tasks for interacting with Firebolt databases via
the firebolt-python-sdk library.
"""

try:
from prefect.tasks.firebolt.firebolt import FireboltQuery
except ImportError as err:
raise ImportError(
'Using `prefect.tasks.firebolt` requires Prefect to be installed with the "firebolt" extra.'
) from err
91 changes: 91 additions & 0 deletions src/prefect/tasks/firebolt/firebolt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
import firebolt.db.connection as firebolt_conn
from typing import List
from prefect import Task
from prefect.utilities.tasks import defaults_from_attrs


class FireboltQuery(Task):
"""
Task for executing a query against a Firebolt database.
Args:
- database (str): name of the database to use.
- username (str): username used to authenticate.
- password (str): password used to authenticate.
- engine_name (str): name of the engine to use.
- query (str): query to execute against database.
- **kwargs (dict, optional): additional keyword arguments to pass to the Task constructor.
"""

def __init__(
self,
database: str = None,
username: str = None,
password: str = None,
engine_name: str = None,
query: str = None,
**kwargs
):
self.database = database
self.username = username
self.password = password
self.engine_name = engine_name
self.query = query
super().__init__(**kwargs)

@defaults_from_attrs("database", "username", "password", "engine_name", "query")
def run(
self,
database: str = None,
username: str = None,
password: str = None,
engine_name: str = None,
query: str = None,
) -> List[List]:
"""
Task run method. Executes a query against Firebolt database.
Args:
- database (str): name of the database to use.
- username (str): username used to authenticate.
- password (str): password used to authenticate.
- engine_name (str): name of the engine to use.
- query (str): query to execute against database.
Returns:
- List[List]: output of 'cursor.fetchall()' if 'cursor.execute(query)' > 0,
else an empty list
Raises:
- ValueError: if a required parameter is not supplied.
- DatabaseError: if exception occurs when executing the query.
"""
if not database:
raise ValueError("A database name must be provided")
if not username:
raise ValueError("User name must be provided")
if not password:
raise ValueError("A password must be provided")
if not engine_name:
raise ValueError("An engine name must be provided")
if not query:
raise ValueError("A query string must be provided")

# build the connection parameter dictionary
conn_config = {
"database": database,
"username": username,
"password": password,
"engine_name": engine_name,
}

# connect to database, return a connection object
conn = firebolt_conn.connect(**conn_config)
result = []

# execute query
with conn:
with conn.cursor() as cursor:
if cursor.execute(query) > 0:
result = cursor.fetchall()
return result
3 changes: 3 additions & 0 deletions tests/tasks/firebolt/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
import pytest

pytest.importorskip("firebolt.db.connection")
154 changes: 154 additions & 0 deletions tests/tasks/firebolt/test_firebolt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
from unittest.mock import MagicMock
import pytest
import firebolt.db.connection as fb_conn
from prefect.tasks.firebolt.firebolt import FireboltQuery


@pytest.fixture
def mock_conn(monkeypatch):

# returns a mocked cursor for testing
firebolt_conn = MagicMock()
connection = MagicMock(spec=fb_conn.Connection)
cursor = MagicMock(spec=fb_conn.Cursor)

# link all the mocks together appropriately
firebolt_conn.return_value = connection
connection.cursor = cursor
firebolt_connection = MagicMock(connect=firebolt_conn)
monkeypatch.setattr(
"prefect.tasks.firebolt.firebolt.firebolt_conn", firebolt_connection
)

return firebolt_connection, cursor


class TestFireboltQuery:
def test_construction(self):
"""
Tests that all required params are present for FireboltQuery Task.
"""
task = FireboltQuery(
database="test",
username="test",
password="test",
engine_name="test",
query="test",
)
assert task.database == "test"
assert task.username == "test"
assert task.password == "test"
assert task.engine_name == "test"
assert task.query == "test"

def test_required_params(self):
"""
Tests to check if there are missing required parameters.
"""

# raises Value error if engine name is not provided
with pytest.raises(ValueError, match="An engine name must be provided"):
FireboltQuery().run(
database="test",
username="test",
password="test",
query="test",
)

# raises Value error if database name is not provided
with pytest.raises(ValueError, match="A database name must be provided"):
FireboltQuery().run(
username="test",
password="test",
engine_name="test",
query="test",
)
# raises Value error if username is not provided
with pytest.raises(ValueError, match="User name must be provided"):
FireboltQuery().run(
database="test",
password="test",
engine_name="test",
query="test",
)

# raises Value error if password is not provided
with pytest.raises(ValueError, match="A password must be provided"):
FireboltQuery().run(
database="test",
username="test",
engine_name="test",
query="test",
)

# raises Value error if query is not provided
with pytest.raises(ValueError, match="A query string must be provided"):
FireboltQuery().run(
database="test",
username="test",
password="test",
engine_name="test",
)

# test to check if the ddl/dml query was executed
def test_execute_query(self, mock_conn):
"""
Tests that the FireboltQuery Task calls the execute method on the cursor.
"""
connection = mock_conn[0]
cursor = mock_conn[1]

# setting execute return
cursor().__enter__().execute.return_value = 0
query = "SHOW DATABASES"
output = FireboltQuery(
database="test",
username="test",
password="test",
engine_name="test",
query=query,
).run()
conn_config = {
"database": "test",
"username": "test",
"password": "test",
"engine_name": "test",
}
connection.connect.assert_called_once_with(**conn_config)
cursor.assert_called_with()
cursor().__enter__().execute.assert_called_once_with(query)
cursor().__enter__().fetchall.assert_not_called()
assert output == []

# test to check if the query was executed and metadata was retrieved from database
def test_execute_fetchall(self, mock_conn):
"""
Tests that the FireboltQuery Task calls the fetchall method on the cursor.
"""
connection = mock_conn[0]
cursor = mock_conn[1]

# setting execute return
cursor().__enter__().execute.return_value = 1

# setting fetchall return
cursor().__enter__().fetchall.return_value = ["TESTDB"]
query = "SHOW DATABASES"
output = FireboltQuery(
database="test",
username="test",
password="test",
engine_name="test",
query=query,
).run()
conn_config = {
"database": "test",
"username": "test",
"password": "test",
"engine_name": "test",
}
connection.connect.assert_called_once_with(**conn_config)
cursor.assert_called_with()
cursor().__enter__().execute.assert_called_once_with(query)
cursor().__enter__().fetchall.assert_called()
assert output == ["TESTDB"]

0 comments on commit 7dbb3c8

Please sign in to comment.