Repository: astronomer/astro-sdk

Issue:
Convert current TruncateOperator into a DropTableOperator
The current TruncateOperator is implemented to drop tables instead of truncating them. This could cause critical issues with users who expect tables to exist after truncation. While not the correct action, there IS use for a DropTableOperator. 

Acceptance criteria:

- [ ] Create an operator that given a table will drop said table.
- [ ] If the user does not have proper permissions to drop the table, give proper warning in the form of an error
- [ ] create a corresponding AQL function


Example:
```
@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)

    with sample_dag:
        aql.drop_table(
            table=test_table,
        )
    test_utils.run_dag(sample_dag)

    assert not database.table_exists(test_table)
```
Issue when truncating after load_file with output_table
**Describe the bug**

Truncate after loading to a file failed:
```
    with DAG(dag_name, schedule_interval=None, start_date=START_DATE) as dag:
        chunk_size = int(os.getenv("ASTRO_CHUNK_SIZE", DEFAULT_CHUNK_SIZE))
        table_metadata = table_args.pop("metadata", {})
        if table_metadata:
            table = Table(metadata=Metadata(**table_metadata), **table_args)
        else:
            table = Table(**table_args)

        table_xcom = aql.load_file(  # noqa: F841
            input_file=File(
                path=dataset_path,
                conn_id=dataset_conn_id,
                filetype=FileType(dataset_filetype),
            ),
            task_id="load",
            output_table=table,
            chunk_size=chunk_size,
        )
        aql.truncate(table_xcom)
```

With the error message:

```
  File "/Users/al-cht01/Code/astro-sdk/src/astro/sql/operators/truncate.py", line 20, in __init__
    task_id = task_id or get_unique_task_id(table.name + "_truncate")
AttributeError: 'XComArg' object has no attribute 'name'
```
This was discovered during:
https://github.com/astronomer/astro-sdk/pull/514

**Version**
* Astro: 0.11.0

**Expected behaviour**
The DAG should succeed without errors



Discussion:

Resolving patch:
diff --git a/example_dags/example_snowflake_partial_table_with_append.py b/example_dags/example_snowflake_partial_table_with_append.py
--- a/example_dags/example_snowflake_partial_table_with_append.py
+++ b/example_dags/example_snowflake_partial_table_with_append.py
@@ -9,10 +9,10 @@
     append,
     cleanup,
     dataframe,
+    drop_table,
     load_file,
     run_raw_sql,
     transform,
-    truncate,
 )
 from astro.sql.table import Metadata, Table
 
@@ -136,7 +136,7 @@ def example_snowflake_partial_table_with_append():
     # We truncate this table only to avoid wasting Snowflake resources
     # Why? Between 2022-03-25 and 2022-04-11 it accumulated 301G (89 million rows) because
     # this example DAG used to append rows without deleting them
-    truncate_results = truncate(
+    truncate_results = drop_table(
         table=Table(name="homes_reporting", conn_id=SNOWFLAKE_CONN_ID)
     )
     truncate_results.set_upstream(record_results)
diff --git a/src/astro/sql/__init__.py b/src/astro/sql/__init__.py
--- a/src/astro/sql/__init__.py
+++ b/src/astro/sql/__init__.py
@@ -12,12 +12,12 @@
 from astro.sql.operators.append import APPEND_COLUMN_TYPE, AppendOperator
 from astro.sql.operators.cleanup import CleanupOperator
 from astro.sql.operators.dataframe import DataframeOperator
+from astro.sql.operators.drop import DropTableOperator
 from astro.sql.operators.export_file import export_file  # noqa: F401
 from astro.sql.operators.load_file import load_file  # noqa: F401
 from astro.sql.operators.merge import MERGE_COLUMN_TYPE, MergeOperator
 from astro.sql.operators.raw_sql import RawSQLOperator
 from astro.sql.operators.transform import TransformOperator  # noqa: F401
-from astro.sql.operators.truncate import TruncateOperator
 from astro.sql.table import Table
 
 
@@ -224,18 +224,18 @@ def merge(
     ).output
 
 
-def truncate(
+def drop_table(
     table: Table,
     **kwargs: Any,
-) -> TruncateOperator:
+):
     """
-    Truncate a table.
+    Drops a table.
 
-    :param table: Table to be truncated
+    :param table: Table to be dropped
     :param kwargs:
     """
 
-    return TruncateOperator(table=table, **kwargs)
+    return DropTableOperator(table=table, **kwargs).output
 
 
 def dataframe(
diff --git a/src/astro/sql/operators/truncate.py b/src/astro/sql/operators/drop.py
similarity index 70%
rename from src/astro/sql/operators/truncate.py
rename to src/astro/sql/operators/drop.py
--- a/src/astro/sql/operators/truncate.py
+++ b/src/astro/sql/operators/drop.py
@@ -7,8 +7,10 @@
 from astro.sql.table import Table
 
 
-class TruncateOperator(BaseOperator):
-    """Airflow Operator for truncating SQL tables."""
+class DropTableOperator(BaseOperator):
+    """Airflow Operator for dropping SQL tables."""
+
+    template_fields = ("table",)
 
     def __init__(
         self,
@@ -17,14 +19,15 @@ def __init__(
         **kwargs,
     ):
         self.table = table
-        task_id = task_id or get_unique_task_id(table.name + "_truncate")
+        task_id = task_id or get_unique_task_id("_drop")
         super().__init__(
             task_id=task_id,
             **kwargs,
         )
 
-    def execute(self, context: Dict) -> None:  # skipcq: PYL-W0613
+    def execute(self, context: Dict) -> Table:  # skipcq: PYL-W0613
         """Method run when the Airflow runner calls the operator."""
         database = create_database(self.table.conn_id)
         self.table = database.populate_table_metadata(self.table)
         database.drop_table(self.table)
+        return self.table


Test patch:
diff --git a/tests/sql/operators/test_truncate.py b/tests/sql/operators/test_drop_table.py
similarity index 63%
rename from tests/sql/operators/test_truncate.py
rename to tests/sql/operators/test_drop_table.py
--- a/tests/sql/operators/test_truncate.py
+++ b/tests/sql/operators/test_drop_table.py
@@ -1,8 +1,10 @@
-"""Tests to cover the truncate decorator"""
+"""Tests to cover the drop table function"""
 
 import pathlib
 
+import pandas
 import pytest
+from airflow.decorators import task
 
 import astro.sql as aql
 from astro.constants import Database
@@ -38,18 +40,31 @@
     indirect=True,
     ids=["sqlite", "postgres", "bigquery", "snowflake"],
 )
-def test_truncate_with_table_metadata(database_table_fixture, sample_dag):
-    """Test truncate operator for all databases."""
+def test_drop_table_with_table_metadata(database_table_fixture, sample_dag):
+    """Test drop table operator for all databases."""
     database, test_table = database_table_fixture
     assert database.table_exists(test_table)
+    tmp_table = test_table.create_similar_table()
+
+    @aql.dataframe
+    def do_nothing(df: pandas.DataFrame):
+        return df
+
+    @task
+    def validate_table_exists(table: Table):
+        assert database.table_exists(table)
+        assert table.name == tmp_table.name
+        return table
 
     with sample_dag:
-        aql.truncate(
-            table=test_table,
+        transformed_operator = do_nothing(test_table, output_table=tmp_table)
+        validated_table = validate_table_exists(transformed_operator)
+        aql.drop_table(
+            table=validated_table,
         )
     test_utils.run_dag(sample_dag)
 
-    assert not database.table_exists(test_table)
+    assert not database.table_exists(tmp_table)
 
 
 @pytest.mark.integration
@@ -65,13 +80,13 @@ def test_truncate_with_table_metadata(database_table_fixture, sample_dag):
     indirect=True,
     ids=["postgres"],
 )
-def test_truncate_without_table_metadata(database_table_fixture, sample_dag):
-    """Test truncate operator for all databases."""
+def test_drop_table_without_table_metadata(database_table_fixture, sample_dag):
+    """Test drop table operator for all databases."""
     database, test_table = database_table_fixture
     assert database.table_exists(test_table)
 
     with sample_dag:
-        aql.truncate(
+        aql.drop_table(
             table=test_table,
         )
     test_utils.run_dag(sample_dag)