From e7f1e94be368b33001f36fd823d4aef9cced58ae Mon Sep 17 00:00:00 2001
From: talevy42 <85844932+talevy42@users.noreply.github.com>
Date: Tue, 20 Jun 2023 10:56:35 -0700
Subject: [PATCH] add exists_ok param to add_column (#839)

* add exists_ok param to add_column

* feat(test): test exists_ok

* chore(linting): Black parsons

* switch exists_ok to if_exists

* chore(test): Fix test to use new param

* chore(linting): Lint tests too

---------

Co-authored-by: Tal Levy <tal@secondspectrum.com>
---
 parsons/etl/etl.py | 14 +++++++++++---
 test/test_etl.py   | 11 +++++++----
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/parsons/etl/etl.py b/parsons/etl/etl.py
index 75f9762fed..7797db0f2b 100644
--- a/parsons/etl/etl.py
+++ b/parsons/etl/etl.py
@@ -1,6 +1,7 @@
-import petl
 import logging
 
+import petl
+
 logger = logging.getLogger(__name__)
 
 
@@ -9,7 +10,7 @@ def __init__(self):
 
         pass
 
-    def add_column(self, column, value=None, index=None):
+    def add_column(self, column, value=None, index=None, if_exists="fail"):
         """
         Add a column to your table
 
@@ -20,12 +21,19 @@ def add_column(self, column, value=None, index=None):
                 A fixed or calculated value
             index: int
                 The position of the new column in the table
+            if_exists: str (options: 'fail', 'replace')
+                If set `replace`, this function will call `fill_column`
+                if the column already exists, rather than raising a `ValueError`
         `Returns:`
             `Parsons Table` and also updates self
         """
 
         if column in self.columns:
-            raise ValueError(f"Column {column} already exists")
+            if if_exists == "replace":
+                self.fill_column(column, value)
+                return self
+            else:
+                raise ValueError(f"Column {column} already exists")
 
         self.table = self.table.addfield(column, value, index)
 
diff --git a/test/test_etl.py b/test/test_etl.py
index 3a3013b522..d17b900bc3 100644
--- a/test/test_etl.py
+++ b/test/test_etl.py
@@ -1,13 +1,12 @@
-import unittest
-import petl
 import os
 import shutil
+import unittest
+from test.utils import assert_matching_tables
 
+import petl
 from parsons import Table
-from test.utils import assert_matching_tables
 from parsons.utilities import zip_archive
 
-
 # Notes :
 # - The `Table.to_postgres()` test is housed in the Postgres tests
 # - The `Table.from_postgres()` test is housed in the Postgres test
@@ -306,6 +305,10 @@ def test_column_add_dupe(self):
         # Test that we can't add an existing column name
         self.assertRaises(ValueError, self.tbl.add_column, "first")
 
+    def test_add_column_if_exists(self):
+        self.tbl.add_column("first", if_exists="replace")
+        self.assertEqual(self.tbl.columns, ["first", "last"])
+
     def test_remove_column(self):
         # Test that column is removed correctly
         self.tbl.remove_column("first")