[BUGFIX] Set application query param for SnowflakeDatasource #9863
Conversation
- This ends up being a no-op, since env var ends up being superceded by snowflake-sqlalchemy
roblim
changed the title
| self._engine = sa.create_engine(connection_string, **kwargs) | ||
| self._engine = sa.create_engine( | ||
| connection_string, | ||
| connect_args={"application": snowflake_partner_application}, |
There was a problem hiding this comment.
Just eyeballing this: does this conflict if the user has connect_args in their kwargs?
There was a problem hiding this comment.
Ah great catch! Yeah, would probably break. Will update to merge if present.
There was a problem hiding this comment.
Or actually, just noticed a test that suggests providing both str and connect args should yield ValueError. Lemme dig deeper and confirm expected behavior.
There was a problem hiding this comment.
Changed impl to add query param to url instead of using connect_args, mostly to allow for easier testing (once an engine is instantiated, there's no way to really check connect_args without actually connecting to Snowflake).
| _ = SnowflakeDatasource( | ||
| name="my_sf_ds", connection_string=connection_string, **connect_args | ||
| ) |
There was a problem hiding this comment.
This test wasn't working properly. It was passing only because same error was being thrown for missing name, and when this was fixed, passing because improper url was also throwing same error.
When test was fixed, it revealed that validation of connection_string XOR connection args wasn't happening.
|
|
||
| if values.get("connection_string") and connection_details: | ||
| raise ValueError( | ||
| "Cannot provide both a connection string or a combination of account, user, and password." | ||
| ) | ||
|
|
There was a problem hiding this comment.
test_conflicting_connection_string_and_args_raises_error was broken. After fixing, found that XOR check wasn't happening, so added this.
@Kilo59 - I assume we still wanted this XOR validation since we still had the above test, and there was a validator called _check_xor_input_args (but XOR check wasn't actually happening). LMK if you meant to remove XOR condition.
| @pydantic.root_validator | ||
| def _check_xor_input_args(cls, values: dict) -> dict: | ||
| @pydantic.validator("connection_string") | ||
| def _check_connection_string( |
There was a problem hiding this comment.
Renamed since it wasn't actually doing XOR check and was just validating connection_string.
| **connection_string | ||
| application=snowflake_partner_application, | ||
| **connection_string, | ||
| **kwargs, |
There was a problem hiding this comment.
@Kilo59 - I noticed **kwargs was missing here. Assumed we still wanted it since we're still passing **kwargs to engine for str connection_string, and we were passing **kwargs here before. Is this correct?
There was a problem hiding this comment.
I think the kwargs are passed elsewhere.
There was a problem hiding this comment.
I thought they might be too, but they end up getting filtered out in _get_connect_args, since kwargs is defined in parent class:
def _get_connect_args(self) -> dict[str, str | bool]:
excluded_fields: set[str] = set(SQLDatasource.__fields__.keys())
# dump as json dict to force serialization of things like AnyUrl
return self._json_dict(exclude=excluded_fields, exclude_none=True)
…ectations/great_expectations into set-sf-partner-by-context-type
| snowflake_partner_application = "great_expectations_oss" | ||
| else: | ||
| snowflake_partner_application = "great_expectations_oss" |
There was a problem hiding this comment.
And is there a way to decouple Cloud logic here? I don't love leaking this logic here but we might need to given our current architecture
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
# Conflicts: # docs/docusaurus/versioned_docs/version-0.18/snippets/redshift_python_example.py # great_expectations/datasource/fluent/snowflake_datasource.py
Set application query param for SnowflakeDatasource when creating engine. This is to ensure GX gets attribution for the Snowflake Partner program
This was attempted previously by setting the
SF_PARTNERenv var, which would typically be picked up by the Snowflake connecter to set the application query param. However, since GX uses snowflake-sqlalchemy to connect to Snowflake, the env var gets superseded, with attribution going to snowflake-sqlalchemy.PR title is prefixed with one of: [BUGFIX], [FEATURE], [DOCS], [MAINTENANCE], [CONTRIB]
Code is linted - run
invoke lint(usesruff format+ruff check)Appropriate tests and docs have been updated
For more information about contributing, see Contribute.
After you submit your PR, keep the page open and monitor the statuses of the various checks made by our continuous integration process at the bottom of the page. Please fix any issues that come up and reach out on Slack if you need help. Thanks for contributing!