Skip to content

Commit 5bfa090

Browse files
committed
fix(impala): make create_database follow the CanCreateDatabase API
1 parent bd926f0 commit 5bfa090

File tree

3 files changed

+25
-35
lines changed

3 files changed

+25
-35
lines changed

ibis/backends/impala/__init__.py

Lines changed: 22 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from typing import TYPE_CHECKING, Any, Literal
1010

1111
import impala.dbapi as impyla
12+
import impala.hiveserver2 as hs2
1213
import sqlglot as sg
1314
import sqlglot.expressions as sge
1415
from impala.error import Error as ImpylaError
@@ -19,7 +20,7 @@
1920
import ibis.expr.schema as sch
2021
import ibis.expr.types as ir
2122
from ibis import util
22-
from ibis.backends import NoExampleLoader
23+
from ibis.backends import CanCreateDatabase, NoExampleLoader
2324
from ibis.backends.impala import ddl, udf
2425
from ibis.backends.impala.udf import (
2526
aggregate_function,
@@ -34,7 +35,6 @@
3435
from pathlib import Path
3536
from urllib.parse import ParseResult
3637

37-
import impala.hiveserver2 as hs2
3838
import pandas as pd
3939
import polars as pl
4040
import pyarrow as pa
@@ -51,7 +51,7 @@
5151
)
5252

5353

54-
class Backend(SQLBackend, NoExampleLoader):
54+
class Backend(SQLBackend, CanCreateDatabase, NoExampleLoader):
5555
name = "impala"
5656
compiler = sc.impala.compiler
5757

@@ -286,34 +286,28 @@ def current_database(self) -> str:
286286
[(db,)] = cur.fetchall()
287287
return db
288288

289-
def create_database(self, name, path=None, force=False):
290-
"""Create a new Impala database.
291-
292-
Parameters
293-
----------
294-
name
295-
Database name
296-
path
297-
Path where to store the database data; otherwise uses the Impala default
298-
force
299-
Forcibly create the database
289+
def table(
290+
self, name, /, *, database: str | tuple[str, str] | None = None
291+
) -> ir.Table:
292+
try:
293+
return super().table(name, database=database)
294+
except hs2.HiveServer2Error as e:
295+
if "AnalysisException: Could not resolve path:" in str(e):
296+
raise com.TableNotFound(name) from e
300297

301-
"""
302-
statement = ddl.CreateDatabase(name, path=path, can_exist=force)
298+
def create_database(
299+
self, name: str, /, *, catalog: str | None = None, force: bool = False
300+
) -> None:
301+
statement = ddl.CreateDatabase(name, path=catalog, can_exist=force)
303302
self._safe_exec_sql(statement)
304303

305-
def drop_database(self, name, force=False):
306-
"""Drop an Impala database.
307-
308-
Parameters
309-
----------
310-
name
311-
Database name
312-
force
313-
If False and there are any tables in this database, raises an
314-
IntegrityError
315-
316-
"""
304+
def drop_database(
305+
self, name: str, /, *, catalog: str | None = None, force: bool = False
306+
) -> None:
307+
if catalog is not None:
308+
raise NotImplementedError(
309+
"Ibis has not yet implemented `catalog` parameter of drop_database() for Impala"
310+
)
317311
if not force or name in self.list_databases():
318312
tables = self.list_tables(database=name)
319313
udfs = self.list_udfs(database=name)

ibis/backends/impala/tests/test_ddl.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ def temp_db(con, tmp_dir):
5353

5454
def test_create_database_with_location(con, temp_db):
5555
name = os.path.basename(temp_db)
56-
con.create_database(name, path=temp_db)
56+
con.create_database(name, catalog=temp_db)
5757
assert name in con.list_databases()
5858

5959

ibis/backends/tests/test_signatures.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,10 @@ def _scrape_methods(modules, params):
4242
marks = {
4343
"compile": pytest.param(BaseBackend, "compile"),
4444
"create_database": pytest.param(
45-
CanCreateDatabase,
46-
"create_database",
47-
marks=pytest.mark.notyet(["impala", "mysql"]),
45+
CanCreateDatabase, "create_database", marks=pytest.mark.notyet(["mysql"])
4846
),
4947
"drop_database": pytest.param(
50-
CanCreateDatabase,
51-
"drop_database",
52-
marks=pytest.mark.notyet(["impala", "mysql"]),
48+
CanCreateDatabase, "drop_database", marks=pytest.mark.notyet(["mysql"])
5349
),
5450
"drop_table": pytest.param(
5551
SQLBackend, "drop_table", marks=pytest.mark.notyet(["druid"])

0 commit comments

Comments
 (0)