From d4453e4a1d34dab0d2f97304a966e13daa2a9ca8 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Thu, 23 Jan 2025 14:59:02 +0100
Subject: [PATCH 1/3] fix: Add support for SQLite and PostgreSQL for inserting
 data in SQLAlchemyAdapter

---
 .../sqlalchemy/SqlAlchemyAdapter.py           | 37 +++++++++++++++----
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
index 68561979d..f09435b4f 100644
--- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
+++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
@@ -77,14 +77,37 @@ class SQLAlchemyAdapter:
                     text(f"DROP TABLE IF EXISTS {schema_name}.{table_name} CASCADE;")
                 )
 
-    async def insert_data(self, schema_name: str, table_name: str, data: list[dict]):
-        columns = ", ".join(data[0].keys())
-        values = ", ".join([f"({', '.join([f':{key}' for key in row.keys()])})" for row in data])
-        insert_query = text(f"INSERT INTO {schema_name}.{table_name} ({columns}) VALUES {values};")
+    async def insert_data(
+        self,
+        table_name: str,
+        data: list[dict],
+        schema_name: Optional[str] = "public",
+    ) -> int:
+        """
+        Insert data into specified table using SQLAlchemy Core with batch optimization
+        Returns number of inserted rows
+        """
+        if not data:
+            logger.info("No data provided for insertion")
+            return 0
 
-        async with self.engine.begin() as connection:
-            await connection.execute(insert_query, data)
-            await connection.close()
+        try:
+            # Dialect-agnostic table reference
+            if self.engine.dialect.name == "sqlite":
+                table = await self.get_table(table_name)  # SQLite ignores schemas
+            else:
+                table = await self.get_table(table_name, schema_name)
+
+            # Use SQLAlchemy Core insert with execution options
+            async with self.engine.begin() as conn:
+                result = await conn.execute(table.insert().values(data))
+
+                # Return rowcount for validation
+                return result.rowcount
+
+        except Exception as e:
+            logger.error(f"Insert failed: {str(e)}")
+            raise e  # Re-raise for error handling upstream
 
     async def get_schema_list(self) -> List[str]:
         """

From de19016494e031c91be034621831ea29de58369c Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Thu, 23 Jan 2025 15:10:27 +0100
Subject: [PATCH 2/3] fix: Add flag to allow SQLite to use foreign keys

---
 .../relational/sqlalchemy/SqlAlchemyAdapter.py    | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
index f09435b4f..f9e2eb6d1 100644
--- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
+++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
@@ -92,14 +92,17 @@ class SQLAlchemyAdapter:
             return 0
 
         try:
-            # Dialect-agnostic table reference
-            if self.engine.dialect.name == "sqlite":
-                table = await self.get_table(table_name)  # SQLite ignores schemas
-            else:
-                table = await self.get_table(table_name, schema_name)
-
             # Use SQLAlchemy Core insert with execution options
             async with self.engine.begin() as conn:
+                # Dialect-agnostic table reference
+                if self.engine.dialect.name == "sqlite":
+                    # Foreign key constraints are disabled by default in SQLite (for backwards compatibility),
+                    # so must be enabled for each database connection/session separately.
+                    await conn.execute(text("PRAGMA foreign_keys=ON"))
+                    table = await self.get_table(table_name)  # SQLite ignores schemas
+                else:
+                    table = await self.get_table(table_name, schema_name)
+
                 result = await conn.execute(table.insert().values(data))
 
                 # Return rowcount for validation

From 2e1a48e22c4ea0eafc1a5b1b6b531400cc5332a1 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Thu, 23 Jan 2025 15:13:46 +0100
Subject: [PATCH 3/3] docs: Add usage example of function

---
 .../relational/sqlalchemy/SqlAlchemyAdapter.py        | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
index f9e2eb6d1..257c89b8f 100644
--- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
+++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
@@ -86,6 +86,17 @@ class SQLAlchemyAdapter:
         """
         Insert data into specified table using SQLAlchemy Core with batch optimization
         Returns number of inserted rows
+
+        Usage Example:
+            from cognee.infrastructure.databases.relational.get_relational_engine import get_relational_engine
+            from uuid import UUID
+            db = get_relational_engine()
+            table_name = "groups"
+            data = {
+                "id": UUID("c70a3cec-3309-44df-8ee6-eced820cf438"),
+                "name": "test"
+            }
+            await db.insert_data(table_name, data)
         """
         if not data:
             logger.info("No data provided for insertion")