queryprofiler import restriction (#1199)

- Added error for the case where the tables already exist inside the schema.
vertica · Mar 29, 2024 · 2753b2c · 2753b2c
1 parent 6b7f3a2
commit 2753b2c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/verticapy/performance/vertica/collection/collection_tables.py b/verticapy/performance/vertica/collection/collection_tables.py
@@ -714,6 +714,15 @@ def export_table(self, tmp_path: Path) -> ExportMetadata:
             file_name=file_name, table_type=self.table_type, exported_rows=pdf_rows
         )
 
+    def check_if_tables_already_exist(self) -> int:
+        # Check if a table already exists
+        return f"""
+        SELECT CASE WHEN COUNT(*) > 0 THEN 1 ELSE 0 END AS table_exists
+        FROM v_catalog.tables 
+        WHERE table_schema = '{self.schema}'
+        AND table_name = '{self.get_import_name_fq().split('.')[-1]}';
+        """
+
 
 def getAllCollectionTables(
     target_schema: str, key: str, version: BundleVersion

diff --git a/verticapy/performance/vertica/collection/profile_import.py b/verticapy/performance/vertica/collection/profile_import.py
@@ -236,6 +236,15 @@ def _create_tables_if_not_exists(self) -> None:
             target_schema=self.target_schema, key=self.key, version=self.bundle_version
         )
         for ctable in all_tables.values():
+            if (
+                _executeSQL(ctable.check_if_tables_already_exist(), method="fetchall")[
+                    0
+                ][0]
+                == 1
+            ):
+                raise ProfileImportError(
+                    f"Schema '{self.target_schema}' and Key ID '{self.key}' are already in the database. Risk of overwriting data."
+                )
             self.logger.info(f"Running create statements for {ctable.name}")
             table_sql = ctable.get_create_table_sql()
             proj_sql = ctable.get_create_projection_sql()