Skip to content

Commit

Permalink
ci(athena): use internal ibis APIs for creating testing data (#10931)
Browse files Browse the repository at this point in the history
Clean up the test data loading in the athena conftest.py module
  • Loading branch information
cpcloud authored Mar 3, 2025
1 parent 1829169 commit 83cb363
Showing 1 changed file with 8 additions and 28 deletions.
36 changes: 8 additions & 28 deletions ibis/backends/athena/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@

import pytest
import sqlglot as sg
import sqlglot.expressions as sge
from sqlglot.dialects import Athena

import ibis
Expand Down Expand Up @@ -46,40 +45,21 @@ def create_table(connection, *, fs: s3fs.S3FileSystem, file: Path, folder: str)

arrow_schema = pq.read_metadata(file).schema.to_arrow_schema()
ibis_schema = PyArrowSchema.to_ibis(arrow_schema)
sg_schema = ibis_schema.to_sqlglot(Athena)
name = file.with_suffix("").name

ddl = sge.Create(
kind="TABLE",
this=sge.Schema(
this=sg.table(name, db=IBIS_ATHENA_TEST_DATABASE, quoted=True),
expressions=sg_schema,
),
properties=sge.Properties(
expressions=[
sge.ExternalProperty(),
sge.FileFormatProperty(this=sge.Var(this="PARQUET")),
sge.LocationProperty(this=sge.convert(f"{folder}/{name}")),
]
),
)

fs.put(str(file), f"{folder.removeprefix('s3://')}/{name}/{file.name}")

drop_query = sge.Drop(
kind="TABLE", this=sg.table(name, db=IBIS_ATHENA_TEST_DATABASE), exists=True
).sql(Athena)

create_query = ddl.sql(Athena)
connection.drop_table(name, database=IBIS_ATHENA_TEST_DATABASE, force=True)

with connection.con.cursor() as cur:
cur.execute(drop_query)
cur.execute(create_query)

assert (
connection.table(name, database=IBIS_ATHENA_TEST_DATABASE).count().execute() > 0
t = connection.create_table(
name,
schema=ibis_schema,
location=f"{folder}/{name}",
database=IBIS_ATHENA_TEST_DATABASE,
)

assert t.count().execute() > 0


class TestConf(BackendTest):
supports_map = False
Expand Down

0 comments on commit 83cb363

Please sign in to comment.