property handle table's schema

Author: plaharanne

crate/operator/restore_backup.py

@@ -370,23 +370,7 @@ def get_restore_keyword(self, *, cursor: Cursor):
         if not tables or (len(tables) == 1 and tables[0].lower() == "all"):
             return "ALL"
 
-        def quote_table(table):
-            """
-            Ensure table names are correctly quoted. If it contains a schema
-            (e.g., 'doc.nyc_taxi'), quote both the schema and the table using
-            psycopg2.extensions.quote_ident.
-            """
-            if "." in table:
-                schema, table_name = table.split(".", 1)
-            else:
-                schema, table_name = None, table
-
-            quoted_schema = quote_ident(schema, cursor._impl) if schema else None
-            quoted_table = quote_ident(table_name, cursor._impl)
-
-            return f"{quoted_schema}.{quoted_table}" if quoted_schema else quoted_table
-
-        formatted_tables = [quote_table(table.strip()) for table in tables]
+        formatted_tables = [quote_table(table.strip(), cursor) for table in tables]
 
         return f'TABLE {",".join(formatted_tables)}'
 
@@ -1157,7 +1141,7 @@ async def remove_duplicated_tables(self, tables: Optional[List[str]] = None):
             async with self.conn_factory() as conn:
                 async with conn.cursor(timeout=120) as cursor:
                     if tables is not None:
-                        gc_tables = self.get_gc_tables(cursor, tables)
+                        gc_tables = [t for t in tables if t.startswith("gc.")]
                         tables_str = ",".join(f"'{table}'" for table in gc_tables)
                         where_stmt = f"t IN ({tables_str})"
                     else:
@@ -1176,8 +1160,10 @@ async def remove_duplicated_tables(self, tables: Optional[List[str]] = None):
                     self.gc_tables = [table[0] for table in tables] if tables else []
                     for table in self.gc_tables:
                         self.logger.info(f"Renaming GC table: {table} to {table}_temp")
+                        table_name = quote_table(table, cursor)
+                        temp_table_name = table_without_schema(f"{table}_temp", cursor)
                         await cursor.execute(
-                            f"ALTER TABLE {table} RENAME TO {table}_temp;"
+                            f"ALTER TABLE {table_name} RENAME TO {temp_table_name};"
                         )
         except DatabaseError as e:
             self.logger.warning(
@@ -1199,8 +1185,10 @@ async def restore_tables(self):
                 async with conn.cursor(timeout=120) as cursor:
                     for table in self.gc_tables:
                         self.logger.info(f"Renaming GC table: {table}_temp to {table}")
+                        table_name = table_without_schema(table, cursor)
+                        temp_table_name = quote_table(f"{table}_temp", cursor)
                         await cursor.execute(
-                            f"ALTER TABLE {table}_temp RENAME TO {table};"
+                            f"ALTER TABLE {temp_table_name} RENAME TO {table_name};"
                         )
         except DatabaseError as e:
             self.logger.warning(
@@ -1220,17 +1208,39 @@ async def cleanup_tables(self):
                 async with conn.cursor(timeout=120) as cursor:
                     for table in self.gc_tables:
                         self.logger.info(f"Dropping old GC table: {table}_temp")
-                        await cursor.execute(f"DROP TABLE {table}_temp;")
+                        temp_table_name = quote_table(f"{table}_temp", cursor)
+                        await cursor.execute(f"DROP TABLE {temp_table_name};")
         except DatabaseError as e:
             self.logger.warning(
                 "DatabaseError in RestoreGCTables.restore_tables", exc_info=e
             )
             raise kopf.PermanentError("grand-central table couldn't be renamed.")
 
-    @staticmethod
-    def get_gc_tables(cursor, tables: list[str]) -> list[str]:
-        return [
-            quote_ident(table, cursor._impl)
-            for table in tables
-            if table.startswith("gc.")
-        ]
+
+def quote_table(table, cursor) -> str:
+    """
+    Ensure table names are correctly quoted. If it contains a schema
+    (e.g., 'doc.nyc_taxi'), quote both the schema and the table using
+    psycopg2.extensions.quote_ident.
+    """
+    if "." in table:
+        schema, table_name = table.split(".", 1)
+    else:
+        schema, table_name = None, table
+
+    quoted_schema = quote_ident(schema, cursor._impl) if schema else None
+    quoted_table = quote_ident(table_name, cursor._impl)
+
+    return f"{quoted_schema}.{quoted_table}" if quoted_schema else quoted_table
+
+
+def table_without_schema(table, cursor) -> str:
+    """
+    Returns the table name without schema, ensuring it's correctly quoted..
+
+    :param table: The full table name, possibly including schema.
+    :param cursor: The database cursor used for quoting.
+    :return: The quoted table name without schema.
+    """
+    table_name = table.split(".")[1] if "." in table else table
+    return quote_ident(table_name, cursor._impl)

tests/test_restore_backup.py

@@ -94,6 +94,20 @@ def backup_repository_data(faker):
     }
 
 
+@pytest.fixture
+def mock_quote_ident():
+
+    def mock_quote_ident(value, connection):
+        if value.startswith('"') and value.endswith('"'):
+            return value
+        return f'"{value}"'
+
+    with mock.patch(
+        "crate.operator.restore_backup.quote_ident", side_effect=mock_quote_ident
+    ):
+        yield
+
+
 @pytest.mark.k8s
 @pytest.mark.asyncio
 @mock.patch("crate.operator.webhooks.webhook_client._send")
@@ -353,6 +367,7 @@ async def test_restore_backup_aws(
 @mock.patch("crate.operator.webhooks.webhook_client._send")
 @mock.patch.object(RestoreBackupSubHandler, "_create_backup_repository")
 @mock.patch.object(RestoreBackupSubHandler, "_ensure_snapshot_exists")
+@mock.patch.object(RestoreBackupSubHandler, "_start_restore_snapshot")
 @mock.patch.object(RestoreInternalTables, "remove_duplicated_tables")
 @mock.patch.object(RestoreInternalTables, "cleanup_tables")
 @pytest.mark.parametrize("gc_enabled", [True, False])
@@ -782,24 +797,17 @@ async def test_restore_backup_create_repo_fails(
         ),
     ],
 )
-def test_get_restore_type_keyword(restore_type, expected_keyword, params):
+def test_get_restore_type_keyword(
+    restore_type, expected_keyword, params, mock_quote_ident
+):
     cursor = mock.AsyncMock()
-
-    def mock_quote_ident(value, connection):
-        if value.startswith('"') and value.endswith('"'):
-            return value
-        return f'"{value}"'
-
-    with mock.patch(
-        "crate.operator.restore_backup.quote_ident", side_effect=mock_quote_ident
-    ):
-        func_kwargs = {}
-        if params:
-            func_kwargs[restore_type.value] = params
-        restore_keyword = RestoreType.create(
-            restore_type.value, **func_kwargs
-        ).get_restore_keyword(cursor=cursor)
-        assert restore_keyword == expected_keyword
+    func_kwargs = {}
+    if params:
+        func_kwargs[restore_type.value] = params
+    restore_keyword = RestoreType.create(
+        restore_type.value, **func_kwargs
+    ).get_restore_keyword(cursor=cursor)
+    assert restore_keyword == expected_keyword
 
 
 @pytest.mark.asyncio
@@ -981,8 +989,10 @@ async def does_credentials_secret_exist(
 def replace_gc_tables_data(faker, mock_cratedb_connection):
     repository = faker.domain_word()
     snapshot = faker.domain_word()
-    table_a = f"gc.{faker.domain_word()}"
-    table_b = f"gc.{faker.domain_word()}"
+    table_a = faker.domain_word()
+    table_b = faker.domain_word()
+    tables = [table_a, table_b]
+    tables_with_schema = [f"gc.{t}" for t in tables]
 
     mock_cursor = mock_cratedb_connection["mock_cursor"]
     mock_logger = mock.Mock(spec=logging.Logger)
@@ -995,8 +1005,8 @@ def replace_gc_tables_data(faker, mock_cratedb_connection):
     return (
         repository,
         snapshot,
-        table_a,
-        table_b,
+        tables,
+        tables_with_schema,
         mock_cursor,
         gc_tables_cls,
     )
@@ -1006,18 +1016,18 @@ def replace_gc_tables_data(faker, mock_cratedb_connection):
 @pytest.mark.parametrize(
     "gc_enabled, type_all", [(True, False), (False, True), (True, True), (False, False)]
 )
-async def test_replace_gc_tables(gc_enabled, type_all, replace_gc_tables_data):
-    repository, snapshot, table_a, table_b, mock_cursor, gc_tables_cls = (
+async def test_replace_gc_tables(
+    gc_enabled, type_all, replace_gc_tables_data, mock_quote_ident
+):
+    repository, snapshot, tables, tables_with_schema, mock_cursor, gc_tables_cls = (
         replace_gc_tables_data
     )
-    mock_cursor.fetchall.return_value = [(table_a,), (table_b,)] if gc_enabled else []
 
-    tables_param = None if type_all else [table_a, table_b]
+    fetch_response = [(t,) for t in tables_with_schema]
+    mock_cursor.fetchall.return_value = fetch_response if gc_enabled else []
+    tables_param = None if type_all else tables_with_schema
 
-    with mock.patch(
-        "crate.operator.restore_backup.quote_ident", side_effect=[table_a, table_b]
-    ):
-        await gc_tables_cls.remove_duplicated_tables(tables_param)
+    await gc_tables_cls.remove_duplicated_tables(tables_param)
 
     if type_all:
         stmts = [
@@ -1039,26 +1049,35 @@ async def test_replace_gc_tables(gc_enabled, type_all, replace_gc_tables_data):
                 "  FROM sys.snapshots "
                 "  WHERE repository=%s AND name=%s"
                 ") "
-                f"SELECT * FROM tables WHERE t IN ('{table_a}','{table_b}');",
+                "SELECT * FROM tables WHERE t IN "
+                f"('{tables_with_schema[0]}','{tables_with_schema[1]}');",
                 (repository, snapshot),
             ),
         ]
 
     if gc_enabled:
-        stmts.append(mock.call(f"ALTER TABLE {table_a} RENAME TO {table_a}_temp;"))
-        stmts.append(mock.call(f"ALTER TABLE {table_b} RENAME TO {table_b}_temp;"))
+        stmts.append(
+            mock.call(f'ALTER TABLE "gc"."{tables[0]}" RENAME TO "{tables[0]}_temp";')
+        )
+        stmts.append(
+            mock.call(f'ALTER TABLE "gc"."{tables[1]}" RENAME TO "{tables[1]}_temp";')
+        )
 
     mock_cursor.execute.assert_has_awaits(stmts)
     assert gc_tables_cls.gc_tables_renamed is True
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("gc_tables_renamed", [True, False])
-async def test_restore_gc_tables(gc_tables_renamed, replace_gc_tables_data):
-    _, _, table_a, table_b, mock_cursor, gc_tables_cls = replace_gc_tables_data
+async def test_restore_gc_tables(
+    gc_tables_renamed, replace_gc_tables_data, mock_quote_ident
+):
+    _, _, tables, tables_with_schema, mock_cursor, gc_tables_cls = (
+        replace_gc_tables_data
+    )
 
     gc_tables_cls.gc_tables_renamed = gc_tables_renamed
-    gc_tables_cls.gc_tables = [table_a, table_b]
+    gc_tables_cls.gc_tables = tables_with_schema
 
     await gc_tables_cls.restore_tables()
 
@@ -1067,19 +1086,27 @@ async def test_restore_gc_tables(gc_tables_renamed, replace_gc_tables_data):
     else:
         mock_cursor.execute.assert_has_awaits(
             [
-                mock.call(f"ALTER TABLE {table_a}_temp RENAME TO {table_a};"),
-                mock.call(f"ALTER TABLE {table_b}_temp RENAME TO {table_b};"),
+                mock.call(
+                    f'ALTER TABLE "gc"."{tables[0]}_temp" RENAME TO "{tables[0]}";'
+                ),
+                mock.call(
+                    f'ALTER TABLE "gc"."{tables[1]}_temp" RENAME TO "{tables[1]}";'
+                ),
             ]
         )
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("gc_tables_renamed", [True, False])
-async def test_cleanup_gc_tables(gc_tables_renamed, replace_gc_tables_data):
-    _, _, table_a, table_b, mock_cursor, gc_tables_cls = replace_gc_tables_data
+async def test_cleanup_gc_tables(
+    gc_tables_renamed, replace_gc_tables_data, mock_quote_ident
+):
+    _, _, tables, tables_with_schema, mock_cursor, gc_tables_cls = (
+        replace_gc_tables_data
+    )
 
     gc_tables_cls.gc_tables_renamed = gc_tables_renamed
-    gc_tables_cls.gc_tables = [table_a, table_b]
+    gc_tables_cls.gc_tables = tables_with_schema
 
     await gc_tables_cls.cleanup_tables()
 
@@ -1088,7 +1115,7 @@ async def test_cleanup_gc_tables(gc_tables_renamed, replace_gc_tables_data):
     else:
         mock_cursor.execute.assert_has_awaits(
             [
-                mock.call(f"DROP TABLE {table_a}_temp;"),
-                mock.call(f"DROP TABLE {table_b}_temp;"),
+                mock.call(f'DROP TABLE "gc"."{tables[0]}_temp";'),
+                mock.call(f'DROP TABLE "gc"."{tables[1]}_temp";'),
             ]
         )