152 changes: 152 additions & 0 deletions ibis/backends/flink/ddl.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
from __future__ import annotations

from typing import TYPE_CHECKING

import sqlglot as sg

from ibis.backends.base.sql.ddl import (
CreateTableWithSchema,
DropObject,
_CreateDDL,
_format_properties,
_is_quoted,
is_fully_qualified,
)
from ibis.backends.base.sql.registry import quote_identifier

if TYPE_CHECKING:
import ibis.expr.schema as sch


class _CatalogAwareBaseQualifiedSQLStatement:
def _get_scoped_name(
self, obj_name: str, database: str | None = None, catalog: str | None = None
) -> str:
if is_fully_qualified(obj_name):
return obj_name
if _is_quoted(obj_name):
obj_name = obj_name[1:-1]
return sg.table(obj_name, db=database, catalog=catalog, quoted=True).sql(
dialect="hive"
)


class CreateTableFromConnector(
_CatalogAwareBaseQualifiedSQLStatement, CreateTableWithSchema
):
def __init__(
self,
table_name: str,
schema: sch.Schema,
tbl_properties: dict,
database: str | None = None,
catalog: str | None = None,
temp: bool = False,
**kwargs,
):
super().__init__(
table_name=table_name,
database=database,
schema=schema,
table_format=None,
format=None,
path=None,
tbl_properties=tbl_properties,
**kwargs,
)
self.catalog = catalog
self.temp = temp

def _storage(self) -> str:
return f"STORED AS {self.format}" if self.format else None

def _format_tbl_properties(self) -> str:
return f"WITH {_format_properties(self.tbl_properties)}"

@property
def _prefix(self) -> str:
# `TEMPORARY` is not documented in Flink's documentation
modifier = " TEMPORARY" if self.temp else ""
return f"CREATE{modifier} TABLE"

def _create_line(self) -> str:
scoped_name = self._get_scoped_name(
self.table_name, self.database, self.catalog
)
return f"{self._prefix} {self._if_exists()}{scoped_name}"

@property
def _pieces(self):
yield from super()._pieces
yield self._format_tbl_properties()


class DropTable(_CatalogAwareBaseQualifiedSQLStatement, DropObject):
_object_type = "TABLE"

def __init__(
self,
table_name: str,
database: str | None = None,
catalog: str | None = None,
must_exist: bool = True,
temp: bool = False,
):
super().__init__(must_exist=must_exist)
self.table_name = table_name
self.database = database
self.catalog = catalog
self.temp = temp

def _object_name(self):
return self._get_scoped_name(self.table_name, self.database, self.catalog)

def compile(self):
temp = "TEMPORARY " if self.temp else ""
if_exists = "" if self.must_exist else "IF EXISTS "
object_name = self._object_name()
return f"DROP {temp}{self._object_type} {if_exists}{object_name}"


class _DatabaseObject:
def _object_name(self):
scoped_name = f"{quote_identifier(self.catalog)}." if self.catalog else ""
scoped_name += quote_identifier(self.name)
return scoped_name


class CreateDatabase(_DatabaseObject, _CreateDDL):
def __init__(
self,
name: str,
db_properties: dict | None,
catalog: str | None = None,
can_exist: bool = False,
):
# TODO(chloeh13q): support COMMENT
self.name = name
self.db_properties = db_properties
self.catalog = catalog
self.can_exist = can_exist

def _format_db_properties(self) -> str:
return (
f"WITH {_format_properties(self.db_properties)}"
if self.db_properties
else ""
)

def compile(self):
create_decl = "CREATE DATABASE"
create_line = f"{create_decl} {self._if_exists()}{self._object_name()}"

return f"{create_line}\n{self._format_db_properties()}"


class DropDatabase(_DatabaseObject, DropObject):
_object_type = "DATABASE"

def __init__(self, name: str, catalog: str | None = None, must_exist: bool = True):
super().__init__(must_exist=must_exist)
self.name = name
self.catalog = catalog
15 changes: 15 additions & 0 deletions ibis/backends/flink/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,3 +49,18 @@ def simple_schema():
@pytest.fixture
def simple_table(simple_schema):
return ibis.table(simple_schema, name="table")


@pytest.fixture(scope="session")
def con(tmp_path_factory, data_dir, worker_id):
return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection


@pytest.fixture(scope="session")
def db(con):
return con.database()


@pytest.fixture(scope="session")
def alltypes(con):
return con.tables.functional_alltypes
84 changes: 84 additions & 0 deletions ibis/backends/flink/tests/test_ddl.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
from __future__ import annotations

import pytest
from py4j.protocol import Py4JJavaError

import ibis.expr.datatypes as dt
import ibis.expr.schema as sch


@pytest.fixture
def awards_players_schema():
return sch.Schema(
{
"playerID": dt.string,
"awardID": dt.string,
"yearID": dt.int32,
"lgID": dt.string,
"tie": dt.string,
"notes": dt.string,
}
)


@pytest.fixture
def awards_players_csv_connector_configs():
return {
"connector": "filesystem",
"path": "ci/ibis-testing-data/csv/awards_players.csv",
"format": "csv",
"csv.ignore-parse-errors": "true",
}


def test_list_tables(con):
assert len(con.list_tables()) == 4
assert (
len(con.list_tables(catalog="default_catalog", database="default_database"))
== 4
)


def test_create_table_from_schema(
con, awards_players_schema, temp_table, awards_players_csv_connector_configs
):
new_table = con.create_table(
temp_table,
schema=awards_players_schema,
tbl_properties=awards_players_csv_connector_configs,
)
assert len(con.list_tables()) == 5
assert temp_table in con.list_tables()
assert new_table.schema() == awards_players_schema


def test_drop_table(
con, awards_players_schema, temp_table, awards_players_csv_connector_configs
):
con.create_table(
temp_table,
schema=awards_players_schema,
tbl_properties=awards_players_csv_connector_configs,
)
assert len(con.list_tables()) == 5
con.drop_table(temp_table)
assert len(con.list_tables()) == 4
assert temp_table not in con.list_tables()


def test_temp_table(
con, awards_players_schema, temp_table, awards_players_csv_connector_configs
):
con.create_table(
temp_table,
schema=awards_players_schema,
tbl_properties=awards_players_csv_connector_configs,
temp=True,
)
assert len(con.list_tables()) == 5
assert temp_table in con.list_tables()
with pytest.raises(Py4JJavaError):
con.drop_table(temp_table)
con.drop_table(temp_table, temp=True)
assert len(con.list_tables()) == 4
assert temp_table not in con.list_tables()
2 changes: 1 addition & 1 deletion ibis/backends/pyspark/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -294,7 +294,7 @@ def create_database(
path
Path where to store the database data; otherwise uses Spark default
force
Whether to append `IF EXISTS` to the database creation SQL
Whether to append `IF NOT EXISTS` to the database creation SQL
"""
statement = CreateDatabase(name, path=path, can_exist=force)
return self.raw_sql(statement.compile())
Expand Down