# Copyright 2025 Google LLC

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import asyncio

import os

import uuid

from typing import Any, Coroutine, Sequence

import asyncpg # type: ignore

import pytest

import pytest_asyncio

from google.cloud.sql.connector import Connector, IPTypes

from langchain_core.embeddings import DeterministicFakeEmbedding

from sqlalchemy import VARCHAR, text

from sqlalchemy.engine import URL

from sqlalchemy.engine.row import RowMapping

from sqlalchemy.ext.asyncio import create_async_engine

from sqlalchemy.pool import NullPool

from langchain_google_cloud_sql_pg import Column, HybridSearchConfig, PostgresEngine

DEFAULT_TABLE = "test_table" + str(uuid.uuid4()).replace("-", "_")

CUSTOM_TABLE = "test_table_custom" + str(uuid.uuid4()).replace("-", "_")

INT_ID_CUSTOM_TABLE = "test_table_custom_int_id" + str(uuid.uuid4()).replace("-", "_")

HYBRID_SEARCH_TABLE = "hybrid" + str(uuid.uuid4()).replace("-", "_")

DEFAULT_TABLE_SYNC = "test_table" + str(uuid.uuid4()).replace("-", "_")

CUSTOM_TABLE_SYNC = "test_table_custom" + str(uuid.uuid4()).replace("-", "_")

INT_ID_CUSTOM_TABLE_SYNC = "test_table_custom_int_id" + str(uuid.uuid4()).replace(

"-", "_"

)

HYBRID_SEARCH_TABLE_SYNC = "hybrid_sync" + str(uuid.uuid4()).replace("-", "_")

VECTOR_SIZE = 768

embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE)

host = os.environ["IP_ADDRESS"]

def get_env_var(key: str, desc: str) -> str:

v = os.environ.get(key)

if v is None:

raise ValueError(f"Must set env var {key} to: {desc}")

return v

# Helper to bridge the Main Test Loop and the Engine Background Loop

async def run_on_background(engine: PostgresEngine, coro: Coroutine) -> Any:

"""Runs a coroutine on the engine's background loop (if it exists)."""

if engine._loop:

return await asyncio.wrap_future(

asyncio.run_coroutine_threadsafe(coro, engine._loop)

)

return await coro

async def aexecute(

engine: PostgresEngine,

query: str,

) -> None:

async def _impl():

async with engine._pool.connect() as conn:

await conn.execute(text(query))

await conn.commit()

await run_on_background(engine, _impl())

async def afetch(engine: PostgresEngine, query: str) -> Sequence[RowMapping]:

async def _impl():

async with engine._pool.connect() as conn:

result = await conn.execute(text(query))

result_map = result.mappings()

return result_map.fetchall()

return await run_on_background(engine, _impl())

@pytest.mark.asyncio(scope="module")

class TestEngineAsync:

@pytest.fixture(scope="module")

def db_project(self) -> str:

return get_env_var("PROJECT_ID", "project id for google cloud")

@pytest.fixture(scope="module")

def db_region(self) -> str:

return get_env_var("REGION", "region for cloud sql instance")

@pytest.fixture(scope="module")

def db_instance(self) -> str:

return get_env_var("INSTANCE_ID", "instance for cloud sql")

@pytest.fixture(scope="module")

def db_name(self) -> str:

return get_env_var("DATABASE_ID", "instance for cloud sql")

@pytest.fixture(scope="module")

def user(self) -> str:

return get_env_var("DB_USER", "database user for cloud sql")

@pytest.fixture(scope="module")

def password(self) -> str:

return get_env_var("DB_PASSWORD", "database password for cloud sql")

@pytest.fixture(scope="module")

def iam_account(self) -> str:

return get_env_var("IAM_ACCOUNT", "Cloud SQL IAM account email")

@pytest_asyncio.fixture(scope="class")

async def engine(self, db_project, db_region, db_instance, db_name):

engine = await PostgresEngine.afrom_instance(

project_id=db_project,

instance=db_instance,

region=db_region,

database=db_name,

engine_args={

# add some connection args to validate engine_args works correctly

"pool_size": 3,

"max_overflow": 2,

},

)

yield engine

await aexecute(engine, f'DROP TABLE "{CUSTOM_TABLE}"')

await aexecute(engine, f'DROP TABLE "{DEFAULT_TABLE}"')

await aexecute(engine, f'DROP TABLE "{INT_ID_CUSTOM_TABLE}"')

await aexecute(engine, f'DROP TABLE "{HYBRID_SEARCH_TABLE}"')

await engine.close()

async def test_engine_args(self, engine):

# Accessing engine._pool.pool.status() is synchronous and safe on main loop objects

# assuming SQLAlchemy pool status doesn't strictly require loop context

assert "Pool size: 3" in engine._pool.pool.status()

async def test_init_table(self, engine):

await run_on_background(

engine, engine.ainit_vectorstore_table(DEFAULT_TABLE, VECTOR_SIZE)

)

id = str(uuid.uuid4())

content = "coffee"

embedding = await embeddings_service.aembed_query(content)

# Note: DeterministicFakeEmbedding generates a numpy array, converting to list a list of float values

embedding_string = [float(dimension) for dimension in embedding]

stmt = f"INSERT INTO {DEFAULT_TABLE} (langchain_id, content, embedding) VALUES ('{id}', '{content}','{embedding_string}');"

await aexecute(engine, stmt)

async def test_init_table_custom(self, engine):

await run_on_background(

engine,

engine.ainit_vectorstore_table(

CUSTOM_TABLE,

VECTOR_SIZE,

id_column="uuid",

content_column="my-content",

embedding_column="my_embedding",

metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],

store_metadata=True,

),

)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{CUSTOM_TABLE}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "uuid", "data_type": "uuid"},

{"column_name": "my_embedding", "data_type": "USER-DEFINED"},

{"column_name": "langchain_metadata", "data_type": "json"},

{"column_name": "my-content", "data_type": "text"},

{"column_name": "page", "data_type": "text"},

{"column_name": "source", "data_type": "text"},

]

for row in results:

assert row in expected

async def test_init_table_with_int_id(self, engine):

await run_on_background(

engine,

engine.ainit_vectorstore_table(

INT_ID_CUSTOM_TABLE,

VECTOR_SIZE,

id_column=Column(

name="integer_id", data_type="INTEGER", nullable="False"

),

content_column="my-content",

embedding_column="my_embedding",

metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],

store_metadata=True,

),

)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{INT_ID_CUSTOM_TABLE}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "integer_id", "data_type": "integer"},

{"column_name": "my_embedding", "data_type": "USER-DEFINED"},

{"column_name": "langchain_metadata", "data_type": "json"},

{"column_name": "my-content", "data_type": "text"},

{"column_name": "page", "data_type": "text"},

{"column_name": "source", "data_type": "text"},

]

for row in results:

assert row in expected

async def test_password(

self,

db_project,

db_region,

db_instance,

db_name,

user,

password,

):

# Note: PostgresEngine._connector is no longer a class attribute in fixed engine.py

# But for test cleanup safety regarding the OLD code structure, we can ignore this.

# PostgresEngine._connector = None

engine = await PostgresEngine.afrom_instance(

project_id=db_project,

instance=db_instance,

region=db_region,

database=db_name,

user=user,

password=password,

)

assert engine

await aexecute(engine, "SELECT 1")

await engine.close()

async def test_from_engine(

self,

db_project,

db_region,

db_instance,

db_name,

user,

password,

):

async with Connector(loop=asyncio.get_running_loop()) as connector:

async def getconn() -> asyncpg.Connection:

conn = await connector.connect_async( # type: ignore

f"{db_project}:{db_region}:{db_instance}",

"asyncpg",

user=user,

password=password,

db=db_name,

enable_iam_auth=False,

ip_type=IPTypes.PUBLIC,

)

return conn

engine_async = create_async_engine(

"postgresql+asyncpg://",

async_creator=getconn,

)

engine = PostgresEngine.from_engine(engine_async)

await aexecute(engine, "SELECT 1")

await engine.close()

async def test_from_connection_string(

self,

db_name,

user,

password,

):

port = "5432"

url = f"postgresql+asyncpg://{user}:{password}@{host}:{port}/{db_name}"

engine = PostgresEngine.from_connection_string(

url,

echo=True,

poolclass=NullPool,

)

await aexecute(engine, "SELECT 1")

await engine.close()

engine = PostgresEngine.from_connection_string(

URL.create("postgresql+asyncpg", user, password, host, port, db_name)

)

await aexecute(engine, "SELECT 1")

await engine.close()

async def test_from_engine_args_url(

self,

db_name,

user,

password,

):

port = "5432"

url = f"postgresql+asyncpg://{user}:{password}@{host}:{port}/{db_name}"

engine = PostgresEngine.from_engine_args(

url,

echo=True,

poolclass=NullPool,

)

await aexecute(engine, "SELECT 1")

await engine.close()

engine = PostgresEngine.from_engine_args(

URL.create("postgresql+asyncpg", user, password, host, port, db_name)

)

await aexecute(engine, "SELECT 1")

await engine.close()

async def test_from_engine_args_url_error(

self,

db_name,

user,

password,

):

port = "5432"

url = f"postgresql+asyncpg://{user}:{password}@{host}:{port}/{db_name}"

with pytest.raises(TypeError):

engine = PostgresEngine.from_engine_args(url, random=False)

with pytest.raises(ValueError):

PostgresEngine.from_engine_args(

f"postgresql+pg8000://{user}:{password}@{host}:{port}/{db_name}",

)

with pytest.raises(ValueError):

PostgresEngine.from_engine_args(

URL.create("postgresql+pg8000", user, password, host, port, db_name)

)

async def test_column(self, engine):

with pytest.raises(ValueError):

Column("test", VARCHAR)

with pytest.raises(ValueError):

Column(1, "INTEGER")

async def test_iam_account_override(

self,

db_project,

db_instance,

db_region,

db_name,

iam_account,

engine,

):

engine = await PostgresEngine.afrom_instance(

project_id=db_project,

instance=db_instance,

region=db_region,

database=db_name,

iam_account_email=iam_account,

)

assert engine

await aexecute(engine, "SELECT 1")

await engine.close()

async def test_ainit_checkpoint_writes_table(self, engine):

table_name = f"checkpoint{uuid.uuid4()}"

table_name_writes = f"{table_name}_writes"

await run_on_background(

engine, engine.ainit_checkpoint_table(table_name=table_name)

)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name_writes}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "thread_id", "data_type": "text"},

{"column_name": "checkpoint_ns", "data_type": "text"},

{"column_name": "checkpoint_id", "data_type": "text"},

{"column_name": "task_id", "data_type": "text"},

{"column_name": "idx", "data_type": "integer"},

{"column_name": "channel", "data_type": "text"},

{"column_name": "type", "data_type": "text"},

{"column_name": "blob", "data_type": "bytea"},

{"column_name": "task_path", "data_type": "text"},

]

for row in results:

assert row in expected

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "thread_id", "data_type": "text"},

{"column_name": "checkpoint_ns", "data_type": "text"},

{"column_name": "checkpoint_id", "data_type": "text"},

{"column_name": "parent_checkpoint_id", "data_type": "text"},

{"column_name": "type", "data_type": "text"},

{"column_name": "checkpoint", "data_type": "bytea"},

{"column_name": "metadata", "data_type": "bytea"},

]

for row in results:

assert row in expected

await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}"')

await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name_writes}"')

async def test_init_table_hybrid_search(self, engine):

await run_on_background(

engine,

engine.ainit_vectorstore_table(

HYBRID_SEARCH_TABLE,

VECTOR_SIZE,

id_column="uuid",

content_column="my-content",

embedding_column="my_embedding",

metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],

store_metadata=True,

hybrid_search_config=HybridSearchConfig(),

),

)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{HYBRID_SEARCH_TABLE}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "uuid", "data_type": "uuid"},

{"column_name": "my_embedding", "data_type": "USER-DEFINED"},

{"column_name": "langchain_metadata", "data_type": "json"},

{"column_name": "my-content", "data_type": "text"},

{"column_name": "my-content_tsv", "data_type": "tsvector"},

{"column_name": "page", "data_type": "text"},

{"column_name": "source", "data_type": "text"},

]

for row in results:

assert row in expected

@pytest.mark.asyncio(scope="module")

class TestEngineSync:

@pytest.fixture(scope="module")

def db_project(self) -> str:

return get_env_var("PROJECT_ID", "project id for google cloud")

@pytest.fixture(scope="module")

def db_region(self) -> str:

return get_env_var("REGION", "region for cloud sql instance")

@pytest.fixture(scope="module")

def db_instance(self) -> str:

return get_env_var("INSTANCE_ID", "instance for cloud sql")

@pytest.fixture(scope="module")

def db_name(self) -> str:

return get_env_var("DATABASE_ID", "instance for cloud sql")

@pytest.fixture(scope="module")

def user(self) -> str:

return get_env_var("DB_USER", "database user for cloud sql")

@pytest.fixture(scope="module")

def password(self) -> str:

return get_env_var("DB_PASSWORD", "database password for cloud sql")

@pytest.fixture(scope="module")

def iam_account(self) -> str:

return get_env_var("IAM_ACCOUNT", "Cloud SQL IAM account email")

@pytest_asyncio.fixture(scope="class")

async def engine(self, db_project, db_region, db_instance, db_name):

engine = PostgresEngine.from_instance(

project_id=db_project,

instance=db_instance,

region=db_region,

database=db_name,

)

yield engine

await aexecute(engine, f'DROP TABLE "{CUSTOM_TABLE_SYNC}"')

await aexecute(engine, f'DROP TABLE "{DEFAULT_TABLE_SYNC}"')

await aexecute(engine, f'DROP TABLE "{INT_ID_CUSTOM_TABLE_SYNC}"')

await aexecute(engine, f'DROP TABLE "{HYBRID_SEARCH_TABLE_SYNC}"')

await engine.close()

async def test_init_table(self, engine):

# Sync method uses _run_as_sync internally -> safe to call on Main Loop

engine.init_vectorstore_table(DEFAULT_TABLE_SYNC, VECTOR_SIZE)

id = str(uuid.uuid4())

content = "coffee"

embedding = await embeddings_service.aembed_query(content)

embedding_string = [float(dimension) for dimension in embedding]

stmt = f"INSERT INTO {DEFAULT_TABLE_SYNC} (langchain_id, content, embedding) VALUES ('{id}', '{content}','{embedding_string}');"

await aexecute(engine, stmt)

async def test_init_table_custom(self, engine):

engine.init_vectorstore_table(

CUSTOM_TABLE_SYNC,

VECTOR_SIZE,

id_column="uuid",

content_column="my-content",

embedding_column="my_embedding",

metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],

store_metadata=True,

)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{CUSTOM_TABLE_SYNC}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "uuid", "data_type": "uuid"},

{"column_name": "my_embedding", "data_type": "USER-DEFINED"},

{"column_name": "langchain_metadata", "data_type": "json"},

{"column_name": "my-content", "data_type": "text"},

{"column_name": "page", "data_type": "text"},

{"column_name": "source", "data_type": "text"},

]

for row in results:

assert row in expected

async def test_init_table_with_int_id(self, engine):

engine.init_vectorstore_table(

INT_ID_CUSTOM_TABLE_SYNC,

VECTOR_SIZE,

id_column=Column(name="integer_id", data_type="INTEGER", nullable=False),

content_column="my-content",

embedding_column="my_embedding",

metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],

store_metadata=True,

)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{INT_ID_CUSTOM_TABLE_SYNC}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "integer_id", "data_type": "integer"},

{"column_name": "my_embedding", "data_type": "USER-DEFINED"},

{"column_name": "langchain_metadata", "data_type": "json"},

{"column_name": "my-content", "data_type": "text"},

{"column_name": "page", "data_type": "text"},

{"column_name": "source", "data_type": "text"},

]

for row in results:

assert row in expected

async def test_password(

self,

db_project,

db_region,

db_instance,

db_name,

user,

password,

):

engine = PostgresEngine.from_instance(

project_id=db_project,

instance=db_instance,

region=db_region,

database=db_name,

user=user,

password=password,

quota_project=db_project,

)

assert engine

await aexecute(engine, "SELECT 1")

await engine.close()

async def test_engine_constructor_key(

self,

engine,

):

key = object()

with pytest.raises(Exception):

PostgresEngine(key, engine, None, None)

async def test_iam_account_override(

self,

db_project,

db_instance,

db_region,

db_name,

iam_account,

engine,

):

engine = PostgresEngine.from_instance(

project_id=db_project,

instance=db_instance,

region=db_region,

database=db_name,

iam_account_email=iam_account,

)

assert engine

await aexecute(engine, "SELECT 1")

await engine.close()

async def test_init_checkpoints_table(self, engine):

table_name = f"checkpoint{uuid.uuid4()}"

table_name_writes = f"{table_name}_writes"

engine.init_checkpoint_table(table_name=table_name)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "thread_id", "data_type": "text"},

{"column_name": "checkpoint_ns", "data_type": "text"},

{"column_name": "checkpoint_id", "data_type": "text"},

{"column_name": "parent_checkpoint_id", "data_type": "text"},

{"column_name": "type", "data_type": "text"},

{"column_name": "checkpoint", "data_type": "bytea"},

{"column_name": "metadata", "data_type": "bytea"},

]

for row in results:

assert row in expected

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name_writes}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "thread_id", "data_type": "text"},

{"column_name": "checkpoint_ns", "data_type": "text"},

{"column_name": "checkpoint_id", "data_type": "text"},

{"column_name": "task_id", "data_type": "text"},

{"column_name": "idx", "data_type": "integer"},

{"column_name": "channel", "data_type": "text"},

{"column_name": "type", "data_type": "text"},

{"column_name": "blob", "data_type": "bytea"},

{"column_name": "task_path", "data_type": "text"},

]

for row in results:

assert row in expected

await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}"')

await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name_writes}"')

async def test_init_table_hybrid_search(self, engine):

engine.init_vectorstore_table(

HYBRID_SEARCH_TABLE_SYNC,

VECTOR_SIZE,

id_column="uuid",

content_column="my-content",

embedding_column="my_embedding",

metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],

store_metadata=True,

hybrid_search_config=HybridSearchConfig(),

)

stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{HYBRID_SEARCH_TABLE_SYNC}';"

results = await afetch(engine, stmt)

expected = [

{"column_name": "uuid", "data_type": "uuid"},

{"column_name": "my_embedding", "data_type": "USER-DEFINED"},

{"column_name": "langchain_metadata", "data_type": "json"},

{"column_name": "my-content", "data_type": "text"},

{"column_name": "my-content_tsv", "data_type": "tsvector"},

{"column_name": "page", "data_type": "text"},

{"column_name": "source", "data_type": "text"},

]

for row in results:

assert row in expected