Skip to content

Commit

Permalink
Properly handle missing schemas/tables in PostgreSQL driver (Cog-Crea…
Browse files Browse the repository at this point in the history
  • Loading branch information
Jackenmen authored and Drapersniper committed Nov 27, 2022
1 parent 1dc88fe commit 7396af8
Show file tree
Hide file tree
Showing 3 changed files with 107 additions and 2 deletions.
10 changes: 10 additions & 0 deletions .github/workflows/tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
32 changes: 30 additions & 2 deletions redbot/core/drivers/postgres/ddl.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -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);

Expand All @@ -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);

Expand Down
67 changes: 67 additions & 0 deletions tests/core/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -602,3 +602,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)

0 comments on commit 7396af8

Please sign in to comment.