Feat: Support Snowflake's travel time #414
Conversation
| assert actual_cmd == expected_cmd | ||
|
|
||
|
|
||
| def _build_mocked_snowflake_processor( |
There was a problem hiding this comment.
Built it there by convenience and for code clarity. Removes some boilerplates from what the tests are actually doing.
There was a problem hiding this comment.
I didn't see the need to write integration tests on top of unit tests.
Let me know if you prefer to test this differently.
📝 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.
| def get_create_table_extra_clauses(self) -> list[str]: | ||
| """Return a list of clauses to append on CREATE TABLE statements.""" | ||
| return [] |
There was a problem hiding this comment.
Should we validate the output of get_create_table_extra_clauses to prevent SQL injection?
Currently, get_create_table_extra_clauses returns a list of clauses that are directly appended to the SQL command. If any of these clauses are built from untrusted input, it might introduce SQL injection vulnerabilities. Would it be a good idea to validate or sanitize these clauses before including them in the CREATE TABLE statement? Wdyt?
There was a problem hiding this comment.
@coderabbitai - This is not a free-form input but it's good to call out anyway. Do you have any suggestion to sanitize/pre-check the SQL?
There was a problem hiding this comment.
@aaronsteers, even though get_create_table_extra_clauses isn't accepting free-form input currently, it might be beneficial to add a validation step to ensure any clauses added are safe, especially if the method is extended in the future. For example, we could validate the clauses against a whitelist of allowed SQL statements:
def get_create_table_extra_clauses(self) -> list[str]:
clauses = [...] # Clauses to be added
allowed_clauses = {'DATA_RETENTION_TIME_IN_DAYS', 'CLUSTER BY', 'PARTITION BY'}
for clause in clauses:
if clause.split()[0].upper() not in allowed_clauses:
raise ValueError(f"Disallowed clause detected: {clause}")
return clausesThis way, we can prevent any unexpected or unsafe SQL clauses from being appended to the CREATE TABLE statement. Would implementing this kind of validation make sense? Wdyt?
There was a problem hiding this comment.
Having a startswith is not gonna be enough to forbid sql injections. If one of the database accepts multi-statement queries, one could do DATA_RETENTION_TIME_IN_DAYS=1; INSERT INTO ... (SELECT ...); CREATE TABLE foo (...) or similar things. That way even if there are multiple statements, they bypass it.
My take on this was to only let implementations of sql_processor specify their value.
However in a separator PR we could rework how query are being sent. Even though it's not a free form, we could still improve the security of the codebase against malicious actors.
def _execute_sql(self, sql: str | TextClause | Executable) -> CursorResultcould be rewritten with the signature
def _execute_sql(self, sql: str | TextClause | Executable, parameters: dict[str, Any]) -> CursorResultand then we update the few queries being sent in the code to use parameters instead of string concatenation (_add_column_to_table, _drop_temp_table, _append_temp_table_to_final_table aren't safe if I'm not wrong, especially because there is no normalizer or anything in _fully_qualified).
What do you think?
Actually, that's not gonna work for a CREATE TABLE statement.
| extra_clauses = "\n".join(self.sql_config.get_create_table_extra_clauses()) | ||
|
|
||
| cmd = f""" | ||
| CREATE TABLE {self._fully_qualified(table_name)} ( | ||
| {column_definition_str} | ||
| ) | ||
| {extra_clauses} |
There was a problem hiding this comment.
Is there a risk of SQL injection when appending extra_clauses to the SQL command?
In the _create_table method, extra_clauses are concatenated directly into the SQL statement. Should we consider using parameterized queries or ensuring that extra_clauses are properly sanitized to prevent potential security issues? What do you think?
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