Merge pull request #155 from DocShow-AI/dev

Dev

.github/workflows/build-and-deploy.yaml

@@ -25,8 +25,15 @@ jobs:
         echo ${{ secrets.REGISTRY_PASSWORD }} | docker login registry.digitalocean.com -u ${{ secrets.REGISTRY_USERNAME }} --password-stdin
         docker push registry.digitalocean.com/docshow-ai/backend-prod:latest
 
+    - name: Build and Push Migration Docker Image
+      run: |
+        docker build -t registry.digitalocean.com/docshow-ai/migration-prod:latest -f backend/Dockerfile.migration backend/
+        echo ${{ secrets.REGISTRY_PASSWORD }} | docker login registry.digitalocean.com -u ${{ secrets.REGISTRY_USERNAME }} --password-stdin
+        docker push registry.digitalocean.com/docshow-ai/migration-prod:latest
+
     - name: Set up Kubectl and Update Deployment
       run: |
         echo "${{ secrets.KUBECONFIG_SECRET }}" | base64 --decode > kubeconfig.yml
         kubectl --kubeconfig=kubeconfig.yml set image deployment/frontend-deployment frontend=registry.digitalocean.com/docshow-ai/frontend-prod:latest
         kubectl --kubeconfig=kubeconfig.yml set image deployment/backend-deployment backend=registry.digitalocean.com/docshow-ai/backend-prod:latest
+        kubectl --kubeconfig=kubeconfig.yml apply -f k8s/migration-job.yaml

.github/workflows/python-quality.yml

@@ -21,10 +21,10 @@ jobs:
       run: pip install black flake8 mypy
 
     - name: Check Python code formatting with Black
-      run: black --check backend/
+      run: black --check --exclude backend/alembic/versions backend/
 
     - name: Check Python code style with Flake8
-      run: flake8 backend/
+      run: flake8 --exclude=backend/alembic/versions backend/
 
     - name: Run Mypy Type Checking
-      run: mypy backend/
+      run: mypy --ignore-missing-imports --exclude backend/alembic/versions backend/

.pre-commit-config.yaml

@@ -5,15 +5,18 @@ repos:
     -   id: black
         language_version: python3  # Ensure black uses the correct version of Python
         files: ^backend/  # Run black only on files in the backend/ directory
+        exclude: ^backend/alembic/versions/  # Exclude alembic migration files
 -   repo: https://github.com/PyCQA/flake8
     rev: 6.1.0
     hooks:
     -   id: flake8
         files: ^backend/  # Run flake8 only on files in the backend/ directory
+        exclude: ^backend/alembic/versions/  # Exclude alembic migration files
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
     rev: 'v1.7.1'  # Use the latest version
     hooks:
     -   id: mypy
         files: ^backend/
+        exclude: ^backend/alembic/versions/  # Exclude alembic migration files
 

backend/Dockerfile.migration

@@ -0,0 +1,22 @@
+# Use an official Python runtime as a parent image
+FROM python:3.8-slim-buster
+
+# Set the working directory in the container to /app
+WORKDIR /app
+
+# Add the current directory contents into the container at /app
+ADD . /app
+
+# Install Python packages
+RUN apt-get update && \
+    apt-get install -y gcc libffi-dev && \
+    pip install --trusted-host pypi.python.org -r requirements.txt
+
+# Cleanup
+RUN apt-get remove -y gcc libffi-dev && \
+    apt-get autoremove -y && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Run alembic migrations when the container launches
+CMD ["alembic", "upgrade", "head"]

backend/alembic.ini

@@ -1,3 +1,4 @@
 [alembic]
 # ... other settings ...
 sqlalchemy.url = driver://user:pass@localhost/dummy
+script_location = alembic

backend/alembic/env.py

@@ -1,15 +1,20 @@
 from alembic import context
 from sqlalchemy import engine_from_config, pool
 
-from models.base import Base
+import sys
+import os
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models.base import Base  # noqa: E402
 
 
 def run_migrations_online():
     connectable = context.config.attributes.get("connection")
 
     # For app database
     if connectable is None:
-        db_url = context.get_x_argument(as_dictionary=True).get("app_db", None)
+        db_url = context.get_x_argument(as_dictionary=True).get("db", None)
         if db_url:
             connectable = engine_from_config(
                 context.config.get_section(context.config.config_ini_section),

backend/alembic/script.py.mako

@@ -0,0 +1,24 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade():
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade():
+    ${downgrades if downgrades else "pass"}

backend/alembic/versions/a7364adb18ea_initial_migration.py

@@ -0,0 +1,24 @@
+"""Initial migration
+
+Revision ID: a7364adb18ea
+Revises: 
+Create Date: 2024-01-03 00:00:34.547254
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "a7364adb18ea"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    pass
+
+
+def downgrade():
+    pass

backend/databases/data_profile_manager.py → backend/database/data_profile_manager.py

@@ -1,34 +1,34 @@
-from models.data_profile import DataProfile, DataProfileCreateRequest
-
-
-class DataProfileManager:
-    def __init__(self, session):
-        self.session = session
-
-    def get_dataprofile_by_name(self, name):
-        """Retrieve a DataProfile by its name."""
-        return self.session.query(DataProfile).filter(DataProfile.name == name).first()
-
-    def get_all_data_profiles(self):
-        """Retrieve all DataProfiles."""
-        return self.session.query(DataProfile).all()
-
-    def create_dataprofile(self, data_profile_data: DataProfileCreateRequest):
-        """Create a new DataProfile."""
-        new_data_profile = DataProfile(
-            name=data_profile_data.name,
-            file_type=data_profile_data.file_type,  # Assuming it's included in the request
-            organization_id=data_profile_data.organization_id,  # Assuming it's included in the request
-            description=data_profile_data.description,
-        )
-        self.session.add(new_data_profile)
-        self.session.commit()
-        return new_data_profile.to_dict()
-
-    def get_dataprofile_by_id(self, data_profile_id: int):
-        """Retrieve a DataProfile by its ID."""
-        return (
-            self.session.query(DataProfile)
-            .filter(DataProfile.id == data_profile_id)
-            .first()
-        )
+from models.data_profile import DataProfile, DataProfileCreateRequest
+
+
+class DataProfileManager:
+    def __init__(self, session):
+        self.session = session
+
+    def get_dataprofile_by_name(self, name):
+        """Retrieve a DataProfile by its name."""
+        return self.session.query(DataProfile).filter(DataProfile.name == name).first()
+
+    def get_all_data_profiles(self):
+        """Retrieve all DataProfiles."""
+        return self.session.query(DataProfile).all()
+
+    def create_dataprofile(self, data_profile_data: DataProfileCreateRequest):
+        """Create a new DataProfile."""
+        new_data_profile = DataProfile(
+            name=data_profile_data.name,
+            file_type=data_profile_data.file_type,
+            organization_id=data_profile_data.organization_id,
+            description=data_profile_data.description,
+        )
+        self.session.add(new_data_profile)
+        self.session.commit()
+        return new_data_profile
+
+    def get_dataprofile_by_id(self, data_profile_id: int):
+        """Retrieve a DataProfile by its ID."""
+        return (
+            self.session.query(DataProfile)
+            .filter(DataProfile.id == data_profile_id)
+            .first()
+        )

backend/databases/database_manager.py → backend/database/database_manager.py

@@ -1,12 +1,12 @@
-from settings import DB_URL
+from settings import DATABASE_URL
 
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
 
 class DatabaseManager:
     def __init__(self):
-        self.engine = create_engine(DB_URL)
+        self.engine = create_engine(DATABASE_URL)
 
     def __enter__(self):
         Session = sessionmaker(bind=self.engine)

backend/databases/table_manager.py → backend/database/table_manager.py

@@ -1,13 +1,14 @@
 from sqlalchemy.orm import Session
-from typing import List, Optional
+from typing import Optional
 
 import pandas as pd
 
 from fastapi import HTTPException
-from databases.sql_executor import SQLExecutor
-from databases.table_metadata_manager import TableMetadataManager
+from database.sql_executor import SQLExecutor
+from database.table_map_manager import TableMapManager
+from database.table_metadata_manager import TableMetadataManager
 from llms.base import BaseLLM
-from models.organization_table_map import OrganizationTableMap
+from models.table_map import TableMap
 from utils.sql_string_manipulator import SQLStringManipulator
 
 
@@ -32,16 +33,13 @@ def _map_table_to_org(
         try:
             print(f"Mapping table {table_name} to organization {org_id}")
             if self.session:
-                print(f"Session: {self.session}")
                 alias = table_name if not alias else alias
-                self.session.add(
-                    OrganizationTableMap(
-                        organization_id=org_id, table_name=table_name, table_alias=alias
-                    )
+                table_map_manager = TableMapManager(self.session)
+                table_map = TableMap(
+                    organization_id=org_id, table_name=table_name, table_alias=alias
                 )
-                self.session.commit()
+                table_map_manager.create_table_map(table_map)
         except Exception as e:
-            self.session.rollback() if self.session else None
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
 
@@ -184,24 +182,6 @@ def execute_select_query(self, query: str, format_as_dict: bool = True):
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
 
-    def get_org_tables(self, org_id: int) -> List:
-        """Returns a list of names of all of the tables associated with an organization."""
-        try:
-            if self.session:
-                table_names = (
-                    self.session.query(OrganizationTableMap.table_name)
-                    .filter(OrganizationTableMap.organization_id == org_id)
-                    .all()
-                )
-                return [
-                    name[0] for name in table_names
-                ]  # Extracting table_name from each tuple
-            return []
-        except Exception as e:
-            self.session.rollback() if self.session else None
-            print(f"An error occurred: {e}")
-            raise HTTPException(status_code=400, detail=str(e))
-
     def get_table_columns(self, table_name: str):
         """Returns a list of all of the columns present within the table."""
         try:

backend/database/table_map_manager.py

@@ -0,0 +1,61 @@
+from fastapi import HTTPException
+from sqlalchemy.orm import Session
+from typing import List
+
+from models.table_map import TableMap
+
+
+class TableMapManager:
+    """
+    A class to manage CRUD operations related to the TableMap model.
+
+    Attributes:
+        db_session (Session): An active database session for performing operations.
+    """
+
+    def __init__(self, session: Session):
+        """
+        Initializes the TableMap with the given database session.
+
+        Args:
+            db_session (Session): The database session to be used for operations.
+        """
+        self.db_session = session
+
+    def create_table_map(self, table_map: TableMap):
+        """
+        Add a new table map to the database.
+
+        Args:
+            table_map (TableMap): The TableMap object to be added.
+
+        Returns:
+            TableMap: The created TableMap object.
+        """
+        try:
+            if self.db_session:
+                self.db_session.add(table_map)
+                self.db_session.commit()
+                return table_map
+        except Exception as e:
+            self.db_session.rollback() if self.db_session else None
+            print(f"An error occurred: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
+
+    def get_org_tables(self, org_id: int) -> List:
+        """Returns a list of names of all of the tables associated with an organization."""
+        try:
+            if self.db_session:
+                table_names = (
+                    self.db_session.query(TableMap.table_name)
+                    .filter(TableMap.organization_id == org_id)
+                    .all()
+                )
+                return [
+                    name[0] for name in table_names
+                ]  # Extracting table_name from each tuple
+            return []
+        except Exception as e:
+            self.db_session.rollback() if self.db_session else None
+            print(f"An error occurred: {e}")
+            raise HTTPException(status_code=400, detail=str(e))