Skip to content

Cog 414 fix postgres database deletion #163

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
borisarzentar merged 7 commits into from
Oct 28, 2024
Merged

Cog 414 fix postgres database deletion #163

Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ee2f1fb
fix: Add deletion of all tables in all schemas for postgres
dexters1 Oct 23, 2024
c1c0dc4
fix: Resolve issue with database deletion
dexters1 Oct 23, 2024
36377c0
fix: Move cleaning of MetaData out of drop table loop
dexters1 Oct 23, 2024
c67639f
refactor: Remove unnecessary print statement
dexters1 Oct 23, 2024
ddaa1e8
fix: Fix table deletion for SqlAlchemyAdapter
dexters1 Oct 25, 2024
09f5fa8
feat: Add deletion by id for SqlAlchemyAdapter
dexters1 Oct 27, 2024
dc8680f
fix: Add support for postgresql syntax for getting table function in …
dexters1 Oct 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 25 additions & 11 deletions cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
from os import path
from typing import Optional
from typing import AsyncGenerator
from contextlib import asynccontextmanager
from sqlalchemy import text, select
from sqlalchemy import text, select, MetaData
from sqlalchemy.orm import joinedload
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker

Expand Down Expand Up @@ -50,11 +51,14 @@ async def create_table(self, schema_name: str, table_name: str, table_config: li
await connection.execute(text(f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} ({', '.join(fields_query_parts)});"))
await connection.close()

async def delete_table(self, table_name: str):
async def delete_table(self, table_name: str, schema_name: Optional[str] = "public"):
async with self.engine.begin() as connection:
await connection.execute(text(f"DROP TABLE IF EXISTS {table_name} CASCADE;"))

await connection.close()
if self.engine.dialect.name == "sqlite":
# SQLite doesn’t support schema namespaces and the CASCADE keyword.
# However, foreign key constraint can be defined with ON DELETE CASCADE during table creation.
await connection.execute(text(f"DROP TABLE IF EXISTS {table_name};"))
else:
await connection.execute(text(f"DROP TABLE IF EXISTS {schema_name}.{table_name} CASCADE;"))
Comment on lines +55 to +62
Copy link
Contributor

@coderabbitai coderabbitai bot Oct 25, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Consider using SQL parameters for table and schema names.

The implementation correctly handles different database types, but direct string interpolation in SQL queries could be vulnerable to SQL injection. Consider using parameters:

-                await connection.execute(text(f"DROP TABLE IF EXISTS {table_name};"))
+                await connection.execute(text("DROP TABLE IF EXISTS :table_name;"), {"table_name": table_name})
-                await connection.execute(text(f"DROP TABLE IF EXISTS {schema_name}.{table_name} CASCADE;"))
+                await connection.execute(
+                    text("DROP TABLE IF EXISTS :schema_name.:table_name CASCADE;"),
+                    {"schema_name": schema_name, "table_name": table_name}
+                )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
async def delete_table(self, table_name: str, schema_name: Optional[str] = "public"):
async with self.engine.begin() as connection:
await connection.execute(text(f"DROP TABLE IF EXISTS {table_name} CASCADE;"))
await connection.close()
if self.engine.dialect.name == "sqlite":
# SQLite doesn’t support schema namespaces and the CASCADE keyword.
# However, foreign key constraint can be defined with ON DELETE CASCADE during table creation.
await connection.execute(text(f"DROP TABLE IF EXISTS {table_name};"))
else:
await connection.execute(text(f"DROP TABLE IF EXISTS {schema_name}.{table_name} CASCADE;"))
async def delete_table(self, table_name: str, schema_name: Optional[str] = "public"):
async with self.engine.begin() as connection:
if self.engine.dialect.name == "sqlite":
# SQLite doesn't support schema namespaces and the CASCADE keyword.
# However, foreign key constraint can be defined with ON DELETE CASCADE during table creation.
await connection.execute(text("DROP TABLE IF EXISTS :table_name;"), {"table_name": table_name})
else:
await connection.execute(
text("DROP TABLE IF EXISTS :schema_name.:table_name CASCADE;"),
{"schema_name": schema_name, "table_name": table_name}
)


async def insert_data(self, schema_name: str, table_name: str, data: list[dict]):
columns = ", ".join(data[0].keys())
Expand Down Expand Up @@ -119,12 +123,22 @@ async def delete_database(self):
self.db_path = None
else:
async with self.engine.begin() as connection:
# Load the schema information into the MetaData object
await connection.run_sync(Base.metadata.reflect)
for table in Base.metadata.sorted_tables:
drop_table_query = text(f"DROP TABLE IF EXISTS {table.name} CASCADE")
await connection.execute(drop_table_query)

result = await connection.execute(
text("""
SELECT schema_name FROM information_schema.schemata
WHERE schema_name NOT IN ('pg_catalog', 'pg_toast', 'information_schema');
""")
)
# Create a MetaData instance to load table information
metadata = MetaData()
# Drop all tables from all schemas
for schema in result.fetchall():
# Load the schema information into the MetaData object
await connection.run_sync(metadata.reflect, schema=schema[0])
for table in metadata.sorted_tables:
drop_table_query = text(f"DROP TABLE IF EXISTS {schema[0]}.{table.name} CASCADE")
await connection.execute(drop_table_query)
metadata.clear()
Copy link
Contributor

@coderabbitai coderabbitai bot Oct 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Consider dropping entire schemas instead of individual tables

Currently, the delete_database method iterates over each table within each schema and drops them individually. This can be inefficient, especially if there are many tables or other objects within the schemas. Dropping the entire schema in a single operation is more efficient and ensures that all objects within the schema are removed.

Apply this diff to simplify the database deletion process:

+                    # Drop entire schemas instead of individual tables
                     for schema in result.fetchall():
-                        # Load the schema information into the MetaData object
-                        await connection.run_sync(metadata.reflect, schema=schema[0])
-                        for table in metadata.sorted_tables:
-                            drop_table_query = text(f"DROP TABLE IF EXISTS {schema[0]}.{table.name} CASCADE")
-                            await connection.execute(drop_table_query)
-                        metadata.clear()
+                        drop_schema_query = text(f'DROP SCHEMA IF EXISTS "{schema[0]}" CASCADE')
+                        await connection.execute(drop_schema_query)

This modification:

  • Drops each schema using the DROP SCHEMA command with the CASCADE option, which removes the schema and all dependent objects.
  • Simplifies the code by eliminating the need to reflect metadata and iterate over tables.
  • Enhances performance by reducing the number of database operations.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
result = await connection.execute(
text("""
SELECT schema_name FROM information_schema.schemata
WHERE schema_name NOT IN ('pg_catalog', 'pg_toast', 'information_schema');
""")
)
# Create a MetaData instance to load table information
metadata = MetaData()
# Drop all tables from all schemas
for schema in result.fetchall():
# Load the schema information into the MetaData object
await connection.run_sync(metadata.reflect, schema=schema[0])
for table in metadata.sorted_tables:
drop_table_query = text(f"DROP TABLE IF EXISTS {schema[0]}.{table.name} CASCADE")
await connection.execute(drop_table_query)
metadata.clear()
result = await connection.execute(
text("""
SELECT schema_name FROM information_schema.schemata
WHERE schema_name NOT IN ('pg_catalog', 'pg_toast', 'information_schema');
""")
)
# Create a MetaData instance to load table information
metadata = MetaData()
# Drop entire schemas instead of individual tables
for schema in result.fetchall():
drop_schema_query = text(f'DROP SCHEMA IF EXISTS "{schema[0]}" CASCADE')
await connection.execute(drop_schema_query)

except Exception as e:
print(f"Error deleting database: {e}")

Expand Down