|
9 | 9 | from typing import TYPE_CHECKING, Any, Literal
|
10 | 10 |
|
11 | 11 | import impala.dbapi as impyla
|
| 12 | +import impala.hiveserver2 as hs2 |
12 | 13 | import sqlglot as sg
|
13 | 14 | import sqlglot.expressions as sge
|
14 | 15 | from impala.error import Error as ImpylaError
|
|
19 | 20 | import ibis.expr.schema as sch
|
20 | 21 | import ibis.expr.types as ir
|
21 | 22 | from ibis import util
|
22 |
| -from ibis.backends import NoExampleLoader |
| 23 | +from ibis.backends import CanCreateDatabase, NoExampleLoader |
23 | 24 | from ibis.backends.impala import ddl, udf
|
24 | 25 | from ibis.backends.impala.udf import (
|
25 | 26 | aggregate_function,
|
|
34 | 35 | from pathlib import Path
|
35 | 36 | from urllib.parse import ParseResult
|
36 | 37 |
|
37 |
| - import impala.hiveserver2 as hs2 |
38 | 38 | import pandas as pd
|
39 | 39 | import polars as pl
|
40 | 40 | import pyarrow as pa
|
|
51 | 51 | )
|
52 | 52 |
|
53 | 53 |
|
54 |
| -class Backend(SQLBackend, NoExampleLoader): |
| 54 | +class Backend(SQLBackend, CanCreateDatabase, NoExampleLoader): |
55 | 55 | name = "impala"
|
56 | 56 | compiler = sc.impala.compiler
|
57 | 57 |
|
@@ -286,34 +286,28 @@ def current_database(self) -> str:
|
286 | 286 | [(db,)] = cur.fetchall()
|
287 | 287 | return db
|
288 | 288 |
|
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 |
300 | 297 |
|
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) |
303 | 302 | self._safe_exec_sql(statement)
|
304 | 303 |
|
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 | + ) |
317 | 311 | if not force or name in self.list_databases():
|
318 | 312 | tables = self.list_tables(database=name)
|
319 | 313 | udfs = self.list_udfs(database=name)
|
|
0 commit comments