From 1b7d90cd13ca2e03d46a8f3c5b95599b35a8cac9 Mon Sep 17 00:00:00 2001
From: Varjitt Jeeva <varjitt.jeeva@autodesk.com>
Date: Mon, 6 Jan 2025 13:45:45 -0500
Subject: [PATCH] fix: support for indexes with quotes in the name #652 (#653)

* fix: support for indexes with quotes in the name #652

* doc: extra comment
---
 pgbelt/util/dump.py                          |  7 +++++--
 pgbelt/util/postgres.py                      |  2 +-
 tests/integration/conftest.py                |  4 +++-
 tests/integration/files/test_schema_data.sql | 22 ++++++++++++++++++++
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/pgbelt/util/dump.py b/pgbelt/util/dump.py
index 0f85c5ba..c299b769 100644
--- a/pgbelt/util/dump.py
+++ b/pgbelt/util/dump.py
@@ -93,7 +93,7 @@ async def dump_source_tables(
                 [
                     "pg_dump",
                     "--data-only",
-                    f"--table={config.schema_name}.{table}",
+                    f'--table={config.schema_name}."{table}"',
                     "-Fc",
                     "-f",
                     table_file(config.db, config.dc, table),
@@ -401,13 +401,16 @@ async def create_target_indexes(
     for c in create_index_statements.split(";"):
         # Get the Index Name
         regex_matches = search(
-            r"CREATE [UNIQUE ]*INDEX (?P<index>[a-zA-Z0-9._]+)+.*",
+            r"CREATE [UNIQUE ]*INDEX (?P<index>[a-zA-Z0-9._\"]+)+.*",
             c,
         )
         if not regex_matches:
             continue
         index = regex_matches.groupdict()["index"]
 
+        # Sometimes the index name is quoted, so remove the quotes
+        index = index.replace('"', "")
+
         # Create the index
         # Note that the host DSN must have a statement timeout of 0.
         # Example DSN: `host=server-hostname user=user dbname=db_name options='-c statement_timeout=3600000'`
diff --git a/pgbelt/util/postgres.py b/pgbelt/util/postgres.py
index c808d1a1..e56dd53f 100644
--- a/pgbelt/util/postgres.py
+++ b/pgbelt/util/postgres.py
@@ -264,7 +264,7 @@ async def table_empty(pool: Pool, table: str, schema: str, logger: Logger) -> bo
     return true if the table is empty
     """
     logger.info(f"Checking if table {table} is empty...")
-    result = await pool.fetch(f"SELECT * FROM {schema}.{table} LIMIT 1;")
+    result = await pool.fetch(f'SELECT * FROM {schema}."{table}" LIMIT 1;')
     return len(result) == 0
 
 
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index 00ea35c7..5ac3bc96 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -53,7 +53,9 @@ async def _create_dbupgradeconfigs() -> dict[str, DbupgradeConfig]:
         db_upgrade_config_kwargs["schema_name"] = (
             "non_public_schema" if "nonpublic" in s else "public"
         )
-        db_upgrade_config_kwargs["tables"] = ["UsersCapital"] if "exodus" in s else None
+        db_upgrade_config_kwargs["tables"] = (
+            ["UsersCapital", "existingSomethingIds"] if "exodus" in s else None
+        )
         db_upgrade_config_kwargs["sequences"] = (
             ["userS_id_seq"] if "exodus" in s else None
         )
diff --git a/tests/integration/files/test_schema_data.sql b/tests/integration/files/test_schema_data.sql
index 3dfb9af2..9948569c 100644
--- a/tests/integration/files/test_schema_data.sql
+++ b/tests/integration/files/test_schema_data.sql
@@ -66,6 +66,17 @@ CREATE INDEX users2_idx ON public."UsersCapital" (
     hash_lastname
 );
 
+-- Addressing the following index statement style: CREATE INDEX "existingEmailIds_email_id_idx" ON public."existingEmailIds" USING btree ("projectId", "emailId");
+-- Issue #652
+-- Did not add a primary key, helped iron out related quoting issues in the dump and load code.
+
+CREATE TABLE public."existingSomethingIds" (
+    "thingId" integer NOT NULL,
+    "somethingId" character varying(255) NOT NULL
+);
+
+CREATE INDEX "existingSomethingIds_something_id_idx" ON public."existingSomethingIds" USING btree ("thingId", "somethingId");
+
 --
 -- Name: userS_id_seq; Type: SEQUENCE; Schema: public; Owner: owner
 --
@@ -130,6 +141,17 @@ INSERT INTO public."UsersCapital2" (id, "hash_firstName", hash_lastname, gender)
 INSERT INTO public.another_test_table ("someThingIDontKnow", "anotherThing")
     VALUES ('0e095b60-ab7d-4892-9a92-6175497fe0f9', '0e095b60-ab7d-4892-9a92-6175497fe0f9');
 
+--
+-- Data for Name: existingSomethingIds; Type: TABLE DATA; Schema: public; Owner: owner
+--
+
+INSERT INTO public."existingSomethingIds" ("thingId", "somethingId")
+    VALUES (1, 'something1'),
+    (2, 'something2'),
+    (3, 'something3'),
+    (4, 'something4');
+
+
 --
 -- Name: userS_id_seq; Type: SEQUENCE SET; Schema: public; Owner: owner
 --