CopyBot/Red-DiscordBot
1
0
Fork 0
mirror of https://github.com/Cog-Creators/Red-DiscordBot.git synced 2025-12-05 17:02:32 -05:00
Code Issues Packages Projects Releases Wiki Activity

Properly handle missing schemas/tables in PostgreSQL driver (#5855)

Browse Source
This commit is contained in:
Jakub Kuczys
2022-10-13 13:38:43 +02:00
committed by GitHub
parent a82c08c9d3
commit a3de616e4d
4 changed files with 113 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
redbot/core/drivers/postgres/ddl.sql
View File

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

20
redbot/core/drivers/postgres/postgres.py
View File

@@ -137,14 +137,11 @@ class PostgresDriver(BaseDriver):
}
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 @@ class PostgresDriver(BaseDriver):
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]