mito-ai 0.1.50__py3-none-any.whl

mito_ai/db/handlers.py

@@ -0,0 +1,168 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import os
+import tornado
+import uuid
+from typing import Any, Final
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.utils.schema import MITO_FOLDER
+from mito_ai.utils.telemetry_utils import (
+    log_db_connection_attempt,
+    log_db_connection_success,
+    log_db_connection_error,
+)
+from mito_ai.db.utils import (
+    setup_database_dir,
+    save_connection,
+    delete_connection,
+    delete_schema,
+    crawl_and_store_schema,
+    install_db_drivers,
+)
+
+DB_DIR_PATH: Final[str] = os.path.join(MITO_FOLDER, "db")
+CONNECTIONS_PATH: Final[str] = os.path.join(DB_DIR_PATH, "connections.json")
+SCHEMAS_PATH: Final[str] = os.path.join(DB_DIR_PATH, "schemas.json")
+
+
+class ConnectionsHandler(APIHandler):
+    """
+    Endpoints for working with connections.json file.
+    """
+
+    @tornado.web.authenticated
+    def get(self) -> None:
+        """Get all connections."""
+        # If the db dir doesn't exist, create it
+        setup_database_dir(DB_DIR_PATH, CONNECTIONS_PATH, SCHEMAS_PATH)
+
+        with open(CONNECTIONS_PATH, "r") as f:
+            connections = json.load(f)
+        self.finish(json.dumps(connections))
+
+    @tornado.web.authenticated
+    def post(self) -> None:
+        """Add a new connection."""
+        try:
+            # If the db dir doesn't exist, create it
+            setup_database_dir(DB_DIR_PATH, CONNECTIONS_PATH, SCHEMAS_PATH)
+
+            # Get the new connection data from the request body
+            connection_details = json.loads(self.request.body)
+
+            # Generate a UUID for the new connection
+            connection_id = str(uuid.uuid4())
+
+            db_type = connection_details["type"]
+            log_db_connection_attempt(db_type)
+
+            # Install database drivers
+            install_result = install_db_drivers(db_type)
+            if not install_result["success"]:
+                log_db_connection_error(db_type, install_result["error"])
+                self.set_status(500)
+                self.write({"error": install_result["error"]})
+                return
+
+            # First, try to validate the connection by building the schema
+            crawl_result = crawl_and_store_schema(
+                SCHEMAS_PATH,
+                connection_id,
+                connection_details,
+            )
+
+            if not crawl_result["success"]:
+                log_db_connection_error(
+                    connection_details["type"], crawl_result["error_message"]
+                )
+                self.set_status(500)
+                self.write({"error": crawl_result["error_message"]})
+                return
+
+            # If schema building succeeded, save the connection
+            save_connection(CONNECTIONS_PATH, connection_id, connection_details)
+
+            log_db_connection_success(connection_details["type"], {})
+
+            self.write(
+                {
+                    "status": "success",
+                    "message": "Added new connection",
+                    "connection_id": connection_id,
+                }
+            )
+
+        except json.JSONDecodeError:
+            self.set_status(400)
+            self.write({"error": "Invalid JSON in request body"})
+        except Exception as e:
+            self.set_status(500)
+            self.write({"error": str(e)})
+        finally:
+            self.finish()
+
+    @tornado.web.authenticated
+    def delete(self, *args: Any, **kwargs: Any) -> None:
+        """Delete a connection by UUID."""
+        try:
+            # Get the connection UUID from the URL
+            connection_id = kwargs.get("uuid")
+            if not connection_id:
+                self.set_status(400)
+                self.write({"error": "Connection UUID is required"})
+                return
+
+            # Delete the connection
+            delete_connection(CONNECTIONS_PATH, connection_id)
+
+            # Delete the schema
+            delete_schema(SCHEMAS_PATH, connection_id)
+
+            self.set_status(200)
+            self.write(
+                {
+                    "status": "success",
+                    "message": "Connection deleted successfully",
+                }
+            )
+
+        except Exception as e:
+            self.set_status(500)
+            self.write({"error": str(e)})
+        finally:
+            self.finish()
+
+
+class SchemaHandler(APIHandler):
+    """
+    Endpoints for working with schemas.json file.
+    """
+
+    @tornado.web.authenticated
+    def get(self) -> None:
+        """Get all schemas."""
+        with open(SCHEMAS_PATH, "r") as f:
+            schemas = json.load(f)
+
+        self.write(schemas)
+        self.finish()
+
+    @tornado.web.authenticated
+    def delete(self, *args: Any, **kwargs: Any) -> None:
+        """Delete a schema by UUID."""
+        # Get the schema UUID from kwargs
+        schema_id = kwargs.get("uuid")
+
+        if not schema_id:
+            self.set_status(400)
+            self.write({"error": "Schema UUID is required"})
+            self.finish()
+            return
+
+        delete_schema(SCHEMAS_PATH, schema_id)
+
+        self.set_status(200)
+        self.write({"status": "success", "message": "Schema deleted successfully"})
+        self.finish()

mito_ai/db/models.py

@@ -0,0 +1,31 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import TypedDict, List, Dict
+
+
+# BASE CRAWLER MODELS
+
+
+class ColumnInfo(TypedDict):
+    name: str
+    type: str
+
+
+class TableSchema(TypedDict):
+    tables: Dict[str, List[ColumnInfo]]
+
+
+# SNOWFLAKE MODELS
+
+
+class SchemaInfo(TypedDict):
+    tables: Dict[str, List[ColumnInfo]]
+
+
+class DatabaseInfo(TypedDict):
+    schemas: Dict[str, SchemaInfo]
+
+
+class WarehouseDetails(TypedDict):
+    databases: Dict[str, DatabaseInfo]

mito_ai/db/urls.py

@@ -0,0 +1,34 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Any
+from jupyter_server.utils import url_path_join
+from mito_ai.db.handlers import ConnectionsHandler, SchemaHandler
+
+
+def get_db_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all database related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai/db"
+    return [
+        (url_path_join(BASE_URL, "connections"), ConnectionsHandler, {}),
+        (
+            # URL for deleting a connection.
+            url_path_join(BASE_URL, "connections", "(?P<uuid>[^/]+)"),
+            ConnectionsHandler,
+            {},
+        ),
+        (url_path_join(BASE_URL, "schemas"), SchemaHandler, {}),
+        (
+            # URL for deleting a schema.
+            url_path_join(BASE_URL, "schemas", "(?P<uuid>[^/]+)"),
+            SchemaHandler,
+            {},
+        ),
+    ]

mito_ai/db/utils.py

@@ -0,0 +1,185 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import os
+from mito_ai.db.crawlers import snowflake, base_crawler
+from mito_ai.db.crawlers.constants import SUPPORTED_DATABASES
+
+
+def setup_database_dir(
+    db_dir_path: str, connections_path: str, schemas_path: str
+) -> None:
+    """
+    Setup the database directory.
+    """
+
+    # Ensure the db directory exists
+    os.makedirs(db_dir_path, exist_ok=True)
+
+    # Create connections.json if it doesn't exist
+    if not os.path.exists(connections_path):
+        with open(connections_path, "w") as f:
+            json.dump({}, f, indent=4)
+
+    # Create schemas.json if it doesn't exist
+    if not os.path.exists(schemas_path):
+        with open(schemas_path, "w") as f:
+            json.dump({}, f, indent=4)
+
+
+def save_connection(
+    connections_path: str, connection_id: str, connection_details: dict
+) -> None:
+    """
+    Save a connection to the connections.json file.
+
+    Args:
+        connections_path (str): The path to the connections.json file.
+        connection_id (str): The UUID of the connection to save.
+        connection_details (dict): The details of the connection to save.
+    """
+
+    with open(connections_path, "r") as f:
+        connections = json.load(f)
+
+    # Add the new connection
+    connections[connection_id] = connection_details
+
+    # Write back to file
+    with open(connections_path, "w") as f:
+        json.dump(connections, f, indent=4)
+
+
+def delete_connection(connections_path: str, connection_id: str) -> None:
+    """
+    Delete a connection by UUID.
+    """
+
+    # Read existing connections
+    with open(connections_path, "r") as f:
+        connections = json.load(f)
+
+    # Remove the connection
+    del connections[connection_id]
+
+    # Write back to file
+    with open(connections_path, "w") as f:
+        json.dump(connections, f, indent=4)
+
+
+def delete_schema(schemas_path: str, schema_id: str) -> None:
+    """
+    Delete a schema by UUID.
+
+    Args:
+        schemas_path (str): The path to the schemas.json file.
+        schema_id (str): The UUID of the schema to delete.
+    """
+
+    with open(schemas_path, "r") as f:
+        schemas = json.load(f)
+
+    del schemas[schema_id]
+
+    with open(schemas_path, "w") as f:
+        json.dump(schemas, f, indent=4)
+
+
+def crawl_and_store_schema(
+    schemas_path: str,
+    connection_id: str,
+    connection_details: dict,
+) -> dict:
+    """
+    Crawl and store schema for a given connection.
+
+    Args:
+        schemas_path (str): The path to the schemas.json file.
+        connection_id (str): The UUID of the connection to crawl.
+        username (str): The username for the connection.
+        password (str): The password for the connection.
+        account (str): The account for the connection.
+        warehouse (str): The warehouse for the connection.
+
+    Returns:
+        tuple[bool, str]: A tuple containing a boolean indicating success and an error message.
+    """
+    if connection_details["type"] == "snowflake":
+        schema = snowflake.crawl_snowflake(
+            connection_details["username"],
+            connection_details["password"],
+            connection_details["account"],
+            connection_details["warehouse"],
+        )
+    elif connection_details["type"] == "postgres":
+        conn_str = f"postgresql+psycopg2://{connection_details['username']}:{connection_details['password']}@{connection_details['host']}:{connection_details['port']}/{connection_details['database']}"
+        schema = base_crawler.crawl_db(conn_str, "postgres")
+    elif connection_details["type"] == "sqlite":
+        conn_str = f"sqlite:///{connection_details['database']}"
+        schema = base_crawler.crawl_db(conn_str, "sqlite")
+    elif connection_details["type"] == "mysql":
+        conn_str = f"mysql+pymysql://{connection_details['username']}:{connection_details['password']}@{connection_details['host']}:{connection_details['port']}/{connection_details['database']}"
+        schema = base_crawler.crawl_db(conn_str, "mysql")
+    elif connection_details["type"] == "mssql":
+        odbc_driver_version = connection_details["odbc_driver_version"]
+        conn_str = f"mssql+pyodbc://{connection_details['username']}:{connection_details['password']}@{connection_details['host']}:{connection_details['port']}/{connection_details['database']}?driver=ODBC+Driver+{odbc_driver_version}+for+SQL+Server"
+        schema = base_crawler.crawl_db(conn_str, "mssql")
+    elif connection_details["type"] == "oracle":
+        conn_str = f"oracle+oracledb://{connection_details['username']}:{connection_details['password']}@{connection_details['host']}:{connection_details['port']}?service_name={connection_details['service_name']}"
+        schema = base_crawler.crawl_db(conn_str, "oracle")
+
+    if schema["error"]:
+        return {
+            "success": False,
+            "error_message": schema["error"],
+            "schema": {},
+        }
+
+    # If we successfully crawled the schema, write it to schemas.json
+    with open(schemas_path, "r+") as f:
+        # Load the existing schemas
+        schemas = json.load(f)
+        # Remove the error key from the schema and add the crawled schema
+        schema.pop("error", None)
+        schemas[connection_id] = schema["schema"]
+        # Move to the beginning of the file and write the new schema
+        f.seek(0)
+        json.dump(schemas, f, indent=4)
+        f.truncate()
+    return {
+        "success": True,
+        "error_message": "",
+        "schema": schema,
+    }
+
+
+def install_db_drivers(db_type: str) -> dict:
+    """
+    Install required database drivers for the given database type.
+
+    Args:
+        db_type (str): The type of database (e.g. 'snowflake', 'postgres')
+
+    Returns:
+        dict: A dictionary containing success status and error message if any
+    """
+    from mito_ai.utils.utils import get_installed_packages, install_packages
+
+    installed_packages = get_installed_packages()
+    required_packages = SUPPORTED_DATABASES[db_type].get("drivers", [])
+    packages_to_install = []
+
+    for package in required_packages:
+        if package not in installed_packages:
+            packages_to_install.append(package)
+
+    if len(packages_to_install) > 0:
+        install_result = install_packages(packages_to_install)
+        if not install_result["success"]:
+            return {
+                "success": False,
+                "error": f"Failed to install {db_type} drivers: {install_result['error']}",
+            }
+
+    return {"success": True, "error": None}

mito_ai/docker/mssql/compose.yml

@@ -0,0 +1,37 @@
+version: '3.8'
+
+services:
+  db:
+    image: mcr.microsoft.com/mssql/server:2022-latest
+    container_name: mssql_test_db
+    ports:
+      - "1433:1433"
+    environment:
+      # Default user is SA (System Administrator)
+      SA_PASSWORD: "YourStrong!Passw0rd" # Password must meet complexity requirements
+      ACCEPT_EULA: "Y"
+    volumes:
+      - mssql_data:/var/opt/mssql
+    restart: unless-stopped
+
+  # Run setup.sql 
+  sqltools:
+    image: mcr.microsoft.com/mssql-tools
+    depends_on:
+      - db
+    volumes:
+      - ./init:/init
+    entrypoint: >
+      /bin/bash -c "
+        echo '⏳ Waiting for SQL Server to be ready...';
+        for i in {1..30}; do
+          /opt/mssql-tools/bin/sqlcmd -S db -U SA -P 'YourStrong!Passw0rd' -Q 'SELECT 1' && break
+          sleep 2
+        done;
+        echo '🚀 Running init script...';
+        /opt/mssql-tools/bin/sqlcmd -S db -U SA -P 'YourStrong!Passw0rd' -i /init/setup.sql
+      "
+    restart: "no"
+
+volumes:
+  mssql_data:

mito_ai/docker/mssql/init/setup.sql

@@ -0,0 +1,21 @@
+-- Copyright (c) Saga Inc.
+-- Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+CREATE DATABASE test_db;
+GO
+
+USE test_db;
+GO
+
+CREATE TABLE users (
+  id INT IDENTITY(1,1) PRIMARY KEY,
+  name NVARCHAR(255) NOT NULL,
+  email NVARCHAR(255) UNIQUE NOT NULL
+);
+GO
+
+INSERT INTO users (name, email) VALUES
+  ('Alice', 'alice@example.com'),
+  ('Bob', 'bob@example.com'),
+  ('Charlie', 'charlie@example.com');
+GO

mito_ai/docker/mysql/compose.yml

@@ -0,0 +1,18 @@
+services:
+  db:
+    image: mysql:8
+    container_name: mysql_test_db
+    ports:
+      - "3306:3306"  # Maps container port 3306 to host port 3306
+    environment:
+      MYSQL_ROOT_PASSWORD: root_pass
+      MYSQL_USER: test_user
+      MYSQL_PASSWORD: test_pass
+      MYSQL_DATABASE: test_db
+    volumes:
+      - mysql_data:/var/lib/mysql                    # persistent data volume
+      - ./init:/docker-entrypoint-initdb.d           # SQL files to run on init
+    restart: unless-stopped
+
+volumes:
+  mysql_data:

mito_ai/docker/mysql/init/setup.sql

@@ -0,0 +1,13 @@
+-- Copyright (c) Saga Inc.
+-- Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+CREATE TABLE users (
+  id INT AUTO_INCREMENT PRIMARY KEY,
+  name VARCHAR(255) NOT NULL,
+  email VARCHAR(255) UNIQUE NOT NULL
+);
+
+INSERT INTO users (name, email) VALUES
+  ('Alice', 'alice@example.com'),
+  ('Bob', 'bob@example.com'),
+  ('Charlie', 'charlie@example.com');

mito_ai/docker/oracle/compose.yml

@@ -0,0 +1,17 @@
+services:
+  db:
+    image: gvenzl/oracle-xe:21.3.0-slim
+    container_name: oracle_test_db
+    ports:
+      - "1521:1521"
+    environment:
+      ORACLE_PASSWORD: test_pass
+      APP_USER: test_user
+      APP_USER_PASSWORD: test_pass
+    volumes:
+      - oracle_data:/opt/oracle/oradata
+      - ./init:/container-entrypoint-startdb.d
+    restart: unless-stopped
+
+volumes:
+  oracle_data:

mito_ai/docker/oracle/init/setup.sql

@@ -0,0 +1,20 @@
+-- Copyright (c) Saga Inc.
+-- Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+connect  test_user/test_pass@xepdb1;
+
+create sequence hibernate_sequence start with 1 increment by 1;
+
+create table department (id integer not null, name varchar(255), primary key (id));
+
+create table employee (id integer not null, email varchar(255), first_name varchar(255), gender varchar(255), last_name varchar(255), salary numeric(19,2), department_id integer, primary key (id));
+
+insert into department (id, name) values (1, 'IT');
+insert into department (id, name) values (2, 'HR');
+insert into department (id, name) values (3, 'Finance');
+
+insert into employee (id, email, first_name, gender, last_name, salary, department_id) values (1, 'john.doe@example.com', 'John', 'M', 'Doe', 50000, 1);
+insert into employee (id, email, first_name, gender, last_name, salary, department_id) values (2, 'jane.smith@example.com', 'Jane', 'F', 'Smith', 55000, 2);
+insert into employee (id, email, first_name, gender, last_name, salary, department_id) values (3, 'jim.beam@example.com', 'Jim', 'M', 'Beam', 60000, 3);
+
+COMMIT;

mito_ai/docker/postgres/compose.yml

@@ -0,0 +1,17 @@
+services:
+  db:
+    image: postgres:17
+    container_name: pg_test_db
+    ports:
+      - "5432:5432"  # Maps container port 5432 to host port 5432
+    environment:
+      POSTGRES_USER: test_user
+      POSTGRES_PASSWORD: test_pass
+      POSTGRES_DB: test_db
+    volumes:
+      - pg_data:/var/lib/postgresql/data        # persistent data volume
+      - ./init:/docker-entrypoint-initdb.d      # SQL files to run on init
+    restart: unless-stopped
+
+volumes:
+  pg_data:

mito_ai/docker/postgres/init/setup.sql

@@ -0,0 +1,13 @@
+-- Copyright (c) Saga Inc.
+-- Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+CREATE TABLE users (
+  id SERIAL PRIMARY KEY,
+  name TEXT NOT NULL,
+  email TEXT UNIQUE NOT NULL
+);
+
+INSERT INTO users (name, email) VALUES
+  ('Alice', 'alice@example.com'),
+  ('Bob', 'bob@example.com'),
+  ('Charlie', 'charlie@example.com');

mito_ai/enterprise/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+

mito_ai/enterprise/utils.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the The Mito Enterprise license.
+
+from mito_ai.utils.version_utils import is_enterprise, is_mitosheet_private
+from mito_ai.constants import AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL
+
+def is_azure_openai_configured() -> bool:
+    """
+    Azure OpenAI is only supported for Mito Enterprise users
+    """
+    is_allowed_to_use_azure = is_enterprise() or is_mitosheet_private()
+    return all([is_allowed_to_use_azure, AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL])

mito_ai/file_uploads/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+