Skip to content

Commit fc39323

Browse files
David Gallaghercpcloud
David Gallagher
authored and
cpcloud
committed
feat(backends/mssql): add backend support for Microsoft Sql Server
1 parent 52ca43e commit fc39323

31 files changed

+781
-148
lines changed

‎.github/workflows/ibis-backends.yml

Lines changed: 32 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ jobs:
3232
test_backends:
3333
name: ${{ matrix.backend.title }} ${{ matrix.os }} python-${{ matrix.python-version }}
3434
runs-on: ${{ matrix.os }}
35+
env:
36+
MSSQL_SA_PASSWORD: "1bis_Testing!"
3537
strategy:
3638
fail-fast: false
3739
matrix:
@@ -82,6 +84,13 @@ jobs:
8284
sys-deps:
8385
- cmake
8486
- ninja-build
87+
- name: mssql
88+
title: MS SQL Server
89+
services:
90+
- mssql
91+
sys-deps:
92+
- libkrb5-dev
93+
- krb5-config
8594
exclude:
8695
- os: windows-latest
8796
backend:
@@ -121,6 +130,15 @@ jobs:
121130
sys-deps:
122131
- cmake
123132
- ninja-build
133+
- os: windows-latest
134+
backend:
135+
name: mssql
136+
title: MS SQL Server
137+
services:
138+
- mssql
139+
sys-deps:
140+
- libkrb5-dev
141+
- krb5-config
124142
steps:
125143
- name: update and install system dependencies
126144
if: ${{ matrix.os == 'ubuntu-latest' && matrix.backend.sys-deps != null }}
@@ -143,6 +161,13 @@ jobs:
143161
- name: checkout
144162
uses: actions/checkout@v3
145163

164+
- uses: extractions/setup-just@v1
165+
env:
166+
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
167+
168+
- name: download backend data
169+
run: just download-data
170+
146171
- name: start services
147172
if: ${{ matrix.backend.services != null }}
148173
run: docker compose up --wait ${{ join(matrix.backend.services, ' ') }}
@@ -173,13 +198,6 @@ jobs:
173198
if: ${{ matrix.backend.has_geo }}
174199
run: poetry install --without dev --without docs --extras ${{ matrix.backend.name }} --extras geospatial
175200

176-
- uses: extractions/setup-just@v1
177-
env:
178-
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
179-
180-
- name: download backend data
181-
run: just download-data
182-
183201
- name: "run parallel tests: ${{ matrix.backend.name }}"
184202
if: ${{ matrix.backend.name != 'pyspark' && matrix.backend.name != 'impala' }}
185203
run: just ci-check -m ${{ matrix.backend.name }} --numprocesses auto --dist=loadgroup
@@ -264,6 +282,13 @@ jobs:
264282
if: ${{ matrix.backend.name == 'postgres' }}
265283
run: sudo apt-get install -qq -y build-essential libgeos-dev
266284

285+
- uses: extractions/setup-just@v1
286+
env:
287+
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
288+
289+
- name: download backend data
290+
run: just download-data
291+
267292
- name: start services
268293
if: ${{ matrix.backend.services != null }}
269294
run: docker compose up --wait ${{ join(matrix.backend.services, ' ') }}
@@ -288,10 +313,6 @@ jobs:
288313
# without updating anything except the requested versions
289314
run: poetry lock --no-update
290315

291-
- uses: extractions/setup-just@v1
292-
env:
293-
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
294-
295316
- name: install ibis
296317
if: ${{ matrix.backend.name != 'postgres' }}
297318
run: poetry install --without dev --without docs --extras ${{ matrix.backend.name }}
@@ -300,9 +321,6 @@ jobs:
300321
if: ${{ matrix.backend.name == 'postgres' }}
301322
run: poetry install --without dev --without docs --extras ${{ matrix.backend.name }} --extras geospatial
302323

303-
- name: download backend data
304-
run: just download-data
305-
306324
- name: run tests
307325
run: just ci-check -m ${{ matrix.backend.name }} --numprocesses auto --dist=loadgroup
308326

‎ci/schema/mssql.sql

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
DROP TABLE IF EXISTS diamonds;
2+
3+
CREATE TABLE diamonds (
4+
carat FLOAT,
5+
cut VARCHAR(MAX),
6+
color VARCHAR(MAX),
7+
clarity VARCHAR(MAX),
8+
depth FLOAT,
9+
"table" FLOAT,
10+
price BIGINT,
11+
x FLOAT,
12+
y FLOAT,
13+
z FLOAT
14+
);
15+
16+
DROP TABLE IF EXISTS batting;
17+
18+
CREATE TABLE batting (
19+
"playerID" VARCHAR(MAX),
20+
"yearID" BIGINT,
21+
stint BIGINT,
22+
"teamID" VARCHAR(MAX),
23+
"lgID" VARCHAR(MAX),
24+
"G" BIGINT,
25+
"AB" BIGINT,
26+
"R" BIGINT,
27+
"H" BIGINT,
28+
"X2B" BIGINT,
29+
"X3B" BIGINT,
30+
"HR" BIGINT,
31+
"RBI" BIGINT,
32+
"SB" BIGINT,
33+
"CS" BIGINT,
34+
"BB" BIGINT,
35+
"SO" BIGINT,
36+
"IBB" BIGINT,
37+
"HBP" BIGINT,
38+
"SH" BIGINT,
39+
"SF" BIGINT,
40+
"GIDP" BIGINT
41+
);
42+
43+
DROP TABLE IF EXISTS awards_players;
44+
45+
CREATE TABLE awards_players (
46+
"playerID" VARCHAR(MAX),
47+
"awardID" VARCHAR(MAX),
48+
"yearID" BIGINT,
49+
"lgID" VARCHAR(MAX),
50+
tie VARCHAR(MAX),
51+
notes VARCHAR(MAX)
52+
);
53+
54+
DROP TABLE IF EXISTS functional_alltypes;
55+
56+
CREATE TABLE functional_alltypes (
57+
"index" BIGINT,
58+
"Unnamed: 0" BIGINT,
59+
id INTEGER,
60+
bool_col BIT,
61+
tinyint_col SMALLINT,
62+
smallint_col SMALLINT,
63+
int_col INTEGER,
64+
bigint_col BIGINT,
65+
float_col REAL,
66+
double_col DOUBLE PRECISION,
67+
date_string_col VARCHAR(MAX),
68+
string_col VARCHAR(MAX),
69+
timestamp_col DATETIME2,
70+
year INTEGER,
71+
month INTEGER
72+
);
73+
74+
CREATE INDEX "ix_functional_alltypes_index" ON functional_alltypes ("index");

‎docker-compose.yml

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,9 +112,28 @@ services:
112112
- 5432:5432
113113
networks:
114114
- postgres
115+
mssql:
116+
environment:
117+
MSSQL_SA_PASSWORD: 1bis_Testing!
118+
ACCEPT_EULA: "Y"
119+
healthcheck:
120+
interval: 10s
121+
retries: 3
122+
test:
123+
- CMD-SHELL
124+
- /opt/mssql-tools/bin/sqlcmd -S localhost -U sa -P "$MSSQL_SA_PASSWORD" -Q "SELECT 1 AS one"
125+
timeout: 10s
126+
image: mcr.microsoft.com/mssql/server:2022-latest
127+
ports:
128+
- 1433:1433
129+
networks:
130+
- mssql
131+
volumes:
132+
- $PWD/ci/ibis-testing-data:/data:ro
115133

116134
networks:
117135
impala:
118136
mysql:
137+
mssql:
119138
clickhouse:
120139
postgres:

‎ibis/backends/conftest.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ def recreate_database(
131131
database: str,
132132
**kwargs: Any,
133133
) -> None:
134-
"""Drop the {database} at {url}, if it exists.
134+
"""Drop the `database` at `url`, if it exists.
135135
136136
Create a new, blank database with the same name.
137137
@@ -142,7 +142,7 @@ def recreate_database(
142142
database : str
143143
Name of the database to be dropped.
144144
"""
145-
engine = sa.create_engine(url, **kwargs)
145+
engine = sa.create_engine(url.set(database=""), **kwargs)
146146

147147
if url.database is not None:
148148
with engine.connect() as conn:
@@ -157,9 +157,9 @@ def init_database(
157157
recreate: bool = True,
158158
**kwargs: Any,
159159
) -> sa.engine.Engine:
160-
"""Initialise {database} at {url} with {schema}.
160+
"""Initialise `database` at `url` with `schema`.
161161
162-
If {recreate}, drop the {database} at {url}, if it exists.
162+
If `recreate`, drop the `database` at `url`, if it exists.
163163
164164
Parameters
165165
----------

‎ibis/backends/mssql/__init__.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
"""The Microsoft Sql Server backend."""
2+
3+
from __future__ import annotations
4+
5+
from typing import Literal
6+
7+
import sqlalchemy as sa
8+
9+
from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
10+
from ibis.backends.mssql.compiler import MsSqlCompiler
11+
12+
13+
class Backend(BaseAlchemyBackend):
14+
name = "mssql"
15+
compiler = MsSqlCompiler
16+
17+
def do_connect(
18+
self,
19+
host: str = "localhost",
20+
user: str | None = None,
21+
password: str | None = None,
22+
port: int = 1433,
23+
database: str | None = None,
24+
url: str | None = None,
25+
driver: Literal["pymssql"] = "pymssql",
26+
) -> None:
27+
if driver != "pymssql":
28+
raise NotImplementedError("pymssql is currently the only supported driver")
29+
alchemy_url = self._build_alchemy_url(
30+
url=url,
31+
host=host,
32+
port=port,
33+
user=user,
34+
password=password,
35+
database=database,
36+
driver=f'mssql+{driver}',
37+
)
38+
self.database_name = alchemy_url.database
39+
super().do_connect(sa.create_engine(alchemy_url))

‎ibis/backends/mssql/compiler.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import sqlalchemy as sa
2+
from sqlalchemy.dialects import mssql
3+
4+
import ibis.expr.datatypes as dt
5+
from ibis.backends.base.sql.alchemy import AlchemyCompiler, AlchemyExprTranslator
6+
from ibis.backends.mssql.registry import operation_registry
7+
8+
9+
class MsSqlExprTranslator(AlchemyExprTranslator):
10+
_registry = operation_registry
11+
_rewrites = AlchemyExprTranslator._rewrites.copy()
12+
_type_map = AlchemyExprTranslator._type_map.copy()
13+
_type_map.update(
14+
{
15+
dt.Boolean: mssql.BIT,
16+
dt.Int8: mssql.TINYINT,
17+
dt.Int16: mssql.SMALLINT,
18+
dt.Int32: mssql.INTEGER,
19+
dt.Int64: mssql.BIGINT,
20+
dt.Float16: mssql.FLOAT,
21+
dt.Float32: mssql.FLOAT,
22+
dt.Float64: mssql.REAL,
23+
dt.String: mssql.NVARCHAR,
24+
}
25+
)
26+
_bool_aggs_need_cast_to_int32 = True
27+
integer_to_timestamp = sa.func.from_unixtime
28+
native_json_type = False
29+
30+
31+
rewrites = MsSqlExprTranslator.rewrites
32+
33+
34+
class MsSqlCompiler(AlchemyCompiler):
35+
translator_class = MsSqlExprTranslator

0 commit comments

Comments
 (0)