code: def upgrade():
      op.drop_table("ai_table")
  
      conn = op.get_bind()
  
      # views was created with unnamed fk. Therefore need recreate it
      op.create_table(
          "view_tmp",
          sa.Column("id", sa.Integer(), nullable=False),
          sa.Column("name", sa.String(), nullable=False),
          sa.Column("company_id", sa.Integer(), nullable=True),
          sa.Column("query", sa.String(), nullable=False),
          sa.Column("integration_id", sa.Integer(), nullable=False),
          sa.ForeignKeyConstraint(
              ["integration_id"],
              ["integration.id"],
              name="fk_integration_id",
          ),
          sa.PrimaryKeyConstraint("id"),
          sa.UniqueConstraint(
              "name", "company_id", name="unique_view_name_company_id"
          ),
      )
      conn.execute(
          text(
              """
          insert into view_tmp (id, name, company_id, query, integration_id)
          select id, name, company_id, query, datasource_id from view;
      """
          )
      )
      op.drop_table("view")
      op.rename_table("view_tmp", "view")
  
      op.create_table(
          "analysis",
          sa.Column("id", sa.Integer(), nullable=False),
          sa.Column(
              "analysis",
              mindsdb.interfaces.storage.db.Json(),
              nullable=False,
          ),
          sa.Column("created_at", sa.DateTime(), nullable=True),
          sa.Column("updated_at", sa.DateTime(), nullable=True),
          sa.PrimaryKeyConstraint("id"),
      )
  
      with op.batch_alter_table("datasource", schema=None) as batch_op:
          batch_op.add_column(
              sa.Column("analysis_id", sa.Integer(), nullable=True)
          )
          batch_op.create_foreign_key(
              "fk_ds_analysis_id", "analysis", ["analysis_id"], ["id"]
          )
          batch_op.add_column(
              sa.Column("ds_class", sa.String(), nullable=True)
          )
  
      session = sa.orm.Session(bind=conn)
      dsatasources = conn.execute(
          sa.text("select id, analysis from datasource")
      ).fetchall()
      for row in dsatasources:
          if row["analysis"] is not None:
              # NOTE 'returning' is relatively new in sqlite, so better will be use select after insert.
              conn.execute(
                  text(
                      """
                      insert into analysis (analysis) select analysis from datasource where id = :id;
                  """
                  ),
                  {"id": row["id"]},
              )
              analysis_id = conn.execute(
                  text(
                      """
                  select id from analysis order by id desc limit 1;
              """
                  )
              ).fetchall()
              conn.execute(
                  text(
                      """
                      update datasource set analysis_id = :analysis_id where id = :id
                  """
                  ),
                  {"analysis_id": analysis_id[0][0], "id": row["id"]},
              )
  
      with op.batch_alter_table("datasource", schema=None) as batch_op:
          batch_op.drop_column("analysis")
  
      op.create_table(
          "file",
          sa.Column("id", sa.Integer(), nullable=False),
          sa.Column("name", sa.String(), nullable=False),
          sa.Column("company_id", sa.Integer(), nullable=True),
          sa.Column("source_file_path", sa.String(), nullable=False),
          sa.Column("file_path", sa.String(), nullable=False),
          sa.Column("row_count", sa.Integer(), nullable=False),
          sa.Column(
              "columns",
              mindsdb.interfaces.storage.db.Json(),
              nullable=False,
          ),
          sa.Column(
              "created_at",
              sa.DateTime(),
              nullable=True,
              server_default=sa.func.current_timestamp(),
          ),
          sa.Column(
              "updated_at",
              sa.DateTime(),
              nullable=True,
              server_default=sa.func.current_timestamp(),
              server_onupdate=sa.func.current_timestamp(),
          ),
          sa.Column("analysis_id", sa.Integer(), nullable=True),
          sa.ForeignKeyConstraint(
              ["analysis_id"], ["analysis.id"], name="fk_analysis_id"
          ),
          sa.PrimaryKeyConstraint("id"),
          sa.UniqueConstraint(
              "name", "company_id", name="unique_file_name_company_id"
          ),
      )
  
      # delete ds where data is none
      dsatasources = conn.execute(
          text("select * from datasource")
      ).fetchall()
      for ds in dsatasources:
          if ds["data"] is None:
              conn.execute(
                  text("delete from datasource where id = :id"),
                  {"id": ds["id"]},
              )
              continue
          ds_data = json.loads(ds["data"])
          creation_info = json.loads(ds["creation_info"])
          datasource_name = ds_data.get("source_type")
          if datasource_name == "file":
              created_at = None
              if isinstance(ds["created_at"], str):
                  created_at = datetime.datetime.fromisoformat(
                      ds["created_at"]
                  )
              elif isinstance(ds["created_at"], [float, int]):
                  created_at = datetime.fromtimestamp(ds["created_at"])
  
              updated_at = None
              if isinstance(ds["updated_at"], str):
                  updated_at = datetime.datetime.fromisoformat(
                      ds["updated_at"]
                  )
              elif isinstance(ds["updated_at"], [float, int]):
                  updated_at = datetime.fromtimestamp(ds["updated_at"])
  
              file = mindsdb.interfaces.storage.db.File(
                  name=ds["name"],
                  company_id=ds["company_id"],
                  source_file_path=ds_data["source"],
                  file_path=creation_info["args"][0],
                  row_count=ds_data["row_count"],
                  columns=ds_data["columns"],
                  created_at=created_at,
                  updated_at=updated_at,
                  analysis_id=ds["analysis_id"],
              )
              session.add(file)
              session.flush()
              # ds_data['file_id'] = file.id
              ds_data["source"] = {
                  "mindsdb_file_name": ds["name"]
                  # 'source': ds_data['source']
              }
              conn.execute(
                  text(
                      """
                      update datasource set data = :ds_data where id = :id;
                  """
                  ),
                  {"id": ds["id"], "ds_data": json.dumps(ds_data)},
              )
  
          conn.execute(
              text(
                  """
                  update datasource
                  set integration_id = (select id from integration where name = :datasource_name and company_id = :company_id),
                      ds_class = :ds_class
                  where id = :id
              """
              ),
              {
                  "datasource_name": datasource_name,
                  "company_id": ds["company_id"],
                  "ds_class": creation_info["class"],
                  "id": ds["id"],
              },
          )
  
      session.commit()
  
      op.rename_table("datasource", "dataset")
  
      with op.batch_alter_table("dataset", schema=None) as batch_op:
          batch_op.create_foreign_key(
              "fk_integration_id",
              "integration",
              ["integration_id"],
              ["id"],
          )
  
      # NOTE two different 'batch' is necessary, in other way FK is not creating
      with op.batch_alter_table("predictor", schema=None) as batch_op:
          batch_op.alter_column(
              "datasource_id", new_column_name="dataset_id"
          )
      with op.batch_alter_table("predictor", schema=None) as batch_op:
          batch_op.create_foreign_key(
              "fk_predictor_dataset_id",
              "dataset",
              ["dataset_id"],
              ["id"],
          )
      with op.batch_alter_table("predictor", schema=None) as batch_op:
          batch_op.create_unique_constraint(
              "unique_predictor_name_company_id", ["name", "company_id"]
          )
  
      with op.batch_alter_table("integration", schema=None) as batch_op:
          batch_op.create_unique_constraint(
              "unique_integration_name_company_id",
              ["name", "company_id"],
          )
  
      with op.batch_alter_table("dataset", schema=None) as batch_op:
          batch_op.create_unique_constraint(
              "unique_dataset_name_company_id", ["name", "company_id"]
          )
  
imports: import json
  import datetime
  from alembic import op
  import sqlalchemy
  import mindsdb.interfaces.storage.db
  from sqlalchemy.sql import text
label_desc: Use image node_id 1 for calling a global function with example usage: upgrade() without return types