From b8c1f84c5ef9699f13265eac91504a9c43966538 Mon Sep 17 00:00:00 2001 From: jiatolentinoDM Date: Mon, 15 Jun 2026 05:12:23 +0800 Subject: [PATCH 1/2] feat: Add support for Databricks and MongoDB connection types --- src/datamasque_cli/commands/connections.py | 15 +++++- tests/commands/test_connections.py | 60 ++++++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/src/datamasque_cli/commands/connections.py b/src/datamasque_cli/commands/connections.py index 1217c62..1957e61 100644 --- a/src/datamasque_cli/commands/connections.py +++ b/src/datamasque_cli/commands/connections.py @@ -13,7 +13,9 @@ ConnectionConfig, DatabaseConnectionConfig, DatabaseType, + DatabricksConnectionConfig, DynamoConnectionConfig, + MongoConnectionConfig, MountedShareConnectionConfig, S3ConnectionConfig, SnowflakeConnectionConfig, @@ -37,6 +39,8 @@ class ConnectionType(StrEnum): MOUNTED_SHARE = "mounted_share" SNOWFLAKE = "snowflake" DYNAMODB = "dynamodb" + DATABRICKS = "databricks" + MONGODB = "mongodb" _FILE_CONNECTION_TYPES = (MountedShareConnectionConfig, S3ConnectionConfig, AzureConnectionConfig) @@ -69,6 +73,8 @@ def _format_role(conn: ConnectionConfig) -> str: ConnectionType.MOUNTED_SHARE: MountedShareConnectionConfig, ConnectionType.SNOWFLAKE: SnowflakeConnectionConfig, ConnectionType.DYNAMODB: DynamoConnectionConfig, + ConnectionType.DATABRICKS: DatabricksConnectionConfig, + ConnectionType.MONGODB: MongoConnectionConfig, } @@ -131,7 +137,9 @@ def get_connection( def create_connection( file: Path | None = typer.Option(None, "--file", "-f", help="JSON file defining the connection"), name: str | None = typer.Option(None, help="Connection name"), - conn_type: str | None = typer.Option(None, "--type", "-t", help="database, s3, azure, mounted_share"), + conn_type: str | None = typer.Option( + None, "--type", "-t", help="database, s3, azure, mounted_share, snowflake, dynamodb, databricks, mongodb" + ), host: str | None = typer.Option(None, help="Database host"), port: str | None = typer.Option(None, help="Database port"), database: str | None = typer.Option(None, help="Database name"), @@ -160,6 +168,11 @@ def create_connection( # Quick mounted share dm connections create --name input --type mounted_share --base-dir my-data --source + + # Databricks, MongoDB, Snowflake and DynamoDB have many fields, so use --file: + # {"type": "databricks", "name": "dbx", "server_hostname": "...", "http_path": "...", + # "access_token": "...", "catalog": "main", "schema": "default"} + dm connections create --file databricks.json """ client = get_client(profile) diff --git a/tests/commands/test_connections.py b/tests/commands/test_connections.py index 82f85b7..e0c001b 100644 --- a/tests/commands/test_connections.py +++ b/tests/commands/test_connections.py @@ -7,6 +7,8 @@ from datamasque.client.models.connection import ( DatabaseConnectionConfig, DatabaseType, + DatabricksConnectionConfig, + MongoConnectionConfig, MountedShareConnectionConfig, ) from typer.testing import CliRunner @@ -164,6 +166,64 @@ def test_create_connection_from_json_file(mock_get_client: MagicMock, runner: Cl client.create_or_update_connection.assert_called_once() +@patch(f"{MODULE}.get_client") +def test_create_databricks_connection_from_json_file( + mock_get_client: MagicMock, runner: CliRunner, tmp_path: MagicMock +) -> None: + client = MagicMock() + mock_get_client.return_value = client + + conn_file = tmp_path / "dbx.json" + conn_file.write_text( + json.dumps( + { + "type": "databricks", + "name": "dbx", + "server_hostname": "dbc-1514b142-1c6c.cloud.databricks.com", + "http_path": "/sql/1.0/warehouses/ea7d918dd5f236f9", + "access_token": "dapiTOKEN", + "catalog": "main", + "schema": "default", + } + ) + ) + + result = runner.invoke(app, ["connections", "create", "--file", str(conn_file)]) + assert result.exit_code == 0 + config = client.create_or_update_connection.call_args.args[0] + assert isinstance(config, DatabricksConnectionConfig) + assert config.db_type == "databricks" + + +@patch(f"{MODULE}.get_client") +def test_create_mongodb_connection_from_json_file( + mock_get_client: MagicMock, runner: CliRunner, tmp_path: MagicMock +) -> None: + client = MagicMock() + mock_get_client.return_value = client + + conn_file = tmp_path / "mongo.json" + conn_file.write_text( + json.dumps( + { + "type": "mongodb", + "name": "mongo", + "host": "localhost", + "port": "27017", + "database": "mydb", + "user": "admin", + "password": "secret", + } + ) + ) + + result = runner.invoke(app, ["connections", "create", "--file", str(conn_file)]) + assert result.exit_code == 0 + config = client.create_or_update_connection.call_args.args[0] + assert isinstance(config, MongoConnectionConfig) + assert config.db_type == "mongodb" + + # -- delete (tests confirmation logic) ------------------------------------ From b51dea7633c4aad7833455c6ffaebb6167252a92 Mon Sep 17 00:00:00 2001 From: jiatolentinoDM Date: Mon, 15 Jun 2026 05:33:09 +0800 Subject: [PATCH 2/2] chore: Bump version to 1.4.0 --- CHANGELOG.md | 7 +++++++ pyproject.toml | 2 +- uv.lock | 2 +- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f9977c..3205275 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## v1.4.0 + +### Added +- `dm connections create --file` now supports Databricks SQL Warehouse + (`"type": "databricks"`) and MongoDB (`"type": "mongodb"`) connections. + Both list, get, create, and delete like the existing connection types. + ## v1.3.0 ### Added diff --git a/pyproject.toml b/pyproject.toml index 4545528..b4dfa47 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "datamasque-cli" -version = "1.3.0" +version = "1.4.0" description = "Official command-line interface for the DataMasque data-masking platform." authors = [ { name = "DataMasque Ltd" }, diff --git a/uv.lock b/uv.lock index 11c3709..13b056f 100644 --- a/uv.lock +++ b/uv.lock @@ -141,7 +141,7 @@ wheels = [ [[package]] name = "datamasque-cli" -version = "1.3.0" +version = "1.4.0" source = { editable = "." } dependencies = [ { name = "datamasque-python" },