Skip to content

Commit 97a470c

Browse files
douenergydouenergy
committed
feat(ibis): Add Oracle connector
1 parent 3c928e0 commit 97a470c

File tree

5 files changed

+169
-0
lines changed

5 files changed

+169
-0
lines changed

‎ibis-server/app/model/__init__.py

+12
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,10 @@ class QueryMySqlDTO(QueryDTO):
4040
connection_info: ConnectionUrl | MySqlConnectionInfo = connection_info_field
4141

4242

43+
class QueryOracleDTO(QueryDTO):
44+
connection_info: ConnectionUrl | OracleConnectionInfo = connection_info_field
45+
46+
4347
class QueryPostgresDTO(QueryDTO):
4448
connection_info: ConnectionUrl | PostgresConnectionInfo = connection_info_field
4549

@@ -131,6 +135,14 @@ class PostgresConnectionInfo(BaseModel):
131135
password: SecretStr | None = None
132136

133137

138+
class OracleConnectionInfo(BaseModel):
139+
host: SecretStr = Field(examples=["localhost"])
140+
port: SecretStr = Field(examples=[1521])
141+
database: SecretStr
142+
user: SecretStr
143+
password: SecretStr | None = None
144+
145+
134146
class SnowflakeConnectionInfo(BaseModel):
135147
user: SecretStr
136148
password: SecretStr

‎ibis-server/app/model/data_source.py

+14
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
ConnectionInfo,
1818
MSSqlConnectionInfo,
1919
MySqlConnectionInfo,
20+
OracleConnectionInfo,
2021
PostgresConnectionInfo,
2122
QueryBigQueryDTO,
2223
QueryCannerDTO,
@@ -27,6 +28,7 @@
2728
QueryMinioFileDTO,
2829
QueryMSSqlDTO,
2930
QueryMySqlDTO,
31+
QueryOracleDTO,
3032
QueryPostgresDTO,
3133
QueryS3FileDTO,
3234
QuerySnowflakeDTO,
@@ -43,6 +45,7 @@ class DataSource(StrEnum):
4345
clickhouse = auto()
4446
mssql = auto()
4547
mysql = auto()
48+
oracle = auto()
4649
postgres = auto()
4750
snowflake = auto()
4851
trino = auto()
@@ -70,6 +73,7 @@ class DataSourceExtension(Enum):
7073
clickhouse = QueryClickHouseDTO
7174
mssql = QueryMSSqlDTO
7275
mysql = QueryMySqlDTO
76+
oracle = QueryOracleDTO
7377
postgres = QueryPostgresDTO
7478
snowflake = QuerySnowflakeDTO
7579
trino = QueryTrinoDTO
@@ -176,6 +180,16 @@ def get_postgres_connection(info: PostgresConnectionInfo) -> BaseBackend:
176180
password=(info.password and info.password.get_secret_value()),
177181
)
178182

183+
@staticmethod
184+
def get_oracle_connection(info: OracleConnectionInfo) -> BaseBackend:
185+
return ibis.oracle.connect(
186+
host=info.host.get_secret_value(),
187+
port=int(info.port.get_secret_value()),
188+
database=info.database.get_secret_value(),
189+
user=info.user.get_secret_value(),
190+
password=(info.password and info.password.get_secret_value()),
191+
)
192+
179193
@staticmethod
180194
def get_snowflake_connection(info: SnowflakeConnectionInfo) -> BaseBackend:
181195
return ibis.snowflake.connect(

‎ibis-server/app/util.py

+6
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@ def base64_to_dict(base64_str: str) -> dict:
1313

1414
def to_json(df: pd.DataFrame) -> dict:
1515
for column in df.columns:
16+
if df[column].dtype == object:
17+
# Convert Oracle LOB objects to string
18+
df[column] = df[column].apply(lambda x: str(x) if hasattr(x, "read") else x)
1619
if is_datetime64_any_dtype(df[column].dtype):
1720
df[column] = _to_datetime_and_format(df[column])
1821
return _to_json_obj(df)
@@ -44,6 +47,9 @@ def default(obj):
4447
return _date_offset_to_str(obj)
4548
if isinstance(obj, datetime.timedelta):
4649
return str(obj)
50+
# Add handling for any remaining LOB objects
51+
if hasattr(obj, "read"): # Check if object is LOB-like
52+
return str(obj)
4753
raise TypeError
4854

4955
json_obj = orjson.loads(

‎ibis-server/pyproject.toml

+3
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ ibis-framework = { version = "9.5.0", extras = [
1515
"clickhouse",
1616
"mssql",
1717
"mysql",
18+
"oracle",
1819
"postgres",
1920
"snowflake",
2021
"trino",
@@ -33,6 +34,7 @@ gql = { extras = ["aiohttp"], version = "3.5.0" }
3334
anyio = "4.8.0"
3435
duckdb = "1.1.3"
3536
opendal = ">=0.45"
37+
oracledb = "2.5.1"
3638

3739
[tool.poetry.group.dev.dependencies]
3840
pytest = "8.3.4"
@@ -61,6 +63,7 @@ markers = [
6163
"functions: mark a test as a functions test",
6264
"mssql: mark a test as a mssql test",
6365
"mysql: mark a test as a mysql test",
66+
"oracle: mark a test as a oracle test",
6467
"postgres: mark a test as a postgres test",
6568
"snowflake: mark a test as a snowflake test",
6669
"trino: mark a test as a trino test",

‎ibis-server/tests/routers/v2/connector/test_oracle.py

+134
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import base64
2+
3+
import orjson
4+
import pandas as pd
5+
import pytest
6+
import sqlalchemy
7+
from sqlalchemy import text
8+
from testcontainers.oracle import OracleDbContainer
9+
10+
from tests.conftest import file_path
11+
12+
pytestmark = pytest.mark.oracle
13+
14+
base_url = "/v2/connector/oracle"
15+
16+
manifest = {
17+
"catalog": "my_catalog",
18+
"schema": "my_schema",
19+
"models": [
20+
{
21+
"name": "Orders",
22+
"tableReference": {
23+
"schema": "SYSTEM",
24+
"table": "ORDERS",
25+
},
26+
"columns": [
27+
{"name": "orderkey", "expression": "O_ORDERKEY", "type": "number"},
28+
{"name": "custkey", "expression": "O_CUSTKEY", "type": "number"},
29+
{
30+
"name": "orderstatus",
31+
"expression": "O_ORDERSTATUS",
32+
"type": "varchar2",
33+
},
34+
{
35+
"name": "totalprice",
36+
"expression": "O_TOTALPRICE",
37+
"type": "number",
38+
},
39+
{"name": "orderdate", "expression": "O_ORDERDATE", "type": "date"},
40+
{
41+
"name": "order_cust_key",
42+
"expression": "O_ORDERKEY || '_' || O_CUSTKEY",
43+
"type": "varchar2",
44+
},
45+
{
46+
"name": "timestamp",
47+
"expression": "CAST('2024-01-01 23:59:59' AS TIMESTAMP)",
48+
"type": "timestamp",
49+
},
50+
{
51+
"name": "test_null_time",
52+
"expression": "CAST(NULL AS TIMESTAMP)",
53+
"type": "timestamp",
54+
},
55+
{
56+
"name": "blob_column",
57+
"expression": "UTL_RAW.CAST_TO_RAW('abc')",
58+
"type": "blob",
59+
},
60+
],
61+
"primaryKey": "orderkey",
62+
}
63+
],
64+
}
65+
66+
67+
@pytest.fixture(scope="module")
68+
def manifest_str():
69+
return base64.b64encode(orjson.dumps(manifest)).decode("utf-8")
70+
71+
72+
@pytest.fixture(scope="module")
73+
def oracle(request) -> OracleDbContainer:
74+
oracle = OracleDbContainer(
75+
"gvenzl/oracle-free:23.6-slim-faststart", oracle_password="Oracle123"
76+
)
77+
78+
oracle.start()
79+
80+
host = oracle.get_container_host_ip()
81+
port = oracle.get_exposed_port(1521)
82+
connection_url = (
83+
f"oracle+oracledb://SYSTEM:Oracle123@{host}:{port}/?service_name=FREEPDB1"
84+
)
85+
engine = sqlalchemy.create_engine(connection_url, echo=True)
86+
87+
with engine.begin() as conn:
88+
pd.read_parquet(file_path("resource/tpch/data/orders.parquet")).to_sql(
89+
"orders", engine, index=False
90+
)
91+
pd.read_parquet(file_path("resource/tpch/data/customer.parquet")).to_sql(
92+
"customer", engine, index=False
93+
)
94+
# Add table and column comments
95+
conn.execute(text("COMMENT ON TABLE orders IS 'This is a table comment'"))
96+
conn.execute(text("COMMENT ON COLUMN orders.o_comment IS 'This is a comment'"))
97+
98+
request.addfinalizer(oracle.stop)
99+
return oracle
100+
101+
102+
async def test_query_with_connection_url(
103+
client, manifest_str, oracle: OracleDbContainer
104+
):
105+
connection_url = _to_connection_url(oracle)
106+
response = await client.post(
107+
url=f"{base_url}/query",
108+
json={
109+
"connectionInfo": {"connectionUrl": connection_url},
110+
"manifestStr": manifest_str,
111+
"sql": "SELECT * FROM SYSTEM.ORDERS LIMIT 1",
112+
},
113+
)
114+
assert response.status_code == 200
115+
result = response.json()
116+
assert len(result["columns"]) == len(manifest["models"][0]["columns"])
117+
assert len(result["data"]) == 1
118+
assert result["data"][0][0] == 1
119+
assert result["dtypes"] is not None
120+
121+
122+
def _to_connection_info(oracle: OracleDbContainer):
123+
return {
124+
"host": oracle.get_container_host_ip(),
125+
"port": oracle.get_exposed_port(oracle.port),
126+
"user": "SYSTEM",
127+
"password": "Oracle123",
128+
"service": "FREEPDB1",
129+
}
130+
131+
132+
def _to_connection_url(oracle: OracleDbContainer):
133+
info = _to_connection_info(oracle)
134+
return f"oracle://{info['user']}:{info['password']}@{info['host']}:{info['port']}/{info['service']}"

0 commit comments

Comments
 (0)