Repository navigation
Expand file tree
/
Copy pathconftest.py
More file actions
128 lines (111 loc) · 4.32 KB
/
Copy pathconftest.py
File metadata and controls
128 lines (111 loc) · 4.32 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
"""PostgreSQL setup for code samples.
This module provides PostgreSQL connection setup for code samples.
It attempts to use testcontainers if available, otherwise provides
utilities to work with environment-configured postgres.
"""
import os
import subprocess
import sys
import time
_DEFAULT_URI = "postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable"
def get_postgres_uri() -> str:
"""Get PostgreSQL connection URI.
Tries multiple approaches in order:
1. Check POSTGRES_URI environment variable
2. Attempt to use testcontainers to spin up postgres
3. Try docker directly
4. Fall back to default local postgres connection
"""
# Check environment variable first
if env_uri := os.environ.get("POSTGRES_URI"):
return env_uri
# Try testcontainers
try:
from testcontainers.postgres import ( # type: ignore[import-not-found]
PostgresContainer,
)
# Store container in a global so it persists
if not hasattr(get_postgres_uri, "_container"):
# Use pgvector image which includes the vector extension
container = PostgresContainer("pgvector/pgvector:pg17")
container.start()
get_postgres_uri._container = container # type: ignore[attr-defined]
# Give it a moment to fully start
time.sleep(2)
return get_postgres_uri._container.get_connection_url() # type: ignore[attr-defined]
except ImportError:
print("conftest: testcontainers not installed, trying docker", file=sys.stderr)
# Try to use docker directly if testcontainers not available
try:
# Check if postgres container is already running
result = subprocess.run(
[
"docker",
"ps",
"--filter",
"name=langchain-docs-postgres",
"--format",
"{{.Names}}",
],
capture_output=True,
text=True,
check=True,
timeout=5,
)
if "langchain-docs-postgres" not in result.stdout:
# Start a postgres container with pgvector extension
subprocess.run(
[
"docker",
"run",
"-d",
"--name",
"langchain-docs-postgres",
"-e",
"POSTGRES_PASSWORD=postgres",
"-e",
"POSTGRES_DB=postgres",
"-p",
"5432:5432",
"pgvector/pgvector:pg17",
],
check=True,
capture_output=True,
timeout=30,
)
# Give it time to start
time.sleep(3)
return _DEFAULT_URI
except FileNotFoundError:
print("conftest: docker not found, using default URI", file=sys.stderr)
except subprocess.TimeoutExpired as e:
print(
f"conftest: docker timed out after {e.timeout}s, using default URI",
file=sys.stderr,
)
except subprocess.CalledProcessError as e:
print(
f"conftest: docker failed (exit {e.returncode}), using default URI",
file=sys.stderr,
)
# Fall back to default (assumes postgres is running locally)
print(f"conftest: falling back to default URI: {_DEFAULT_URI}", file=sys.stderr)
return _DEFAULT_URI
def prepare_postgres_store(uri: str) -> None:
"""Drop existing store tables so setup() creates a fresh schema.
Use before PostgresStore.from_conn_string when tests share a database
(e.g. CI) and may see leftover tables from a different schema version.
"""
import psycopg
try:
with psycopg.connect(uri, autocommit=True) as conn:
with conn.cursor() as cur:
cur.execute(
"DROP TABLE IF EXISTS public.store_vectors CASCADE; "
"DROP TABLE IF EXISTS public.store CASCADE; "
"DROP TABLE IF EXISTS public.store_migrations CASCADE;"
)
except psycopg.OperationalError as e:
print(f"conftest: could not connect to clean tables: {e}", file=sys.stderr)
except psycopg.Error as e:
print(f"conftest: failed to drop store tables: {e}", file=sys.stderr)