Skip to content

Commit abb0a43

Browse files
authored
chore(provision): make Spark integration test provisioning idempotent (#198)
1 parent 9e59782 commit abb0a43

1 file changed

Lines changed: 38 additions & 0 deletions

File tree

dev/spark/provision.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,50 @@
2020
# Provisions Paimon tables into the warehouse (file:/tmp/paimon-warehouse)
2121
# for paimon-rust integration tests to read.
2222

23+
import shutil
24+
from pathlib import Path
25+
from urllib.parse import unquote, urlparse
26+
2327
from pyspark.sql import SparkSession
2428

2529

30+
def _warehouse_path_from_spark_conf(spark: SparkSession) -> Path:
31+
warehouse_uri = spark.conf.get("spark.sql.catalog.paimon.warehouse")
32+
parsed = urlparse(warehouse_uri)
33+
34+
if parsed.scheme not in ("", "file"):
35+
raise ValueError(
36+
f"Unsupported Paimon warehouse URI scheme {parsed.scheme!r}: {warehouse_uri}"
37+
)
38+
39+
if parsed.netloc not in ("", "localhost"):
40+
raise ValueError(
41+
f"Unsupported remote Paimon warehouse location {parsed.netloc!r}: {warehouse_uri}"
42+
)
43+
44+
warehouse_path = Path(unquote(parsed.path if parsed.scheme else warehouse_uri))
45+
if not warehouse_path.is_absolute() or str(warehouse_path) == "/":
46+
raise ValueError(f"Refusing to clear unsafe warehouse path: {warehouse_path}")
47+
48+
return warehouse_path
49+
50+
51+
def _reset_warehouse_dir(warehouse_path: Path) -> None:
52+
warehouse_path.mkdir(parents=True, exist_ok=True)
53+
54+
for child in warehouse_path.iterdir():
55+
if child.is_symlink() or child.is_file():
56+
child.unlink()
57+
else:
58+
shutil.rmtree(child)
59+
60+
2661
def main():
2762
spark = SparkSession.builder.getOrCreate()
2863

64+
warehouse_path = _warehouse_path_from_spark_conf(spark)
65+
_reset_warehouse_dir(warehouse_path)
66+
2967
# Use Paimon catalog (configured in spark-defaults.conf with warehouse file:/tmp/paimon-warehouse)
3068
spark.sql("USE paimon.default")
3169

0 commit comments

Comments
 (0)