astro-sdk
astro-sdk copied to clipboard
Add a session parameter to all astro SDK decorators
Please describe the feature you'd like to see
There are certain session modifications that users might want to perform against SQL sessions before they run their SQL query. For example, users might want to add query tags to Snowflake queries for lineage and cost control. Other databases might have other options we can expand to in the future. Currently, the only alternative is to create a connection per set of query tags, which is unsustainable for companies that have lots of query tag combinations.
Describe the solution you'd like We propose to create a "session" parameter that takes a "SessionOptions" object similar to how we have "LoadOptions." For example:
@aql.transform(session=SnowflakeSession(query_tags={"team": "a"}))
def transform_table(input_table):
....
This SnowflakeSession class would contain a list of potential options, and internal logic in the form of an apply_session_parameters function that can be called at runtime. This would offer a good mix of session flexibility and simplicity (and users who want new session functionality can contribute back to the project or submit ticket requests, improving the experience for all other users).
Since our first use-case is exposing query_tags for snowflake users, a SnowflakeSession class with support for query tags will be sufficient for first-round support.
Are there any alternatives to this feature?
Another option would be to allow the session parameter to take a function that takes in a session object (similar to the pod_mutation_hook in the KubernetesExecutor). We don't want to expose this function since a) different databases have different session types and b) it might make the SDK too complicated.
Acceptance Criteria
- [ ] All checks and tests in the CI should pass
- [ ] Unit tests (90% code coverage or more, once available)
- [ ] Integration tests (if the feature relates to a new database or external service)
- [ ] Example DAG
- [ ] Docstrings in reStructuredText for each of methods, classes, functions and module-level attributes (including Example DAG on how it should be used)
- [ ] Exception handling in case of errors
- [ ] Logging (are we exposing useful information to the user? e.g. source and destination)
- [ ] Improve the documentation (README, Sphinx, and any other relevant)
- [ ] How to use Guide for the feature (example)