Merge pull request #302 from RamiAwar/S18.1-enable-disable-schemas-and-tables

Toggle schemas and tables. FE, BE, the whole package.

Author: RamiAwar

backend/README.md

@@ -91,3 +91,22 @@ They stay if you ship something you're proud of, you've shipped too late.
 ```
 pre-commit install
 ```
+
+# Example DBMS-based databases
+
+## SQL Server:
+
+`docker run -p 1433:1433 -e 'ACCEPT_EULA=Y' -e 'SA_PASSWORD=My_password1' -d chriseaton/adventureworks:latest`
+DSN: `mssql://SA:My_password1@localhost/AdventureWorks?TrustServerCertificate=yes&driver=ODBC+Driver+18+for+SQL+Server`
+
+## PostgreSQL:
+
+Build & Run image locally: `bash ./scripts/postgres_img_with_sample_data.sh`
+
+DSN: `postgres://postgres:dvdrental@localhost:5432/dvdrental`
+
+## MySQL:
+
+`docker run -p 3306:3306 -d sakiladb/mysql`
+
+DSN: `mysql://sakila:p_ssW0rd@127.0.0.1:3306/sakila`

backend/alembic/versions/2024_08_15_1302-fa9cefccac47_add_connection_options_column_json.py

@@ -0,0 +1,35 @@
+"""Add Connection options column (JSON)
+
+Revision ID: fa9cefccac47
+Revises: 3f6e32040035
+Create Date: 2024-08-15 13:02:24.137632
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "fa9cefccac47"
+down_revision: Union[str, None] = "3f6e32040035"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("connections", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("options", sa.JSON(), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("connections", schema=None) as batch_op:
+        batch_op.drop_column("options")
+
+    # ### end Alembic commands ###

backend/dataline/api/conversation/router.py

@@ -132,7 +132,7 @@ async def execute_sql(
     connection = await connection_service.get_connection(session, connection_id)
 
     # Refresh chart data
-    db = SQLDatabase.from_uri(connection.dsn)
+    db = SQLDatabase.from_dataline_connection(connection)
     query_run_data = execute_sql_query(db, sql)
 
     # Execute query

backend/dataline/models/connection/model.py

@@ -1,6 +1,6 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypedDict
 
-from sqlalchemy import Boolean, String
+from sqlalchemy import Boolean, JSON, String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from dataline.models.base import DBModel, UUIDMixin
@@ -9,6 +9,21 @@
     from dataline.models.conversation.model import ConversationModel
 
 
+class ConnectionSchemaTable(TypedDict):
+    name: str
+    enabled: bool
+
+
+class ConnectionSchema(TypedDict):
+    name: str
+    tables: list[ConnectionSchemaTable]
+    enabled: bool
+
+
+class ConnectionOptions(TypedDict):
+    schemas: list[ConnectionSchema]
+
+
 class ConnectionModel(DBModel, UUIDMixin, kw_only=True):
     __tablename__ = "connections"
     dsn: Mapped[str] = mapped_column("dsn", String, nullable=False, unique=True)
@@ -17,6 +32,7 @@ class ConnectionModel(DBModel, UUIDMixin, kw_only=True):
     type: Mapped[str] = mapped_column("type", String, nullable=False)
     dialect: Mapped[str | None] = mapped_column("dialect", String)
     is_sample: Mapped[bool] = mapped_column("is_sample", Boolean, nullable=False, default=False, server_default="false")
+    options: Mapped[ConnectionOptions | None] = mapped_column("options", JSON, nullable=True)
 
     # Relationships
     conversations: Mapped[list["ConversationModel"]] = relationship("ConversationModel", back_populates="connection")

backend/dataline/models/connection/schema.py

@@ -8,6 +8,21 @@
 from dataline.config import config
 
 
+class ConnecitonSchemaTable(BaseModel):
+    name: str
+    enabled: bool
+
+
+class ConnectionSchema(BaseModel):
+    name: str
+    tables: list[ConnecitonSchemaTable]
+    enabled: bool
+
+
+class ConnectionOptions(BaseModel):
+    schemas: list[ConnectionSchema]
+
+
 class Connection(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
@@ -18,6 +33,7 @@ class Connection(BaseModel):
     dialect: str
     type: str
     is_sample: bool
+    options: Optional[ConnectionOptions] = None
 
 
 class ConnectionOut(Connection):
@@ -141,6 +157,7 @@ def validate_dsn_format(cls, value: str) -> str:
 class ConnectionUpdateIn(BaseModel):
     name: Optional[str] = None
     dsn: Optional[str] = None
+    options: Optional[ConnectionOptions] = None
 
     @field_validator("dsn")
     def validate_dsn_format(cls, value: str) -> str:

backend/dataline/repositories/connection.py

@@ -6,6 +6,8 @@
 
 from dataline.models.connection.model import ConnectionModel
 from dataline.repositories.base import AsyncSession, BaseRepository
+from dataline.models.connection.schema import ConnectionOptions
+
 
 class ConnectionType(Enum):
     csv = "csv"
@@ -16,6 +18,7 @@ class ConnectionType(Enum):
     snowflake = "snowflake"
     sas = "sas"
 
+
 class ConnectionCreate(BaseModel):
     model_config = ConfigDict(from_attributes=True, extra="ignore")
 
@@ -25,6 +28,7 @@ class ConnectionCreate(BaseModel):
     dialect: str
     type: str
     is_sample: bool = False
+    options: ConnectionOptions | None = None
 
 
 class ConnectionUpdate(BaseModel):
@@ -36,6 +40,7 @@ class ConnectionUpdate(BaseModel):
     dialect: str | None = None
     type: str | None = None
     is_sample: bool | None = None
+    options: ConnectionOptions | None = None
 
 
 class ConnectionRepository(BaseRepository[ConnectionModel, ConnectionCreate, ConnectionUpdate]):

backend/dataline/repositories/result.py

@@ -18,19 +18,23 @@ def model(self) -> Type[ResultModel]:
         return ResultModel
 
     async def get_dsn_from_result(self, session: AsyncSession, result_id: UUID) -> str:
+        connection = await self.get_connection_from_result(session, result_id)
+        return connection.dsn
+
+    async def get_connection_from_result(self, session: AsyncSession, result_id: UUID) -> ConnectionModel:
         query = (
-            select(ConnectionModel.dsn)
+            select(ConnectionModel)
             .join(ConversationModel)
             .join(MessageModel)
             .join(ResultModel)
             .where(ResultModel.id == result_id)
         )
         result = await session.execute(query)
-        dsn = result.fetchone()
-        if not dsn:
-            raise ValueError(f"Could not find DSN for result_id: {result_id}")
+        connection = result.fetchone()
+        if not connection:
+            raise ValueError(f"Could not find connection for result_id: {result_id}")
 
-        return dsn[0]
+        return connection[0]
 
     async def get_chart_from_sql_query(self, session: AsyncSession, sql_string_result_id: UUID) -> ResultModel:
         query = (

backend/dataline/services/connection.py

@@ -9,13 +9,18 @@
 import pandas as pd
 import pyreadstat
 from fastapi import Depends, UploadFile
-from sqlalchemy import create_engine
 from sqlalchemy.exc import OperationalError
 
 from dataline.config import config
 from dataline.errors import ValidationError
 from dataline.models.connection.model import ConnectionModel
-from dataline.models.connection.schema import ConnectionOut, ConnectionUpdateIn
+from dataline.models.connection.schema import (
+    ConnecitonSchemaTable,
+    ConnectionOptions,
+    ConnectionOut,
+    ConnectionSchema,
+    ConnectionUpdateIn,
+)
 from dataline.repositories.base import AsyncSession, NotFoundError, NotUniqueError
 from dataline.repositories.connection import (
     ConnectionCreate,
@@ -24,6 +29,7 @@
     ConnectionUpdate,
 )
 from dataline.services.file_parsers.excel_parser import ExcelParserService
+from dataline.services.llm_flow.utils import DatalineSQLDatabase as SQLDatabase
 from dataline.utils.utils import (
     forward_connection_errors,
     generate_short_uuid,
@@ -54,37 +60,29 @@ async def get_connections(self, session: AsyncSession) -> list[ConnectionOut]:
     async def delete_connection(self, session: AsyncSession, connection_id: UUID) -> None:
         await self.connection_repo.delete_by_uuid(session, connection_id)
 
-    async def get_connection_details(self, dsn: str) -> tuple[str, str, str]:
+    async def get_connection_details(self, dsn: str) -> SQLDatabase:
         # Check if connection can be established before saving it
         try:
-            engine = create_engine(dsn)
-            with engine.connect():
-                pass
-
-            dialect = engine.url.get_dialect().name
-            database = engine.url.database
+            db = SQLDatabase.from_uri(dsn)
+            database = db._engine.url.database
 
             if not database:
                 raise ValidationError("Invalid DSN. Database name is missing, append '/DBNAME'.")
 
-            return dialect, database, dsn
+            return db
 
         except OperationalError as exc:
             # Try again replacing localhost with host.docker.internal to connect with DBs running in docker
             if "localhost" in dsn:
                 dsn = dsn.replace("localhost", "host.docker.internal")
                 try:
-                    engine = create_engine(dsn)
-                    with engine.connect():
-                        pass
-
-                    dialect = engine.url.get_dialect().name
-                    database = engine.url.database
+                    db = SQLDatabase.from_uri(dsn)
+                    database = db._engine.url.database
 
                     if not database:
                         raise ValidationError("Invalid DSN. Database name is missing, append '/DBNAME'.")
 
-                    return dialect, database, dsn
+                    return db
                 except OperationalError as e:
                     logger.error(e)
                     raise ValidationError("Failed to connect to database, please check your DSN.")
@@ -124,10 +122,14 @@ async def update_connection(
                 raise NotUniqueError("Connection DSN already exists.")
 
             # Check if connection can be established before saving it
-            dialect, database, dsn = await self.get_connection_details(data.dsn)
-            update.dsn = dsn
-            update.database = database
-            update.dialect = dialect
+            db = await self.get_connection_details(data.dsn)
+            update.dsn = str(db._engine.url.render_as_string(hide_password=False))
+            update.database = db._engine.url.database
+            update.dialect = db.dialect
+            # TODO: What do we do for the options? Enable everything?
+        elif data.options:
+            # only modify options if dsn hasn't changed
+            update.options = data.options
 
         if data.name:
             update.name = data.name
@@ -140,20 +142,35 @@ async def create_connection(
         session: AsyncSession,
         dsn: str,
         name: str,
-        type: str | None = None,
+        connection_type: str | None = None,
         is_sample: bool = False,
     ) -> ConnectionOut:
         # Check if connection can be established before saving it
-        dialect, database, dsn = await self.get_connection_details(dsn)
-        if not type:
-            type = dialect
+        db = await self.get_connection_details(dsn)
+        if not connection_type:
+            connection_type = db.dialect
 
         # Check if connection already exists
         await self.check_dsn_already_exists(session, dsn)
-
+        connection_schemas: list[ConnectionSchema] = [
+            ConnectionSchema(
+                name=schema,
+                tables=[ConnecitonSchemaTable(name=table, enabled=True) for table in tables],
+                enabled=True,
+            )
+            for schema, tables in db._all_tables_per_schema.items()
+        ]
         connection = await self.connection_repo.create(
             session,
-            ConnectionCreate(dsn=dsn, database=database, name=name, dialect=dialect, type=type, is_sample=is_sample),
+            ConnectionCreate(
+                dsn=dsn,
+                database=db._engine.url.database,
+                name=name,
+                dialect=db.dialect,
+                type=connection_type,
+                is_sample=is_sample,
+                options=ConnectionOptions(schemas=connection_schemas),
+            ),
         )
         return ConnectionOut.model_validate(connection)
 
@@ -186,7 +203,9 @@ async def create_csv_connection(self, session: AsyncSession, file: UploadFile, n
 
         # Create connection with the locally copied file
         dsn = get_sqlite_dsn(str(file_path.absolute()))
-        return await self.create_connection(session, dsn=dsn, name=name, type=ConnectionType.csv.value, is_sample=False)
+        return await self.create_connection(
+            session, dsn=dsn, name=name, connection_type=ConnectionType.csv.value, is_sample=False
+        )
 
     async def create_excel_connection(self, session: AsyncSession, file: UploadFile, name: str) -> ConnectionOut:
         generated_name = generate_short_uuid() + ".sqlite"
@@ -201,7 +220,7 @@ async def create_excel_connection(self, session: AsyncSession, file: UploadFile,
         # Create connection with the locally copied file
         dsn = get_sqlite_dsn(str(file_path.absolute()))
         return await self.create_connection(
-            session, dsn=dsn, name=name, type=ConnectionType.excel.value, is_sample=False
+            session, dsn=dsn, name=name, connection_type=ConnectionType.excel.value, is_sample=False
         )
 
     async def create_sas7bdat_connection(self, session: AsyncSession, file: UploadFile, name: str) -> ConnectionOut:
@@ -244,7 +263,7 @@ async def create_sas7bdat_connection(self, session: AsyncSession, file: UploadFi
             # Create connection with the locally copied file
             dsn = get_sqlite_dsn(str(file_path.absolute()))
             return await self.create_connection(
-                session, dsn=dsn, name=name, type=ConnectionType.sas.value, is_sample=False
+                session, dsn=dsn, name=name, connection_type=ConnectionType.sas.value, is_sample=False
             )
         finally:
             # Clean up the temporary file

backend/dataline/services/conversation.py

@@ -114,9 +114,7 @@ async def query(
         user_with_model_details = await self.settings_service.get_model_details(session)
 
         # Create query graph
-        query_graph = QueryGraphService(
-            dsn=connection.dsn,
-        )
+        query_graph = QueryGraphService(connection=connection)
         history = await self.get_conversation_history(session, conversation_id)
 
         messages: list[BaseMessage] = []