From a3de616e4de6bcc97c8ae1dd09e1c1fb11a23537 Mon Sep 17 00:00:00 2001 From: Jakub Kuczys Date: Thu, 13 Oct 2022 13:38:43 +0200 Subject: [PATCH] Properly handle missing schemas/tables in PostgreSQL driver (#5855) --- .github/workflows/tests.yml | 10 ++++ redbot/core/drivers/postgres/ddl.sql | 32 ++++++++++- redbot/core/drivers/postgres/postgres.py | 20 +++---- tests/core/test_config.py | 67 ++++++++++++++++++++++++ 4 files changed, 113 insertions(+), 16 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index efd40b4bdd8..a34ad0b7214 100644 --- a/.github/workflows/tests.yml +++ b/.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 diff --git a/redbot/core/drivers/postgres/ddl.sql b/redbot/core/drivers/postgres/ddl.sql index 9d1c3c7b017..d7fa56b8b2e 100644 --- a/redbot/core/drivers/postgres/ddl.sql +++ b/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,6 +332,9 @@ 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); @@ -317,6 +342,9 @@ CREATE OR REPLACE FUNCTION 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); diff --git a/redbot/core/drivers/postgres/postgres.py b/redbot/core/drivers/postgres/postgres.py index 2b00e0cf281..c9c4c17cf0c 100644 --- a/redbot/core/drivers/postgres/postgres.py +++ b/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] diff --git a/tests/core/test_config.py b/tests/core/test_config.py index 6ad442ae6b5..473b4da2891 100644 --- a/tests/core/test_config.py +++ b/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)