Support for passing arguments to SqlSensor underlying hooks #17592
Conversation
|
Any help fixing the pre-commit will be much appreciated. |
|
Looks like some tests are broken. Could anyone give a little context on the tests? |
|
Please rebase to latest |
dinigo
force-pushed
the
pass-arguments-to-dbapihook
branch
from
526b4ea to
f87a45c
Compare
|
Done rebasing |
|
|
||
| def get_hook(self): | ||
| """Return hook based on conn_type.""" | ||
| def get_hook(self, hook_kwargs: Dict): |
There was a problem hiding this comment.
Should this not be **hook_kwargs
| + f"Supported connection types: {list(allowed_conn_type)}" | ||
| ) | ||
| return conn.get_hook() | ||
| return conn.get_hook(**self.hook_kwargs) |
There was a problem hiding this comment.
… or this should be conn.get_hook(self.hook_kwargs)?
There was a problem hiding this comment.
I changed my mind about who should receive the destructured dict. I suppose you are right. because then we don't risk passing an argument to get_hook that is later popped out of the args, or unintendedly defined and taken out of the get_hook kwargs. I'll switch it the other way around. Just like discussed in the issue
| @mock.patch('airflow.sensors.sql.BaseHook') | ||
| def test_sql_sensor_bigquery_hook_kwargs(self, mock_hook): | ||
| op = SqlSensor( | ||
| task_id='sql_sensor_check', | ||
| conn_id='postgres_default', | ||
| sql="SELECT 1", | ||
| hook_kwargs={ | ||
| 'use_legacy_sql': False, | ||
| 'location': 'test_location', | ||
| }, | ||
| ) | ||
|
|
||
| mock_hook.get_connection('google_cloud_default').conn_type = "google_cloud_platform" | ||
| assert op._get_hook().use_legacy_sql | ||
| assert op._get_hook().location == 'test_location' |
There was a problem hiding this comment.
And this test should ensure the kwargs are correctly passed onto the hook.
There was a problem hiding this comment.
You mean mocking also the BaseHook and checking that call arguments match the ** hook_kwargs? you are right
|
Is this PR in progress? |
| success=None, | ||
| failure=None, | ||
| fail_on_empty=False, | ||
| hook_kwargs=None, |
There was a problem hiding this comment.
| hook_kwargs=None, | |
| hook_kwargs={}, |
And then the line below could be reduced to
self.hook_kwargs = hook_kwargs
|
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions. |
github-actions
bot
added
the
stale
|
Hi @SamWheating. |
github-actions
bot
removed
the
stale
|
Suppressed by #18431 |
6 participants
closes: #13750
related: #17315
SqlSensorrelies on underlying hooks that are automatically injected (DI) depending on the connection type provided to theSqlSensor. However, from time to time (but mainly inBigqueryHook) it is needed to pre-configure the "to-be-injected" hook because some default params do not fit the current config.For example, if using SqlSensor with Bigquery it by default configures the SQL dialect to "legacy SQL", and there is no way to change this through configuration. Instead the SqlSensor has to be extended and some functions overwritten (as mentioned in the #13750).
To be able to customise the underlying hook there are 2 approaches to fix this:
DbApiHookbut doesn't try to instantiate on its own.Either one has it's drawbacks:
Knowing that, and that @uranusjr approves this second approach, and also given the known limitations, this PR makes the SqlSensor pass over the hook_kwargs to the connection hook locator that returns an instantiated hook.