From cf9edf2663e87cfa9a77972015f1f39beb4a462f Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Tue, 25 Nov 2025 18:03:35 +0100
Subject: [PATCH] chore: Add migration for new dataset database model field

---
 ...d2b2_expand_dataset_database_with_json_.py | 66 +++++++++++++++++++
 .../modules/users/models/DatasetDatabase.py   |  4 +-
 2 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 alembic/versions/46a6ce2bd2b2_expand_dataset_database_with_json_.py

diff --git a/alembic/versions/46a6ce2bd2b2_expand_dataset_database_with_json_.py b/alembic/versions/46a6ce2bd2b2_expand_dataset_database_with_json_.py
new file mode 100644
index 000000000..becd29226
--- /dev/null
+++ b/alembic/versions/46a6ce2bd2b2_expand_dataset_database_with_json_.py
@@ -0,0 +1,66 @@
+"""Expand dataset database with json connection field
+
+Revision ID: 46a6ce2bd2b2
+Revises: 76625596c5c3
+Create Date: 2025-11-25 17:56:28.938931
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "46a6ce2bd2b2"
+down_revision: Union[str, None] = "76625596c5c3"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def _get_column(inspector, table, name, schema=None):
+    for col in inspector.get_columns(table, schema=schema):
+        if col["name"] == name:
+            return col
+    return None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    insp = sa.inspect(conn)
+
+    vector_database_connection_info_column = _get_column(
+        insp, "dataset_database", "vector_database_connection_info"
+    )
+    if not vector_database_connection_info_column:
+        op.add_column(
+            "dataset_database",
+            sa.Column(
+                "vector_database_connection_info",
+                sa.JSON(),
+                unique=False,
+                nullable=False,
+                default={},
+            ),
+        )
+
+    graph_database_connection_info_column = _get_column(
+        insp, "dataset_database", "graph_database_connection_info"
+    )
+    if not graph_database_connection_info_column:
+        op.add_column(
+            "dataset_database",
+            sa.Column(
+                "graph_database_connection_info",
+                sa.JSON(),
+                unique=False,
+                nullable=False,
+                default={},
+            ),
+        )
+
+
+def downgrade() -> None:
+    op.drop_column("dataset_database", "vector_database_connection_info")
+    op.drop_column("dataset_database", "graph_database_connection_info")
diff --git a/cognee/modules/users/models/DatasetDatabase.py b/cognee/modules/users/models/DatasetDatabase.py
index b864fb951..fee323d2f 100644
--- a/cognee/modules/users/models/DatasetDatabase.py
+++ b/cognee/modules/users/models/DatasetDatabase.py
@@ -27,8 +27,8 @@ class DatasetDatabase(Base):
     # TODO: Instead of specifying and forwawrding all these individual fields, consider using a JSON field to store
     # configuration details for different database types. This would make it more flexible to add new database types
     # without changing the database schema.
-    graph_database_connection_info = Column(JSON, unique=False, nullable=True)
-    vector_database_connection_info = Column(JSON, unique=False, nullable=True)
+    graph_database_connection_info = Column(JSON, unique=False, nullable=False, default={})
+    vector_database_connection_info = Column(JSON, unique=False, nullable=False, default={})
 
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
     updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))