Properly handle missing schemas/tables in PostgreSQL driver (Cog-Crea…

…tors#5855)

.github/workflows/tests.yml

@@ -91,3 +91,13 @@ jobs:
             PGPASSWORD: postgres
             PGPORT: 5432
           run: tox
+        - name: Verify no errors in PostgreSQL logs.
+          run: |
+            logs="$(docker logs "${{ job.services.postgresql.id }}" 2>&1)"
+            echo "---- PostgreSQL logs ----"
+            echo "$logs"
+            echo "---- PostgreSQL logs ----"
+            error_count="$(echo "$logs" | { grep -c 'ERROR: ' || true; })"
+            if [[ $error_count -gt 0 ]]; then
+              exit 1
+            fi

redbot/core/drivers/postgres/ddl.sql

@@ -137,10 +137,17 @@ CREATE OR REPLACE FUNCTION
     pkey_type CONSTANT text := red_utils.get_pkey_type(id_data.is_custom);
     whereclause CONSTANT text := red_utils.gen_whereclause(num_pkeys, pkey_type);
 
+    table_exists CONSTANT boolean := exists(
+      SELECT 1
+      FROM information_schema.tables
+      WHERE table_schema = schemaname AND table_name = id_data.category);
+
     missing_pkey_columns text;
 
   BEGIN
-    IF num_missing_pkeys <= 0 THEN
+    IF NOT table_exists THEN
+      -- If the table doesn't exist, just don't do anything to prevent SQL errors.
+    ELSIF num_missing_pkeys <= 0 THEN
       -- No missing primary keys: we're getting all or part of a document.
       EXECUTE format(
         'SELECT json_data #> $2 FROM %I.%I WHERE %s',
@@ -290,10 +297,25 @@ CREATE OR REPLACE FUNCTION
     num_identifiers CONSTANT integer := coalesce(array_length(id_data.identifiers, 1), 0);
     pkey_type CONSTANT text := red_utils.get_pkey_type(id_data.is_custom);
 
+    schema_exists CONSTANT boolean := exists(
+      SELECT 1
+      FROM red_config.red_cogs t
+      WHERE t.cog_name = id_data.cog_name AND t.cog_id = id_data.cog_id);
+    table_exists CONSTANT boolean := schema_exists AND exists(
+      SELECT 1
+      FROM information_schema.tables
+      WHERE table_schema = schemaname AND table_name = id_data.category);
+
     whereclause text;
 
   BEGIN
-    IF num_identifiers > 0 THEN
+    -- If the schema or table doesn't exist, just don't do anything to prevent SQL errors.
+    IF NOT schema_exists THEN
+      -- pass
+    ELSIF num_identifiers > 0 THEN
+      IF NOT table_exists THEN
+        RETURN;
+      END IF;
       -- Popping a key from a document or nested document.
       whereclause := red_utils.gen_whereclause(num_pkeys, pkey_type);
 
@@ -310,13 +332,19 @@ CREATE OR REPLACE FUNCTION
       USING id_data.pkeys, id_data.identifiers;
 
     ELSIF num_pkeys > 0 THEN
+      IF NOT table_exists THEN
+        RETURN;
+      END IF;
       -- Deleting one or many documents
       whereclause := red_utils.gen_whereclause(num_pkeys, pkey_type);
 
       EXECUTE format('DELETE FROM %I.%I WHERE %s', schemaname, id_data.category, whereclause)
       USING id_data.pkeys;
 
     ELSIF id_data.category IS NOT NULL AND id_data.category != '' THEN
+      IF NOT table_exists THEN
+        RETURN;
+      END IF;
       -- Deleting an entire category
       EXECUTE format('DROP TABLE %I.%I CASCADE', schemaname, id_data.category);
 

redbot/core/drivers/postgres/postgres.py

@@ -137,14 +137,11 @@ def get_config_details():
         }
 
     async def get(self, identifier_data: IdentifierData):
-        try:
-            result = await self._execute(
-                "SELECT red_config.get($1)",
-                encode_identifier_data(identifier_data),
-                method=self._pool.fetchval,
-            )
-        except asyncpg.UndefinedTableError:
-            raise KeyError from None
+        result = await self._execute(
+            "SELECT red_config.get($1)",
+            encode_identifier_data(identifier_data),
+            method=self._pool.fetchval,
+        )
 
         if result is None:
             # The result is None both when postgres yields no results, or when it yields a NULL row
@@ -163,12 +160,7 @@ async def set(self, identifier_data: IdentifierData, value=None):
             raise errors.CannotSetSubfield
 
     async def clear(self, identifier_data: IdentifierData):
-        try:
-            await self._execute(
-                "SELECT red_config.clear($1)", encode_identifier_data(identifier_data)
-            )
-        except asyncpg.UndefinedTableError:
-            pass
+        await self._execute("SELECT red_config.clear($1)", encode_identifier_data(identifier_data))
 
     async def inc(
         self, identifier_data: IdentifierData, value: Union[int, float], default: Union[int, float]

tests/core/test_config.py

@@ -660,3 +660,70 @@ async def test_config_custom_partial_pkeys_set(config, pkeys, raw_args, result):
     group = config.custom("TEST", *pkeys)
     await group.set_raw(*raw_args, value=result)
     assert await group.get_raw(*raw_args) == result
+
+
+@pytest.mark.asyncio
+async def test_config_custom_get_raw_with_default_on_whole_scope(config):
+    config.init_custom("TEST", 3)
+    config.register_custom("TEST")
+
+    group = config.custom("TEST")
+    assert await group.get_raw(default=True) is True
+
+
+@pytest.mark.parametrize(
+    "pkeys,raw_args,to_set",
+    (
+        # no config data for (cog_name, cog_id) is present
+        ((), (), None),
+        ((1,), (), None),
+        ((1, 2), (), None),
+        ((1, 2, 3), (), None),
+        ((1, 2, 3), ("key1",), None),
+        ((1, 2, 3), ("key1", "key2"), None),
+        # config data for (cog_name, cog_id) is present but scope does not exist
+        ((), (), ()),
+        ((1,), (), ()),
+        ((1, 2), (), ()),
+        ((1, 2, 3), (), ()),
+        ((1, 2, 3), ("key1",), ()),
+        ((1, 2, 3), ("key1", "key2"), ()),
+        # the scope exists with no records
+        ((1,), (), ("1",)),
+        ((1, 2), (), ("1",)),
+        ((1, 2, 3), (), ("1",)),
+        ((1, 2, 3), ("key1",), ("1",)),
+        ((1, 2, 3), ("key1", "key2"), ("1",)),
+        # scope with partial primary key (1,) exists
+        ((1, 2), (), ("1", "2")),
+        ((1, 2, 3), (), ("1", "2")),
+        ((1, 2, 3), ("key1",), ("1", "2")),
+        ((1, 2, 3), ("key1", "key2"), ("1", "2")),
+        # scope with partial primary key (1, 2) exists
+        ((1, 2, 3), (), ("1", "2", "3")),
+        ((1, 2, 3), ("key1",), ("1", "2", "3")),
+        ((1, 2, 3), ("key1", "key2"), ("1", "2", "3")),
+        # scope with full primary key (1, 2, 3)
+        ((1, 2, 3), ("key1",), ("1", "2", "3", "key1")),
+        ((1, 2, 3), ("key1", "key2"), ("1", "2", "3", "key1")),
+        # scope with full primary key (1, 2, 3) and a group named "key1" exists
+        ((1, 2, 3), ("key1", "key2"), ("1", "2", "3", "key1", "key2")),
+    ),
+)
+@pytest.mark.asyncio
+async def test_config_custom_clear_identifiers_that_do_not_exist(config, pkeys, raw_args, to_set):
+    config.init_custom("TEST", 3)
+    config.register_custom("TEST")
+
+    group = config.custom("TEST", *pkeys)
+    if to_set is not None:
+        data = {}
+        partial = data
+        for key in to_set:
+            partial[key] = {}
+            partial = partial[key]
+        scope = config.custom("TEST")
+        await scope.set(data)
+        # Clear needed to be able to differ between missing config data and missing scope data
+        await scope.clear_raw(*to_set)
+    await group.clear_raw(*raw_args)