Implemented basic test of fast import

This commit is contained in:
Gleb Novikov
2025-01-03 14:41:49 +00:00
parent e291fb7edc
commit 6e5a0add43
3 changed files with 90 additions and 146 deletions

View File

@@ -15,4 +15,5 @@ pytest_plugins = (
"fixtures.compare_fixtures",
"fixtures.slow",
"fixtures.reruns",
"fixtures.fast_import",
)

View File

@@ -1,13 +1,28 @@
import shutil
import subprocess
import tempfile
from collections.abc import Iterator
from pathlib import Path
import pytest
from fixtures.log_helper import log
from fixtures.neon_cli import AbstractNeonCli
from fixtures.pg_version import PgVersion
class FastImport(AbstractNeonCli):
COMMAND = "fast_import"
cmd: subprocess.CompletedProcess[str] | None = None
def __init__(self, extra_env: dict[str, str] | None, binpath: Path, pg_dir: Path):
def __init__(
self,
extra_env: dict[str, str] | None,
binpath: Path,
pg_distrib_dir: Path,
pg_version: PgVersion,
workdir: Path
):
if extra_env is None:
env_vars = {}
else:
@@ -17,23 +32,31 @@ class FastImport(AbstractNeonCli):
raise Exception(f"{self.COMMAND} binary not found at '{binpath}'")
super().__init__(env_vars, binpath)
pg_dir = pg_distrib_dir / pg_version.v_prefixed
self.pg_distrib_dir = pg_distrib_dir
self.pg_version = pg_version
self.pg_bin = pg_dir / "bin"
if not (self.pg_bin / "postgres") .exists():
if not (self.pg_bin / "postgres").exists():
raise Exception(f"postgres binary was not found at '{self.pg_bin}'")
self.pg_lib = pg_dir / "lib"
if not workdir.exists():
raise Exception(f"Working directory '{workdir}' does not exist")
self.workdir = workdir
def run(self,
workdir: Path,
pg_port: int,
source_connection_string: str | None = None,
s3prefix: str | None = None,
interactive: bool = False,
) -> subprocess.CompletedProcess[str]:
) -> subprocess.CompletedProcess[str]:
if self.cmd is not None:
raise Exception("Command already executed")
args = [
f"--pg-bin-dir={self.pg_bin}",
f"--pg-lib-dir={self.pg_lib}",
f"--pg-port={pg_port}",
f"--working-directory={workdir}",
f"--working-directory={self.workdir}",
]
if source_connection_string is not None:
args.append(f"--source-connection-string={source_connection_string}")
@@ -42,4 +65,35 @@ class FastImport(AbstractNeonCli):
if interactive:
args.append("--interactive")
return self.raw_cli(args)
self.cmd = self.raw_cli(args)
return self.cmd
def __enter__(self):
return self
def __exit__(self, *args):
if self.workdir.exists():
shutil.rmtree(self.workdir)
@pytest.fixture(scope="function")
def fast_import(
pg_version: PgVersion,
test_output_dir: Path,
neon_binpath: Path,
pg_distrib_dir: Path,
) -> Iterator[FastImport]:
workdir = Path(tempfile.mkdtemp())
with FastImport(None, neon_binpath, pg_distrib_dir, pg_version, workdir) as fi:
yield fi
if fi.cmd is None:
return
# dump stdout & stderr into test log dir
with open(test_output_dir / "fast_import.stdout", "w") as f:
f.write(fi.cmd.stdout)
with open(test_output_dir / "fast_import.stderr", "w") as f:
f.write(fi.cmd.stderr)
log.info('Written logs to %s', test_output_dir)

View File

@@ -1,150 +1,39 @@
from __future__ import annotations
import shutil
import tempfile
from pathlib import Path
from fixtures.fast_import import FastImport
from fixtures.log_helper import log
from fixtures.neon_fixtures import VanillaPostgres, PgProtocol, PgBin
from fixtures.port_distributor import PortDistributor
#
# Create ancestor branches off the main branch.
#
def test_fast_import(port_distributor, test_output_dir: Path, neon_binpath: Path, pg_distrib_dir: Path):
def test_fast_import(
test_output_dir,
vanilla_pg: VanillaPostgres,
port_distributor: PortDistributor,
fast_import: FastImport,
):
vanilla_pg.start()
vanilla_pg.safe_psql("CREATE TABLE foo (a int); INSERT INTO foo SELECT generate_series(1, 10);")
# TODO:
# 1. fast_import fixture
# 2. maybe run vanilla postgres and insert some data into it
# 3. run fast_import in interactive mode in the background
# - wait for "interactive mode" message
# 4. run some query against the imported data (in running postgres inside fast_import, port is known)
pg_port = port_distributor.get_port()
fast_import.run(pg_port, vanilla_pg.connstr())
vanilla_pg.stop()
# Maybe test with pageserver?
# 1. run whole neon env
# 2. create timeline with some s3 path???
# 3. run fast_import with s3 prefix
# 4. ??? mock http where pageserver will report progress
# 5. run compute on this timeline and check if data is there
pgbin = PgBin(test_output_dir, fast_import.pg_distrib_dir, fast_import.pg_version)
new_pgdata_vanilla_pg = VanillaPostgres(fast_import.workdir / "pgdata", pgbin, pg_port, False)
new_pgdata_vanilla_pg.start()
# database name and user are hardcoded in fast_import binary, and they are different from normal vanilla postgres
conn = PgProtocol(dsn=f"postgresql://cloud_admin@localhost:{pg_port}/neondb")
res = conn.safe_psql("SELECT count(*) FROM foo;")
log.info(f"Result: {res}")
assert res[0][0] == 10
new_pgdata_vanilla_pg.stop()
fast_import = FastImport(
None,
neon_binpath,
pg_distrib_dir / "v16",
)
workdir = Path(tempfile.mkdtemp())
cmd = fast_import.run(
Path(tempfile.mkdtemp()),
port_distributor.get_port(),
"postgresql://gleb:O9UM4tPHEafC@ep-autumn-rain-a3smlr75.eu-central-1.aws.neon.tech/neondb?sslmode=require",
None,
)
# dump stdout & stderr into test log dir
with open(test_output_dir / "fast_import.stdout", "w") as f:
f.write(cmd.stdout)
with open(test_output_dir / "fast_import.stderr", "w") as f:
f.write(cmd.stderr)
log.info('Written logs to %s', test_output_dir)
# clean workdir
shutil.rmtree(workdir)
# env = neon_env_builder.init_start()
# pageserver_http = env.pageserver.http_client()
# # Override defaults: 4M checkpoint_distance, disable background compaction and gc.
# tenant, _ = env.create_tenant(
# conf={
# "checkpoint_distance": "4194304",
# "gc_period": "0s",
# "compaction_period": "0s",
# }
# )
#
# failpoint = "flush-frozen-pausable"
#
# pageserver_http.configure_failpoints((failpoint, "sleep(10000)"))
#
# endpoint_branch0 = env.endpoints.create_start("main", tenant_id=tenant)
# branch0_cur = endpoint_branch0.connect().cursor()
# branch0_timeline = TimelineId(query_scalar(branch0_cur, "SHOW neon.timeline_id"))
# log.info(f"b0 timeline {branch0_timeline}")
#
# # Create table, and insert 100k rows.
# branch0_lsn = query_scalar(branch0_cur, "SELECT pg_current_wal_insert_lsn()")
# log.info(f"b0 at lsn {branch0_lsn}")
#
# branch0_cur.execute("CREATE TABLE foo (t text) WITH (autovacuum_enabled = off)")
# branch0_cur.execute(
# """
# INSERT INTO foo
# SELECT '00112233445566778899AABBCCDDEEFF' || ':branch0:' || g
# FROM generate_series(1, 100000) g
# """
# )
# lsn_100 = query_scalar(branch0_cur, "SELECT pg_current_wal_insert_lsn()")
# log.info(f"LSN after 100k rows: {lsn_100}")
#
# # Create branch1.
# env.create_branch(
# "branch1", ancestor_branch_name="main", ancestor_start_lsn=lsn_100, tenant_id=tenant
# )
# endpoint_branch1 = env.endpoints.create_start("branch1", tenant_id=tenant)
#
# branch1_cur = endpoint_branch1.connect().cursor()
# branch1_timeline = TimelineId(query_scalar(branch1_cur, "SHOW neon.timeline_id"))
# log.info(f"b1 timeline {branch1_timeline}")
#
# branch1_lsn = query_scalar(branch1_cur, "SELECT pg_current_wal_insert_lsn()")
# log.info(f"b1 at lsn {branch1_lsn}")
#
# # Insert 100k rows.
# branch1_cur.execute(
# """
# INSERT INTO foo
# SELECT '00112233445566778899AABBCCDDEEFF' || ':branch1:' || g
# FROM generate_series(1, 100000) g
# """
# )
# lsn_200 = query_scalar(branch1_cur, "SELECT pg_current_wal_insert_lsn()")
# log.info(f"LSN after 200k rows: {lsn_200}")
#
# # Create branch2.
# env.create_branch(
# "branch2", ancestor_branch_name="branch1", ancestor_start_lsn=lsn_200, tenant_id=tenant
# )
# endpoint_branch2 = env.endpoints.create_start("branch2", tenant_id=tenant)
# branch2_cur = endpoint_branch2.connect().cursor()
#
# branch2_timeline = TimelineId(query_scalar(branch2_cur, "SHOW neon.timeline_id"))
# log.info(f"b2 timeline {branch2_timeline}")
#
# branch2_lsn = query_scalar(branch2_cur, "SELECT pg_current_wal_insert_lsn()")
# log.info(f"b2 at lsn {branch2_lsn}")
#
# # Insert 100k rows.
# branch2_cur.execute(
# """
# INSERT INTO foo
# SELECT '00112233445566778899AABBCCDDEEFF' || ':branch2:' || g
# FROM generate_series(1, 100000) g
# """
# )
# lsn_300 = query_scalar(branch2_cur, "SELECT pg_current_wal_insert_lsn()")
# log.info(f"LSN after 300k rows: {lsn_300}")
#
# # Run compaction on branch1.
# compact = f"compact {tenant} {branch1_timeline}"
# log.info(compact)
# pageserver_http.timeline_compact(tenant, branch1_timeline)
#
# assert query_scalar(branch0_cur, "SELECT count(*) FROM foo") == 100000
#
# assert query_scalar(branch1_cur, "SELECT count(*) FROM foo") == 200000
#
# assert query_scalar(branch2_cur, "SELECT count(*) FROM foo") == 300000
#
# pageserver_http.configure_failpoints((failpoint, "off"))
# TODO: Maybe test with pageserver?
# 1. run whole neon env
# 2. create timeline with some s3 path???
# 3. run fast_import with s3 prefix
# 4. ??? mock http where pageserver will report progress
# 5. run compute on this timeline and check if data is there