From 9016dfa528cfc28f9c9023087dcdab8be7eb1aac Mon Sep 17 00:00:00 2001
From: vidit <vidit733@gmail.com>
Date: Fri, 13 Sep 2024 20:18:03 +0530
Subject: [PATCH] chore: fix for postgres migration script

---
 .../migrations/versions/4c70dae378f2_add_comment_feature.py | 6 +++++-
 .../versions/7f6cdb3621f7_add_data_element_support.py       | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/querybook/migrations/versions/4c70dae378f2_add_comment_feature.py b/querybook/migrations/versions/4c70dae378f2_add_comment_feature.py
index 31069c2b7..22b0d1036 100644
--- a/querybook/migrations/versions/4c70dae378f2_add_comment_feature.py
+++ b/querybook/migrations/versions/4c70dae378f2_add_comment_feature.py
@@ -18,13 +18,17 @@
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    MediumText = sa.Text(length=16777215)
+    conn = op.get_bind()
+    if conn.dialect.name = "postgresql":
+        MediumText = sa.Text()
     op.create_table(
         "comment",
         sa.Column("id", sa.Integer(), nullable=False),
         sa.Column("created_at", sa.DateTime(), nullable=True),
         sa.Column("updated_at", sa.DateTime(), nullable=True),
         sa.Column("created_by", sa.Integer(), nullable=True),
-        sa.Column("text", sa.Text(length=16777215), nullable=True),
+        sa.Column("text", MediumText, nullable=True),
         sa.Column("parent_comment_id", sa.Integer(), nullable=True),
         sa.Column("archived", sa.Boolean(), nullable=False),
         sa.ForeignKeyConstraint(["created_by"], ["user.id"], ondelete="SET NULL"),
diff --git a/querybook/migrations/versions/7f6cdb3621f7_add_data_element_support.py b/querybook/migrations/versions/7f6cdb3621f7_add_data_element_support.py
index 43e3e4c23..00f6ee8bf 100644
--- a/querybook/migrations/versions/7f6cdb3621f7_add_data_element_support.py
+++ b/querybook/migrations/versions/7f6cdb3621f7_add_data_element_support.py
@@ -17,6 +17,10 @@
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    DescriptionText = sa.Text(length=16777215)
+    conn = op.get_bind()
+    if conn.dialect.name = "postgresql":
+        DescriptionText = sa.Text()
     op.create_table(
         "data_element",
         sa.Column("id", sa.Integer(), nullable=False),
@@ -26,7 +30,7 @@ def upgrade():
         sa.Column("metastore_id", sa.Integer(), nullable=False),
         sa.Column("name", sa.String(length=255), nullable=False),
         sa.Column("type", sa.String(length=4096), nullable=False),
-        sa.Column("description", sa.Text(length=16777215), nullable=True),
+        sa.Column("description", DescriptionText, nullable=True),
         sa.Column("properties", sa.JSON(), nullable=True),
         sa.ForeignKeyConstraint(["created_by"], ["user.id"], ondelete="SET NULL"),
         sa.ForeignKeyConstraint(