Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(mssql): support connecting with a url #9894

Merged
merged 3 commits into from
Aug 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 @@
# 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 @@
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"]

Check warning on line 189 in ibis/backends/mssql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mssql/__init__.py#L189

Added line #L189 was not covered by tests

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

Check warning on line 192 in ibis/backends/mssql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mssql/__init__.py#L192

Added line #L192 was not covered by tests

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

Check warning on line 195 in ibis/backends/mssql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mssql/__init__.py#L195

Added line #L195 was not covered by tests

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

Check warning on line 198 in ibis/backends/mssql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mssql/__init__.py#L198

Added line #L198 was not covered by tests

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

Check warning on line 201 in ibis/backends/mssql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mssql/__init__.py#L201

Added line #L201 was not covered by tests

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

Check warning on line 204 in ibis/backends/mssql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mssql/__init__.py#L204

Added line #L204 was not covered by tests

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
Loading