added databricks integration

.github/workflows/run_snowflake_cli_tests.yml

@@ -1,17 +1,27 @@
-name: Run Snowflake CLI Tests
+name: Run Database CLI Tests
+run-name: Run CLI Tests (${{ inputs.database }})
 
-# This workflow is designed to be run manually, as it executes the complete suite of CLI tests against a snowflake database and may take a while to complete.
+# This workflow is designed to be run manually, as it executes the complete suite of CLI tests against the selected database and may take a while to complete.
 
 
 on:
-  workflow_dispatch: {}
+  workflow_dispatch:
+    inputs:
+      database:
+        description: "Database to test against"
+        required: true
+        default: "snowflake"
+        type: choice
+        options:
+          - snowflake
+          - databricks
 
 jobs:
-  snowflake-cli-tests:
+  database-cli-tests:
     runs-on: ubuntu-latest
 
     env:
-      JETBASE_SQLALCHEMY_URL: ${{ secrets.SNOWFLAKE_USER_PW_SQLALCHEMY_URL }}
+      JETBASE_SQLALCHEMY_URL: ${{ inputs.database == 'snowflake' && secrets.SNOWFLAKE_USER_PW_SQLALCHEMY_URL || secrets.DATABRICKS_SQLALCHEMY_URL }}
 
     steps:
       - uses: actions/checkout@v5
@@ -25,8 +35,7 @@ jobs:
         uses: astral-sh/setup-uv@v6
 
       - name: Install dependencies
-        run: uv sync --dev
+        run: uv sync --dev --extra ${{ inputs.database }}
 
       - name: Run CLI tests
         run: uv run pytest tests/cli/ -v
-

README.md

@@ -37,11 +37,12 @@ pip install jetbase
 uv add jetbase
 ```
 
-> **Note for Snowflake Users:**  
-> To use Jetbase with Snowflake, install the Snowflake extras:
+> **Note for Snowflake and Databricks Users:**  
+> To use Jetbase with Snowflake or Databricks, install the appropriate extras:
 >
 > ```shell
 > pip install "jetbase[snowflake]"
+> pip install "jetbase[databricks]"
 > ```
 
 
@@ -132,6 +133,7 @@ Jetbase currently supports:
 - ✅ PostgreSQL
 - ✅ SQLite
 - ✅ Snowflake
+- ✅ Databricks
 - ✅ MySQL
 
 ## Need Help?

jetbase/database/queries/databricks.py

@@ -0,0 +1,108 @@
+from sqlalchemy import TextClause, text
+
+from jetbase.database.queries.base import BaseQueries
+
+
+class DatabricksQueries(BaseQueries):
+    """
+    Databricks-specific SQL queries.
+
+    Provides Databricks-compatible implementations for queries that differ
+    from the default PostgreSQL syntax.
+    """
+
+    @staticmethod
+    def create_migrations_table_stmt() -> TextClause:
+        return text(
+            """
+            CREATE TABLE IF NOT EXISTS jetbase_migrations (
+                order_executed BIGINT GENERATED ALWAYS AS IDENTITY,
+                version STRING,
+                description STRING NOT NULL,
+                filename STRING NOT NULL,
+                migration_type STRING NOT NULL,
+                applied_at TIMESTAMP NOT NULL,
+                checksum STRING NOT NULL
+            )
+            """
+        )
+
+    @staticmethod
+    def insert_version_stmt() -> TextClause:
+        return text(
+            """
+            INSERT INTO jetbase_migrations (version, description, filename, migration_type, checksum, applied_at) 
+            VALUES (:version, :description, :filename, :migration_type, :checksum, CURRENT_TIMESTAMP())
+            """
+        )
+
+    @staticmethod
+    def check_if_migrations_table_exists_query() -> TextClause:
+        return text(
+            """
+            SELECT COUNT(*) > 0 AS table_exists
+            FROM INFORMATION_SCHEMA.TABLES
+            WHERE TABLE_SCHEMA = CURRENT_SCHEMA()
+              AND LOWER(TABLE_NAME) = 'jetbase_migrations'
+            """
+        )
+
+    @staticmethod
+    def check_if_lock_table_exists_query() -> TextClause:
+        return text(
+            """
+            SELECT COUNT(*) > 0 AS table_exists
+            FROM INFORMATION_SCHEMA.TABLES
+            WHERE TABLE_SCHEMA = CURRENT_SCHEMA()
+              AND LOWER(TABLE_NAME) = 'jetbase_lock'
+            """
+        )
+
+    @staticmethod
+    def create_lock_table_stmt() -> TextClause:
+        return text(
+            """
+            CREATE TABLE IF NOT EXISTS jetbase_lock (
+                id INT,
+                is_locked BOOLEAN NOT NULL,
+                locked_at TIMESTAMP,
+                process_id STRING
+            )
+            """
+        )
+
+    @staticmethod
+    def initialize_lock_record_stmt() -> TextClause:
+        return text(
+            """
+            MERGE INTO jetbase_lock AS target
+            USING (SELECT 1 AS id, FALSE AS is_locked) AS source
+            ON target.id = source.id
+            WHEN NOT MATCHED THEN
+                INSERT (id, is_locked) VALUES (source.id, source.is_locked)
+            """
+        )
+
+    @staticmethod
+    def acquire_lock_stmt() -> TextClause:
+        return text(
+            """
+            UPDATE jetbase_lock
+            SET is_locked = TRUE,
+                locked_at = CURRENT_TIMESTAMP(),
+                process_id = :process_id
+            WHERE id = 1 AND is_locked = FALSE
+            """
+        )
+
+    @staticmethod
+    def update_repeatable_migration_stmt() -> TextClause:
+        return text(
+            """
+            UPDATE jetbase_migrations
+            SET checksum = :checksum,
+                applied_at = CURRENT_TIMESTAMP()
+            WHERE filename = :filename
+            AND migration_type = :migration_type
+            """
+        )

jetbase/database/queries/query_loader.py

@@ -2,10 +2,11 @@
 
 from jetbase.config import get_config
 from jetbase.database.queries.base import BaseQueries, QueryMethod
+from jetbase.database.queries.databricks import DatabricksQueries
 from jetbase.database.queries.mysql import MySQLQueries
 from jetbase.database.queries.postgres import PostgresQueries
-from jetbase.database.queries.sqlite import SQLiteQueries
 from jetbase.database.queries.snowflake import SnowflakeQueries
+from jetbase.database.queries.sqlite import SQLiteQueries
 from jetbase.enums import DatabaseType
 
 
@@ -34,6 +35,8 @@ def get_database_type() -> DatabaseType:
         return DatabaseType.SNOWFLAKE
     elif dialect_name == "mysql":
         return DatabaseType.MYSQL
+    elif dialect_name == "databricks":
+        return DatabaseType.DATABRICKS
     else:
         raise ValueError(f"Unsupported database type: {dialect_name}")
 
@@ -61,6 +64,8 @@ def get_queries() -> type[BaseQueries]:
         return SnowflakeQueries
     elif db_type == DatabaseType.MYSQL:
         return MySQLQueries
+    elif db_type == DatabaseType.DATABRICKS:
+        return DatabricksQueries
     else:
         raise ValueError(f"Unsupported database type: {db_type}")
 

jetbase/enums.py

@@ -23,3 +23,4 @@ class DatabaseType(Enum):
     SQLITE = "sqlite"
     SNOWFLAKE = "snowflake"
     MYSQL = "mysql"
+    DATABRICKS = "databricks"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jetbase"
-version = "0.16.0"
+version = "0.17.0"
 description = "Jetbase is a Python database migration tool"
 readme = "README.md"
 authors = [
@@ -41,6 +41,9 @@ snowflake = [
     "snowflake-sqlalchemy>=1.8.2",
     "cryptography>=46.0.3",
 ]
+databricks = [
+    "databricks-sqlalchemy>=2.0.1",
+]
 
 
 [tool.pytest.ini_options]

tests/cli/migrations_databricks/RA__ra.sql

@@ -0,0 +1 @@
+INSERT INTO users (name) VALUES ('always mike');

tests/cli/migrations_databricks/ROC__roc.sql

@@ -0,0 +1 @@
+INSERT INTO users (name) VALUES ('on change mike');

tests/cli/migrations_databricks/V1__m1.sql

@@ -0,0 +1,4 @@
+CREATE TABLE IF NOT EXISTS users (
+    id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
+    name VARCHAR(255) NOT NULL
+);

tests/cli/migrations_databricks/V21__mi21.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike21');
+
+--rollback
+DELETE FROM users WHERE name = 'mike21';

tests/cli/migrations_databricks/V2__m2.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike2');
+
+--rollback
+DELETE FROM users WHERE name = 'mike2';

tests/cli/migrations_databricks/V3__m3.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike3');
+
+--rollback
+DELETE FROM users WHERE name = 'mike3';

tests/cli/migrations_databricks/V4__m4.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike4');
+
+--rollback
+DELETE FROM users WHERE name = 'mike4';

tests/cli/migrations_databricks_versions_only/V1__m1.sql

@@ -0,0 +1,4 @@
+CREATE TABLE IF NOT EXISTS users (
+    id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
+    name VARCHAR(255) NOT NULL
+);

tests/cli/migrations_databricks_versions_only/V21__mi21.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike21');
+
+--rollback
+DELETE FROM users WHERE name = 'mike21';

tests/cli/migrations_databricks_versions_only/V2__m2.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike2');
+
+--rollback
+DELETE FROM users WHERE name = 'mike2';

tests/cli/migrations_databricks_versions_only/V3__m3.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike3');
+
+--rollback
+DELETE FROM users WHERE name = 'mike3';

tests/cli/migrations_databricks_versions_only/V4__m4.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name) VALUES ('mike4');
+
+--rollback
+DELETE FROM users WHERE name = 'mike4';

tests/conftest.py

@@ -5,8 +5,9 @@
 import pytest
 from sqlalchemy import create_engine, text
 from typer.testing import CliRunner
-from jetbase.enums import DatabaseType
+
 from jetbase.database.queries.base import detect_db
+from jetbase.enums import DatabaseType
 
 
 @pytest.fixture
@@ -46,6 +47,9 @@ def migrations_fixture_dir(test_db_url):
     if detect_db(test_db_url) == DatabaseType.MYSQL:
         return base_path / "migrations_mysql"
 
+    if detect_db(test_db_url) == DatabaseType.DATABRICKS:
+        return base_path / "migrations_databricks"
+
     return base_path / "migrations"
 
 
@@ -60,6 +64,9 @@ def migrations_versions_only_fixture_dir(test_db_url):
     if detect_db(test_db_url) == DatabaseType.MYSQL:
         return base_path / "migrations_mysql_versions_only"
 
+    if detect_db(test_db_url) == DatabaseType.DATABRICKS:
+        return base_path / "migrations_databricks_versions_only"
+
     return base_path / "migrations_versions_only"