Issue in connecting the GX operator with Snowflake data

[Bowrna]

I am currently using the following version:
airflow-provider-great-expectations==0.2.0

I am trying to run Great Expectations operator and pass the snowflake connection id and execute the query.

My code throws out the following error trace:

AIRFLOW_CTX_DAG_RUN_ID=manual__2023-01-30T10:18:21.654701+00:00
[2023-01-30, 10:18:36 UTC] {great_expectations.py:470} INFO - Running validation with Great Expectations...
[2023-01-30, 10:18:36 UTC] {great_expectations.py:472} INFO - Instantiating Data Context...
[2023-01-30, 10:18:36 UTC] {base.py:71} INFO - Using connection ID 'snowflake_conn' for task execution.
[2023-01-30, 10:18:36 UTC] {taskinstance.py:1851} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 474, in execute
    self.build_runtime_datasources()
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 368, in build_runtime_datasources
    self.build_runtime_sql_datasource_config_from_conn_id()
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 304, in build_runtime_sql_datasource_config_from_conn_id
    "connection_string": self.make_connection_string(),
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 251, in make_connection_string
    uri_string = f"snowflake://{self.conn.login}:{self.conn.password}@{self.conn.extra_dejson['extra__snowflake__account']}.{self.conn.extra_dejson['extra__snowflake__region']}/{self.conn.extra_dejson['extra__snowflake__database']}/{self.conn.schema}?warehouse={self.conn.extra_dejson['extra__snowflake__warehouse']}&role={self.conn.extra_dejson['extra__snowflake__role']}"  # noqa
KeyError: 'extra__snowflake__account'

On logging the conn in Airflow using the below code:

  @task
  def test_conn():
    from airflow.hooks.base import BaseHook
    conn = BaseHook.get_connection(SNOWFLAKE_CONN_ID)
    task_logger.info(f"connection info {conn.extra_dejson} %s")

Below is the information logged:

[2023-01-30, 10:18:28 UTC] {base.py:71} INFO - Using connection ID 'snowflake_conn' for task execution.
[2023-01-30, 10:18:28 UTC] {test_gx_snowflake.py:265} INFO - connection info {'account': 'rtb82372.us-east-1', 'insecure_mode': False, 'database': 'SFSALES_SFC_SAMPLES_VA3_SAMPLE_DATA', 'warehouse': 'XLARGEWH'} %s
[2023-01-30, 10:18:28 UTC] {python.py:177} INFO - Done. Returned value was: None

Reference issue at airflow side:
apache/airflow#26764

[Bowrna]

The above issue occur in the latest version of the great-expectations-operator too:
Following is the log trace for it

Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 474, in execute
    self.build_runtime_datasources()
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 383, in build_runtime_datasources
    self.datasource = self.build_runtime_sql_datasource_config_from_conn_id()
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 320, in build_runtime_sql_datasource_config_from_conn_id
    "connection_string": self.make_connection_string(),
  File "/usr/local/lib/python3.8/site-packages/great_expectations_provider/operators/great_expectations.py", line 255, in make_connection_string
    "account", self.conn.extra_dejson["extra__snowflake__account"]
KeyError: 'extra__snowflake__account'

[denimalpaca]

#84 should have closed this issue - can you please confirm you're using the operator version 0.2.4?

[denimalpaca]

@Bowrna can you confirm this is still an issue with the newest version? Thanks

[Bowrna]

@denimalpaca this issue is persisting in the version airflow-provider-great-expectations==0.2.4

I have verified it again

[Bowrna]

I have verified the code in this link https://github.com/astronomer/airflow-provider-great-expectations/blob/main/great_expectations_provider/operators/great_expectations.py#L253-L266

snowflake_account = self.conn.extra_dejson.get(
                "account", self.conn.extra_dejson["extra__snowflake__account"]
            )

in latest version of airflow, self.conn.extra_dejson dict doesn't have the key "extra__snowflake__account". so the above code starts to throw error.

We could handle it by changing the code like below

snowflake_account = self.conn.extra_dejson.get(
                "extra__snowflake__account", self.conn.extra_dejson.get("account", None)
            )

[denimalpaca]

I think that code should throw and error when no account is specified; do you think the wrong error is being thrown? The key error KeyError: 'extra__snowflake__account' should only occur now if no account is given. I don't think your fix would provide the correct behavior, as the error would then be propagated to the Great Expectations libraries when an invalid URI is passed.

Are you unable to use the operator because of the current behavior?

[mpgreg]

As Tamara noted... snowflake connector changed in 4.0.2: https://airflow.apache.org/docs/apache-airflow-providers-snowflake/stable/index.html#id2. Will post a PR here shortly.

[Bowrna]

I think that code should throw and error when no account is specified; do you think the wrong error is being thrown? The key error KeyError: 'extra__snowflake__account' should only occur now if no account is given. I don't think your fix would provide the correct behavior, as the error would then be propagated to the Great Expectations libraries when an invalid URI is passed.

Are you unable to use the operator because of the current behavior?

Yes, i am unable to use the operator because of the current behavior. I had to run the snowflake query in separate task and pass the result to the GX operator rather than passing the snowflake query directly to GX operator

[mpgreg]

PR for this #95

[denimalpaca]

@Bowrna I'm releasing a fix from @mpgreg 's PR, this should resolve the issue.

I don't feel we've come to a resolution about the account key error. I do think that as long as a Snowflake account is provided to the connection, this shouldn't be a problem in any Airflow version.

[Bowrna]

thanks @denimalpaca
yes this fix should solve the issue

[denimalpaca]

Hi @Bowrna I just released version 0.2.5 with this fix. Will close the issue!