Drop columns for sql client impls

Author: anuunchin

dlt/destinations/impl/athena/athena.py

@@ -12,6 +12,7 @@
     Iterable,
     Type,
     cast,
+    Union,
 )
 from copy import deepcopy
 import re
@@ -231,6 +232,20 @@ def drop_tables(self, *tables: str) -> None:
         ]
         self.execute_many(statements)
 
+    def drop_columns(self, from_tables_drop_cols: List[Dict[str, Union[str, List[str]]]]) -> None:
+        """Drops specified columns from specified tables. The statement will fail if the columns don't exist."""
+
+        statements = []
+        for from_table_drop_cols in from_tables_drop_cols:
+            table = cast(str, from_table_drop_cols["from_table"])
+            for column in from_table_drop_cols["drop_columns"]:
+                statements.append(
+                    f"ALTER TABLE {self.make_qualified_ddl_table_name(table)} DROP COLUMN"
+                    f" {self.escape_column_name(column)};"
+                )
+
+        self.execute_many(statements)
+
     @contextmanager
     @raise_database_error
     def begin_transaction(self) -> Iterator[DBTransaction]:

dlt/destinations/impl/sqlalchemy/db_api_client.py

@@ -312,6 +312,19 @@ def drop_tables(self, *tables: str) -> None:
             tbl = sa.Table(table, self.metadata, schema=self.dataset_name, keep_existing=True)
             self.execute_sql(sa.schema.DropTable(tbl, if_exists=True))
 
+    def drop_columns(self, from_tables_drop_cols: List[Dict[str, Union[str, List[str]]]]) -> None:
+        for from_table_drop_cols in from_tables_drop_cols:
+            table_name = from_table_drop_cols["from_table"]
+            drop_columns = from_table_drop_cols["drop_columns"]
+
+            tbl = sa.Table(table_name, self.metadata, schema=self.dataset_name, keep_existing=True)
+            existing_cols = {col.name for col in tbl.columns}
+
+            for column in drop_columns:
+                if column in existing_cols:
+                    ddl = f"ALTER TABLE {self.make_qualified_table_name(table_name)} DROP COLUMN {self.escape_column_name(column)}"  # type: ignore[arg-type]
+                    self.execute_sql(ddl)
+
     def execute_sql(
         self, sql: Union[AnyStr, sa.sql.Executable], *args: Any, **kwargs: Any
     ) -> Optional[Sequence[Sequence[Any]]]: