Skip to content

Commit

Permalink
feat(mssql): support connecting with a url (#9894)
Browse files Browse the repository at this point in the history
Closes #9856.
  • Loading branch information
cpcloud authored Aug 22, 2024
1 parent d112c15 commit 8bb12e1
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 3 deletions.
40 changes: 38 additions & 2 deletions ibis/backends/mssql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from contextlib import closing
from operator import itemgetter
from typing import TYPE_CHECKING, Any
from urllib.parse import unquote_plus

import pyodbc
import sqlglot as sg
Expand All @@ -21,12 +22,13 @@
import ibis.expr.schema as sch
import ibis.expr.types as ir
from ibis import util
from ibis.backends import CanCreateCatalog, CanCreateDatabase, CanCreateSchema, NoUrl
from ibis.backends import CanCreateCatalog, CanCreateDatabase, CanCreateSchema
from ibis.backends.sql import SQLBackend
from ibis.backends.sql.compilers.base import STAR, C

if TYPE_CHECKING:
from collections.abc import Iterable, Mapping
from urllib.parse import ParseResult

import pandas as pd
import polars as pl
Expand Down Expand Up @@ -73,7 +75,7 @@ def datetimeoffset_to_datetime(value):
# Databases: sys.schemas


class Backend(SQLBackend, CanCreateCatalog, CanCreateDatabase, CanCreateSchema, NoUrl):
class Backend(SQLBackend, CanCreateCatalog, CanCreateDatabase, CanCreateSchema):
name = "mssql"
compiler = sc.mssql.compiler
supports_create_or_replace = False
Expand Down Expand Up @@ -169,6 +171,40 @@ def _post_connect(self):
with closing(self.con.cursor()) as cur:
cur.execute("SET DATEFIRST 1")

def _from_url(self, url: ParseResult, **kwargs):
database, *_ = url.path[1:].split("/", 1)
kwargs.update(
{
"user": url.username,
"password": unquote_plus(url.password or ""),
"host": url.hostname,
"database": database or "",
"port": url.port or None,
}
)

self._convert_kwargs(kwargs)

if "host" in kwargs and not kwargs["host"]:
del kwargs["host"]

if "user" in kwargs and not kwargs["user"]:
del kwargs["user"]

if "password" in kwargs and kwargs["password"] is None:
del kwargs["password"]

if "port" in kwargs and kwargs["port"] is None:
del kwargs["port"]

if "database" in kwargs and not kwargs["database"]:
del kwargs["database"]

if "driver" in kwargs and not kwargs["driver"]:
del kwargs["driver"]

return self.connect(**kwargs)

def get_schema(
self, name: str, *, catalog: str | None = None, database: str | None = None
) -> sch.Schema:
Expand Down
24 changes: 24 additions & 0 deletions ibis/backends/mssql/tests/test_client.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from __future__ import annotations

from urllib.parse import urlencode

import pytest
import sqlglot as sg
import sqlglot.expressions as sge
Expand All @@ -8,6 +10,14 @@
import ibis
import ibis.expr.datatypes as dt
from ibis import udf
from ibis.backends.mssql.tests.conftest import (
IBIS_TEST_MSSQL_DB,
MSSQL_HOST,
MSSQL_PASS,
MSSQL_PORT,
MSSQL_PYODBC_DRIVER,
MSSQL_USER,
)

RAW_DB_TYPES = [
# Exact numbers
Expand Down Expand Up @@ -204,3 +214,17 @@ def test_create_temp_table_from_obj(con):
assert persisted_from_temp.to_pyarrow().equals(t2.to_pyarrow())

con.drop_table("fuhreal")


def test_from_url():
user = MSSQL_USER
password = MSSQL_PASS
host = MSSQL_HOST
port = MSSQL_PORT
database = IBIS_TEST_MSSQL_DB
driver = MSSQL_PYODBC_DRIVER
new_con = ibis.connect(
f"mssql://{user}:{password}@{host}:{port}/{database}?{urlencode(dict(driver=driver))}"
)
result = new_con.sql("SELECT 1 AS [a]").to_pandas().a.iat[0]
assert result == 1
2 changes: 1 addition & 1 deletion ibis/backends/oracle/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ def test_list_tables_schema_warning_refactor(con):
assert con.list_tables(database="SYS", like="EXU8OPT") == ["EXU8OPT"]


def test_from_url(con):
def test_from_url():
new_con = ibis.connect("oracle://ibis:ibis@localhost:1521/IBIS_TESTING")

assert new_con.list_tables()
Expand Down

0 comments on commit 8bb12e1

Please sign in to comment.