Merge pull request #1 from bruin-data/feature/v0.3.0-new-connections

feat: add Fabric, Oracle, DB2, HANA, Spanner, Vertica connections (v0.3.0)

Author: sabrikaragonen

.github/workflows/release.yml

@@ -0,0 +1,47 @@
+name: Auto Release
+
+on:
+  push:
+    branches: [main]
+
+jobs:
+  release:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Get version from pyproject.toml
+        id: version
+        run: |
+          VERSION=$(python3 -c "
+          import re
+          with open('pyproject.toml') as f:
+              match = re.search(r'version\s*=\s*\"(.+?)\"', f.read())
+              print(match.group(1))
+          ")
+          echo "version=$VERSION" >> "$GITHUB_OUTPUT"
+          echo "tag=v$VERSION" >> "$GITHUB_OUTPUT"
+
+      - name: Check if tag exists
+        id: check
+        run: |
+          if git rev-parse "v${{ steps.version.outputs.version }}" >/dev/null 2>&1; then
+            echo "exists=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "exists=false" >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: Create tag and release
+        if: steps.check.outputs.exists == 'false'
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          git tag "v${{ steps.version.outputs.version }}"
+          git push origin "v${{ steps.version.outputs.version }}"
+          gh release create "v${{ steps.version.outputs.version }}" \
+            --title "v${{ steps.version.outputs.version }}" \
+            --generate-notes

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "bruin-sdk"
-version = "0.2.1"
+version = "0.3.0"
 description = "Python SDK for Bruin CLI — query databases, parse context, and access connections with zero boilerplate."
 readme = "docs/README.md"
 license = "Apache-2.0"
@@ -42,6 +42,12 @@ athena = ["pyathena"]
 trino = ["trino"]
 motherduck = ["duckdb"]
 synapse = ["pymssql"]
+fabric = ["pymssql"]
+oracle = ["oracledb"]
+db2 = ["ibm-db"]
+hana = ["hdbcli"]
+spanner = ["google-cloud-spanner", "google-auth"]
+vertica = ["vertica-python"]
 all = [
     "bruin-sdk[bigquery]",
     "bruin-sdk[snowflake]",
@@ -57,6 +63,12 @@ all = [
     "bruin-sdk[trino]",
     "bruin-sdk[motherduck]",
     "bruin-sdk[synapse]",
+    "bruin-sdk[fabric]",
+    "bruin-sdk[oracle]",
+    "bruin-sdk[db2]",
+    "bruin-sdk[hana]",
+    "bruin-sdk[spanner]",
+    "bruin-sdk[vertica]",
 ]
 dev = ["pytest", "pytest-cov"]
 

src/bruin/__init__.py

@@ -4,5 +4,5 @@
 from bruin._context import context
 from bruin._query import query
 
-__version__ = "0.2.1"
+__version__ = "0.3.0"
 __all__ = ["__version__", "context", "get_connection", "query"]

src/bruin/_connection.py

@@ -178,9 +178,10 @@ def _create_client(conn_type: str, raw):
     factories = {
         "snowflake": _create_snowflake,
         "postgres": _create_postgres,
-        "redshift": _create_postgres,
+        "redshift": _create_redshift,
         "mssql": _create_mssql,
         "synapse": _create_mssql,
+        "fabric": _create_mssql,
         "mysql": _create_mysql,
         "duckdb": _create_duckdb,
         "databricks": _create_databricks,
@@ -189,6 +190,11 @@ def _create_client(conn_type: str, raw):
         "trino": _create_trino,
         "sqlite": _create_sqlite,
         "motherduck": _create_motherduck,
+        "oracle": _create_oracle,
+        "db2": _create_db2,
+        "hana": _create_hana,
+        "spanner": _create_spanner,
+        "vertica": _create_vertica,
     }
     factory = factories.get(conn_type)
     if factory is None:
@@ -249,7 +255,25 @@ def _create_postgres(raw: dict):
         dbname=raw.get("database", ""),
         user=raw["username"],
         password=raw["password"],
-        sslmode=raw.get("ssl_mode", "disable"),
+        sslmode=raw.get("ssl_mode", "allow"),
+    )
+
+
+def _create_redshift(raw: dict):
+    try:
+        import psycopg2
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[redshift] to use Redshift connections: "
+            "pip install 'bruin-sdk[redshift]'"
+        )
+    return psycopg2.connect(
+        host=raw["host"],
+        port=raw.get("port", 5439),
+        dbname=raw.get("database", ""),
+        user=raw["username"],
+        password=raw["password"],
+        sslmode=raw.get("ssl_mode", "allow"),
     )
 
 
@@ -398,6 +422,107 @@ def _create_motherduck(raw: dict):
     return conn
 
 
+def _create_oracle(raw: dict):
+    try:
+        import oracledb
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[oracle] to use Oracle connections: "
+            "pip install 'bruin-sdk[oracle]'"
+        )
+    kwargs = {
+        "user": raw["username"],
+        "password": raw["password"],
+        "host": raw["host"],
+        "port": int(raw.get("port", 1521)),
+    }
+    if raw.get("service_name"):
+        kwargs["service_name"] = raw["service_name"]
+    elif raw.get("sid"):
+        kwargs["sid"] = raw["sid"]
+    return oracledb.connect(**kwargs)
+
+
+def _create_db2(raw: dict):
+    try:
+        import ibm_db_dbi
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[db2] to use DB2 connections: "
+            "pip install 'bruin-sdk[db2]'"
+        )
+    conn_str = (
+        f"DATABASE={raw.get('database', '')};"
+        f"HOSTNAME={raw['host']};"
+        f"PORT={raw.get('port', 50000)};"
+        f"PROTOCOL=TCPIP;"
+        f"UID={raw['username']};"
+        f"PWD={raw['password']};"
+    )
+    return ibm_db_dbi.connect(conn_str, "", "")
+
+
+def _create_hana(raw: dict):
+    try:
+        from hdbcli import dbapi as hana_dbapi
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[hana] to use SAP HANA connections: "
+            "pip install 'bruin-sdk[hana]'"
+        )
+    return hana_dbapi.connect(
+        address=raw["host"],
+        port=int(raw.get("port", 30015)),
+        user=raw["username"],
+        password=raw["password"],
+        databaseName=raw.get("database", ""),
+    )
+
+
+def _create_spanner(raw: dict):
+    try:
+        from google.cloud.spanner_dbapi import connect as spanner_connect
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[spanner] to use Cloud Spanner connections: "
+            "pip install 'bruin-sdk[spanner]'"
+        )
+    kwargs = {
+        "instance_id": raw.get("instance_id", ""),
+        "database_id": raw.get("database", ""),
+        "project": raw.get("project_id", ""),
+    }
+    sa_json = raw.get("service_account_json", "")
+    if sa_json:
+        try:
+            from google.oauth2 import service_account
+        except ImportError:
+            raise ImportError(
+                "Install bruin-sdk[spanner] to use Spanner credentials: "
+                "pip install 'bruin-sdk[spanner]'"
+            )
+        sa_info = json.loads(sa_json)
+        kwargs["credentials"] = service_account.Credentials.from_service_account_info(sa_info)
+    return spanner_connect(**kwargs)
+
+
+def _create_vertica(raw: dict):
+    try:
+        import vertica_python
+    except ImportError:
+        raise ImportError(
+            "Install bruin-sdk[vertica] to use Vertica connections: "
+            "pip install 'bruin-sdk[vertica]'"
+        )
+    return vertica_python.connect(
+        host=raw["host"],
+        port=int(raw.get("port", 5433)),
+        user=raw["username"],
+        password=raw["password"],
+        database=raw.get("database", ""),
+    )
+
+
 def get_connection(name: str) -> "Connection | GCPConnection":
     """Look up a Bruin-managed connection by name and return a Connection object.
 

src/bruin/_query.py

@@ -16,12 +16,14 @@
 
 # Connection types that use the generic PEP 249 (DBAPI) cursor/read_sql path.
 _DBAPI_TYPES = frozenset((
-    "postgres", "redshift", "mssql", "synapse", "mysql", "athena", "trino", "sqlite",
+    "postgres", "redshift", "mssql", "synapse", "fabric", "mysql",
+    "athena", "trino", "sqlite", "oracle", "db2", "hana", "spanner", "vertica",
 ))
 
 # Subset of _DBAPI_TYPES that require an explicit commit() for DDL/DML.
 _TRANSACTIONAL = frozenset((
-    "postgres", "redshift", "mssql", "synapse", "mysql", "sqlite",
+    "postgres", "redshift", "mssql", "synapse", "fabric", "mysql",
+    "sqlite", "oracle", "db2", "hana", "spanner", "vertica",
 ))
 
 

tests/conftest.py

@@ -149,3 +149,78 @@ def motherduck_connection_json():
         "token": "md_abc123",
         "database": "my_db",
     }
+
+
+@pytest.fixture
+def fabric_connection_json():
+    return {
+        "host": "fabric.example.com",
+        "port": 1433,
+        "username": "admin",
+        "password": "s3cret",
+        "database": "warehouse",
+    }
+
+
+@pytest.fixture
+def oracle_connection_json():
+    return {
+        "host": "oracle.example.com",
+        "port": 1521,
+        "username": "system",
+        "password": "s3cret",
+        "service_name": "ORCL",
+    }
+
+
+@pytest.fixture
+def oracle_sid_connection_json():
+    return {
+        "host": "oracle.example.com",
+        "port": 1521,
+        "username": "system",
+        "password": "s3cret",
+        "sid": "XE",
+    }
+
+
+@pytest.fixture
+def db2_connection_json():
+    return {
+        "host": "db2.example.com",
+        "port": 50000,
+        "username": "db2admin",
+        "password": "s3cret",
+        "database": "SAMPLE",
+    }
+
+
+@pytest.fixture
+def hana_connection_json():
+    return {
+        "host": "hana.example.com",
+        "port": 30015,
+        "username": "SYSTEM",
+        "password": "s3cret",
+        "database": "HDB",
+    }
+
+
+@pytest.fixture
+def spanner_connection_json():
+    return {
+        "project_id": "my-gcp-project",
+        "instance_id": "my-instance",
+        "database": "my-db",
+    }
+
+
+@pytest.fixture
+def vertica_connection_json():
+    return {
+        "host": "vertica.example.com",
+        "port": 5433,
+        "username": "dbadmin",
+        "password": "s3cret",
+        "database": "analytics",
+    }