Skip to content

Commit

Permalink
feat: implement force argument in sqlalchemy backend base class
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Mar 1, 2022
1 parent a0aa65b commit 9df7f1b
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 13 deletions.
46 changes: 39 additions & 7 deletions ibis/backends/base/sql/alchemy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import contextlib
import getpass
from typing import Literal
from typing import Any, Literal

import pandas as pd
import sqlalchemy as sa
Expand Down Expand Up @@ -152,6 +152,7 @@ def create_table(
expr: pd.DataFrame | ir.TableExpr | None = None,
schema: sch.Schema | None = None,
database: str | None = None,
force: bool = False,
) -> None:
"""Create a table.
Expand All @@ -165,6 +166,8 @@ def create_table(
An ibis schema
database
A database
force
Check whether a table exists before creating it
"""
if database == self.current_database:
# avoid fully qualified name
Expand Down Expand Up @@ -194,7 +197,7 @@ def create_table(
)

with self.begin() as bind:
t.create(bind=bind)
t.create(bind=bind, checkfirst=force)
if expr is not None:
bind.execute(
t.insert().from_select(list(expr.columns), expr.compile())
Expand Down Expand Up @@ -310,7 +313,7 @@ def truncate_table(
t.delete().execute()

def schema(self, name: str) -> sch.Schema:
"""Get a schema object from the current database for the table `name`.
"""Get an ibis schema from the current database for the table `name`.
Parameters
----------
Expand All @@ -320,7 +323,7 @@ def schema(self, name: str) -> sch.Schema:
Returns
-------
Schema
The schema of the object `name`.
The ibis schema of `name`
"""
return self.database().schema(name)

Expand All @@ -341,13 +344,42 @@ def _log(self, sql):
else:
util.log(query_str)

def _get_sqla_table(self, name, schema=None, autoload=True):
def _get_sqla_table(
self,
name: str,
schema: str | None = None,
autoload: bool = True,
**kwargs: Any,
) -> sa.Table:
return sa.Table(name, self.meta, schema=schema, autoload=autoload)

def _sqla_table_to_expr(self, table):
node = self.table_class(table, self)
def _sqla_table_to_expr(self, table: sa.Table) -> ir.TableExpr:
schema = self._schemas.get(table.name)
node = self.table_class(table, self, schema)
return self.table_expr_class(node)

def table(
self,
name: str,
**kwargs: Any,
) -> ir.TableExpr:
"""Create a table expression from a table in the SQLite database.
Parameters
----------
name
Table name
kwargs
Backend specific arguments.
Returns
-------
TableExpr
Table expression
"""
sqla_table = self._get_sqla_table(name, **kwargs)
return self._sqla_table_to_expr(sqla_table)

def insert(
self,
table_name: str,
Expand Down
8 changes: 2 additions & 6 deletions ibis/backends/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -272,20 +272,16 @@ def test_insert_overwrite_from_expr(
def test_list_databases(alchemy_backend, alchemy_con):
# Every backend has its own databases
TEST_DATABASES = {
'sqlite': ['main', 'base'],
'sqlite': ['main'],
'postgres': ['postgres', 'ibis_testing'],
'mysql': ['ibis_testing', 'information_schema'],
}
assert alchemy_con.list_databases() == TEST_DATABASES[alchemy_con.name]


@mark.never(
["postgres"], reason="schemas and databases are different in postgres"
)
def test_list_schemas(alchemy_backend, alchemy_con):
with pytest.warns(FutureWarning):
schemas = alchemy_con.list_schemas()
assert schemas == alchemy_con.list_databases()
alchemy_con.list_schemas()


def test_verify(ddl_backend, ddl_con):
Expand Down

0 comments on commit 9df7f1b

Please sign in to comment.