astro-sdk
astro-sdk copied to clipboard
Create a TruncateOperator
trafficstars
Since the current TruncateOperator is essentially a DropTableOperator, we will need to build a TruncateOperator from scratch.
A truncate operator should not delete a table. Rather, it should clear out all rows and leave the table empty.
Acceptance criteria:
- [ ] A TruncateOperator class that clears all rows without deleting the table
- [ ] A corresponding
aql.truncatefunction the enacts the TruncateOperator - [ ] A new set of tests for the described functionality
Example test:
@pytest.mark.integration
@pytest.mark.parametrize(
"database_table_fixture",
[
{
"database": Database.SQLITE,
"file": File(DEFAULT_FILEPATH),
},
{
"database": Database.POSTGRES,
"file": File(DEFAULT_FILEPATH),
},
{
"database": Database.BIGQUERY,
"file": File(DEFAULT_FILEPATH),
},
{
"database": Database.SNOWFLAKE,
"file": File(DEFAULT_FILEPATH),
},
],
indirect=True,
ids=["sqlite", "postgres", "bigquery", "snowflake"],
)
def test_truncate_with_table_metadata(database_table_fixture, sample_dag):
"""Test truncate operator for all databases."""
database, test_table = database_table_fixture
assert database.table_exists(test_table)
@aql.dataframe
def validate(df: pandas.DataFrame):
assert len(df) == 0
with sample_dag:
aql.truncate(
table=test_table,
)
test_utils.run_dag(sample_dag)
assert database.table_exists(test_table)