Feat: Support Snowflake's travel time #414
Conversation
📝 WalkthroughWalkthroughThe changes introduce enhancements to the Changes
Possibly related PRs
Wdyt? Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (6)
tests/unit_tests/test_processors.py (4)
13-31: How about improving test isolation and SQL comparison?The test looks good overall! A couple of suggestions to make it even better:
Instead of using a global variable for
actual_cmd, we could use a local variable and return it from the_execute_sqlfunction. This would improve test isolation.The indentation in the
expected_cmdstring might cause comparison issues. We could use a heredoc-style string ortextwrap.dedent()to make it cleaner.What do you think about these changes? Here's a possible refactor:
import textwrap def test_snowflake_cache_config_data_retention_time_in_days( mocker: pytest_mock.MockFixture, ): expected_cmd = textwrap.dedent(""" CREATE TABLE airbyte_raw."table_name" ( col_name type ) DATA_RETENTION_TIME_IN_DAYS = 1 """).strip() actual_cmd = None def _execute_sql(cmd): nonlocal actual_cmd actual_cmd = cmd mocker.patch.object(SnowflakeSqlProcessor, "_execute_sql", side_effect=_execute_sql) config = _build_mocked_snowflake_processor(mocker, data_retention_time_in_days=1) config._create_table(table_name="table_name", column_definition_str="col_name type") assert actual_cmd.strip() == expected_cmdWDYT? This approach might make the test more robust and easier to maintain.
34-51: Shall we apply similar improvements here and address the extra newline?Great to see a test for the case without data retention time! A couple of suggestions:
We could apply the same improvements suggested for the previous test (using a local variable instead of global, and using
textwrap.dedent()).There's an extra newline character at the end of the
expected_cmdstring. We might want to remove it to avoid potential whitespace comparison issues.Here's a possible refactor:
import textwrap def test_snowflake_cache_config_no_data_retention_time_in_days( mocker: pytest_mock.MockFixture, ): expected_cmd = textwrap.dedent(""" CREATE TABLE airbyte_raw."table_name" ( col_name type ) """).strip() actual_cmd = None def _execute_sql(cmd): nonlocal actual_cmd actual_cmd = cmd mocker.patch.object(SnowflakeSqlProcessor, "_execute_sql", side_effect=_execute_sql) config = _build_mocked_snowflake_processor(mocker) config._create_table(table_name="table_name", column_definition_str="col_name type") assert actual_cmd.strip() == expected_cmdWhat do you think about these changes? They should make the test more consistent with the previous one and avoid potential whitespace issues.
54-75: Looks good! How about adding a quick comment for the mock?This helper function is well-structured and nicely parameterized. Good job on using SecretString for the password too!
One small suggestion: It might be helpful to add a quick comment explaining why we're mocking the
_ensure_schema_existsmethod. Something like:# Mock _ensure_schema_exists to isolate the test from actual schema creation mocker.patch.object( SnowflakeSqlProcessor, "_ensure_schema_exists", return_value=None )What do you think? This could help future developers understand the test setup more quickly.
1-75: Great job on these tests! How about some additional scenarios?These unit tests are well-structured and cover the main scenarios for the SnowflakeSqlProcessor's handling of data retention time. The use of a helper function for creating the mocked processor is a nice touch for code reuse.
To make the test suite even more robust, we might consider adding a few more test cases. Some ideas:
- Test with a large data retention time (e.g., 365 days)
- Test with a zero or negative data retention time (if these should be handled differently)
- Test error scenarios, like invalid data types for data_retention_time_in_days
What do you think about adding these? They could help catch edge cases and ensure the processor behaves correctly under various conditions.
airbyte/_processors/sql/snowflake.py (1)
47-55: New method looks great! Small suggestion for readability.The
get_create_table_extra_clausesmethod is well-implemented and aligns perfectly with the PR objective. It correctly adds the DATA_RETENTION_TIME_IN_DAYS clause when needed. Great job!A tiny suggestion: wdyt about using a list comprehension for a more concise implementation? Something like:
def get_create_table_extra_clauses(self) -> list[str]: return [f"DATA_RETENTION_TIME_IN_DAYS = {self.data_retention_time_in_days}"] if self.data_retention_time_in_days is not None else []This could make the method even more readable at a glance. What do you think? 🤔
airbyte/shared/sql_processor.py (1)
127-129: Would adding unit tests for the new method enhance test coverage?Adding unit tests for
get_create_table_extra_clausesand its integration within_create_tablecould help ensure that the extra clauses are handled correctly in different scenarios. Do you think it would be beneficial to include such tests? Wdyt?Also applies to: 660-666
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (3)
- airbyte/_processors/sql/snowflake.py (1 hunks)
- airbyte/shared/sql_processor.py (2 hunks)
- tests/unit_tests/test_processors.py (1 hunks)
🧰 Additional context used
🔇 Additional comments (2)
airbyte/_processors/sql/snowflake.py (2)
45-45: LGTM! New field for data retention looks good.The new
data_retention_time_in_daysfield aligns perfectly with the PR objective to support Snowflake's Time Travel feature. The type and default value seem appropriate for an optional configuration. Nicely done! 👍
Line range hint
1-255: Verify usage of new configuration in SnowflakeSqlProcessorThe changes look good and are well-contained within the SnowflakeConfig class. However, I noticed that the SnowflakeSqlProcessor class doesn't seem to be using the new
data_retention_time_in_daysconfiguration yet.Should we update the SnowflakeSqlProcessor to utilize this new configuration when creating tables? For example, in the
_write_files_to_new_tablemethod? What are your thoughts on this? 🤔To help verify this, we can run the following script:
This will help us confirm if and where we might need to update the SnowflakeSqlProcessor to use the new configuration.
|
/test-pr
|
There was a problem hiding this comment.
Thanks @yohannj for contributing! We might refactor the internals in the future (to reduce potential SQL injection footprint), but this looks solid to me and a great increment. The config input from the user won't be affected by future refactoring, so I'm happy to approve this.
Thanks so much for the contribution - and for your patience in us getting this merged.
Summary
Update sql_processor to allow adding "custom" extra clause at the end of the CREATE TABLE statement.
Update Snowflake processor to add an extra clause
DATA_RETENTION_TIME_IN_DAYSto configure Time Travel.This is a follow up of a conversation on Slack to have feature parity with Airbyte (PR)
Summary by CodeRabbit
New Features
Bug Fixes
Tests