forked from langchain-ai/docs
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathconftest.py
More file actions
132 lines (115 loc) · 4.62 KB
/
Copy pathconftest.py
File metadata and controls
132 lines (115 loc) · 4.62 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
129
130
131
132
"""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:
# Drop vector_migrations too. setup() skips VECTOR_MIGRATIONS when
# that table is present, which leaves store_vectors missing after a
# partial reset (put then fails with UndefinedTable).
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; "
"DROP TABLE IF EXISTS public.vector_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)