Skip to content

Commit

Permalink
feat(backends/mssql): add backend support for Microsoft Sql Server
Browse files Browse the repository at this point in the history
  • Loading branch information
David Gallagher authored and cpcloud committed Nov 22, 2022
1 parent df91515 commit a5b0eb4
Show file tree
Hide file tree
Showing 31 changed files with 781 additions and 148 deletions.
46 changes: 32 additions & 14 deletions .github/workflows/ibis-backends.yml
Expand Up @@ -32,6 +32,8 @@ jobs:
test_backends:
name: ${{ matrix.backend.title }} ${{ matrix.os }} python-${{ matrix.python-version }}
runs-on: ${{ matrix.os }}
env:
MSSQL_SA_PASSWORD: "1bis_Testing!"
strategy:
fail-fast: false
matrix:
Expand Down Expand Up @@ -82,6 +84,13 @@ jobs:
sys-deps:
- cmake
- ninja-build
- name: mssql
title: MS SQL Server
services:
- mssql
sys-deps:
- libkrb5-dev
- krb5-config
exclude:
- os: windows-latest
backend:
Expand Down Expand Up @@ -121,6 +130,15 @@ jobs:
sys-deps:
- cmake
- ninja-build
- os: windows-latest
backend:
name: mssql
title: MS SQL Server
services:
- mssql
sys-deps:
- libkrb5-dev
- krb5-config
steps:
- name: update and install system dependencies
if: ${{ matrix.os == 'ubuntu-latest' && matrix.backend.sys-deps != null }}
Expand All @@ -143,6 +161,13 @@ jobs:
- name: checkout
uses: actions/checkout@v3

- uses: extractions/setup-just@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

- name: download backend data
run: just download-data

- name: start services
if: ${{ matrix.backend.services != null }}
run: docker compose up --wait ${{ join(matrix.backend.services, ' ') }}
Expand Down Expand Up @@ -173,13 +198,6 @@ jobs:
if: ${{ matrix.backend.has_geo }}
run: poetry install --without dev --without docs --extras ${{ matrix.backend.name }} --extras geospatial

- uses: extractions/setup-just@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

- name: download backend data
run: just download-data

- name: "run parallel tests: ${{ matrix.backend.name }}"
if: ${{ matrix.backend.name != 'pyspark' && matrix.backend.name != 'impala' }}
run: just ci-check -m ${{ matrix.backend.name }} --numprocesses auto --dist=loadgroup
Expand Down Expand Up @@ -264,6 +282,13 @@ jobs:
if: ${{ matrix.backend.name == 'postgres' }}
run: sudo apt-get install -qq -y build-essential libgeos-dev

- uses: extractions/setup-just@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

- name: download backend data
run: just download-data

- name: start services
if: ${{ matrix.backend.services != null }}
run: docker compose up --wait ${{ join(matrix.backend.services, ' ') }}
Expand All @@ -288,10 +313,6 @@ jobs:
# without updating anything except the requested versions
run: poetry lock --no-update

- uses: extractions/setup-just@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

- name: install ibis
if: ${{ matrix.backend.name != 'postgres' }}
run: poetry install --without dev --without docs --extras ${{ matrix.backend.name }}
Expand All @@ -300,9 +321,6 @@ jobs:
if: ${{ matrix.backend.name == 'postgres' }}
run: poetry install --without dev --without docs --extras ${{ matrix.backend.name }} --extras geospatial

- name: download backend data
run: just download-data

- name: run tests
run: just ci-check -m ${{ matrix.backend.name }} --numprocesses auto --dist=loadgroup

Expand Down
74 changes: 74 additions & 0 deletions ci/schema/mssql.sql
@@ -0,0 +1,74 @@
DROP TABLE IF EXISTS diamonds;

CREATE TABLE diamonds (
carat FLOAT,
cut VARCHAR(MAX),
color VARCHAR(MAX),
clarity VARCHAR(MAX),
depth FLOAT,
"table" FLOAT,
price BIGINT,
x FLOAT,
y FLOAT,
z FLOAT
);

DROP TABLE IF EXISTS batting;

CREATE TABLE batting (
"playerID" VARCHAR(MAX),
"yearID" BIGINT,
stint BIGINT,
"teamID" VARCHAR(MAX),
"lgID" VARCHAR(MAX),
"G" BIGINT,
"AB" BIGINT,
"R" BIGINT,
"H" BIGINT,
"X2B" BIGINT,
"X3B" BIGINT,
"HR" BIGINT,
"RBI" BIGINT,
"SB" BIGINT,
"CS" BIGINT,
"BB" BIGINT,
"SO" BIGINT,
"IBB" BIGINT,
"HBP" BIGINT,
"SH" BIGINT,
"SF" BIGINT,
"GIDP" BIGINT
);

DROP TABLE IF EXISTS awards_players;

CREATE TABLE awards_players (
"playerID" VARCHAR(MAX),
"awardID" VARCHAR(MAX),
"yearID" BIGINT,
"lgID" VARCHAR(MAX),
tie VARCHAR(MAX),
notes VARCHAR(MAX)
);

DROP TABLE IF EXISTS functional_alltypes;

CREATE TABLE functional_alltypes (
"index" BIGINT,
"Unnamed: 0" BIGINT,
id INTEGER,
bool_col BIT,
tinyint_col SMALLINT,
smallint_col SMALLINT,
int_col INTEGER,
bigint_col BIGINT,
float_col REAL,
double_col DOUBLE PRECISION,
date_string_col VARCHAR(MAX),
string_col VARCHAR(MAX),
timestamp_col DATETIME2,
year INTEGER,
month INTEGER
);

CREATE INDEX "ix_functional_alltypes_index" ON functional_alltypes ("index");
19 changes: 19 additions & 0 deletions docker-compose.yml
Expand Up @@ -112,9 +112,28 @@ services:
- 5432:5432
networks:
- postgres
mssql:
environment:
MSSQL_SA_PASSWORD: 1bis_Testing!
ACCEPT_EULA: "Y"
healthcheck:
interval: 10s
retries: 3
test:
- CMD-SHELL
- /opt/mssql-tools/bin/sqlcmd -S localhost -U sa -P "$MSSQL_SA_PASSWORD" -Q "SELECT 1 AS one"
timeout: 10s
image: mcr.microsoft.com/mssql/server:2022-latest
ports:
- 1433:1433
networks:
- mssql
volumes:
- $PWD/ci/ibis-testing-data:/data:ro

networks:
impala:
mysql:
mssql:
clickhouse:
postgres:
8 changes: 4 additions & 4 deletions ibis/backends/conftest.py
Expand Up @@ -131,7 +131,7 @@ def recreate_database(
database: str,
**kwargs: Any,
) -> None:
"""Drop the {database} at {url}, if it exists.
"""Drop the `database` at `url`, if it exists.
Create a new, blank database with the same name.
Expand All @@ -142,7 +142,7 @@ def recreate_database(
database : str
Name of the database to be dropped.
"""
engine = sa.create_engine(url, **kwargs)
engine = sa.create_engine(url.set(database=""), **kwargs)

if url.database is not None:
with engine.connect() as conn:
Expand All @@ -157,9 +157,9 @@ def init_database(
recreate: bool = True,
**kwargs: Any,
) -> sa.engine.Engine:
"""Initialise {database} at {url} with {schema}.
"""Initialise `database` at `url` with `schema`.
If {recreate}, drop the {database} at {url}, if it exists.
If `recreate`, drop the `database` at `url`, if it exists.
Parameters
----------
Expand Down
39 changes: 39 additions & 0 deletions ibis/backends/mssql/__init__.py
@@ -0,0 +1,39 @@
"""The Microsoft Sql Server backend."""

from __future__ import annotations

from typing import Literal

import sqlalchemy as sa

from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
from ibis.backends.mssql.compiler import MsSqlCompiler


class Backend(BaseAlchemyBackend):
name = "mssql"
compiler = MsSqlCompiler

def do_connect(
self,
host: str = "localhost",
user: str | None = None,
password: str | None = None,
port: int = 1433,
database: str | None = None,
url: str | None = None,
driver: Literal["pymssql"] = "pymssql",
) -> None:
if driver != "pymssql":
raise NotImplementedError("pymssql is currently the only supported driver")
alchemy_url = self._build_alchemy_url(
url=url,
host=host,
port=port,
user=user,
password=password,
database=database,
driver=f'mssql+{driver}',
)
self.database_name = alchemy_url.database
super().do_connect(sa.create_engine(alchemy_url))
35 changes: 35 additions & 0 deletions ibis/backends/mssql/compiler.py
@@ -0,0 +1,35 @@
import sqlalchemy as sa
from sqlalchemy.dialects import mssql

import ibis.expr.datatypes as dt
from ibis.backends.base.sql.alchemy import AlchemyCompiler, AlchemyExprTranslator
from ibis.backends.mssql.registry import operation_registry


class MsSqlExprTranslator(AlchemyExprTranslator):
_registry = operation_registry
_rewrites = AlchemyExprTranslator._rewrites.copy()
_type_map = AlchemyExprTranslator._type_map.copy()
_type_map.update(
{
dt.Boolean: mssql.BIT,
dt.Int8: mssql.TINYINT,
dt.Int16: mssql.SMALLINT,
dt.Int32: mssql.INTEGER,
dt.Int64: mssql.BIGINT,
dt.Float16: mssql.FLOAT,
dt.Float32: mssql.FLOAT,
dt.Float64: mssql.REAL,
dt.String: mssql.NVARCHAR,
}
)
_bool_aggs_need_cast_to_int32 = True
integer_to_timestamp = sa.func.from_unixtime
native_json_type = False


rewrites = MsSqlExprTranslator.rewrites


class MsSqlCompiler(AlchemyCompiler):
translator_class = MsSqlExprTranslator

0 comments on commit a5b0eb4

Please sign in to comment.