bring test PR in line with 12721

Author: glasnt

alloydb/notebooks/embeddings_batch_processing.ipynb

@@ -73,7 +73,14 @@
    },
    "outputs": [],
    "source": [
-    "%pip install google-cloud-alloydb-connector[asyncpg]==1.4.0 sqlalchemy==2.0.36 pandas==2.2.3 vertexai==1.70.0 asyncio==3.4.3 greenlet==3.1.1 --quiet"
+    "%pip install \\\n",
+    "    google-cloud-alloydb-connector[asyncpg]==1.4.0 \\\n",
+    "    sqlalchemy==2.0.36 \\\n",
+    "    pandas==2.2.3 \\\n",
+    "    vertexai==1.70.0 \\\n",
+    "    asyncio==3.4.3 \\\n",
+    "    greenlet==3.1.1 \\\n",
+    "    --quiet"
    ]
   },
   {
@@ -191,7 +198,6 @@
     "instance_name = \"my-primary\"  # @param {type:\"string\"}\n",
     "database_name = \"test_db\"  # @param {type:\"string\"}\n",
     "table_name = \"investments\"\n",
-    "username = \"postgres\"\n",
     "password = input(\"Please provide a password to be used for 'postgres' database user: \")"
    ]
   },
@@ -326,7 +332,7 @@
     "        conn: asyncpg.Connection = await connector.connect(\n",
     "            connection_string,\n",
     "            \"asyncpg\",\n",
-    "            user=username,\n",
+    "            user=\"postgres\",\n",
     "            password=password,\n",
     "            db=db_name,\n",
     "            ip_type=IPTypes.PUBLIC,\n",
@@ -367,11 +373,9 @@
    },
    "outputs": [],
    "source": [
-    "# debug: remove custom database.\n",
-    "\"\"\"\n",
     "from sqlalchemy import text, exc\n",
     "\n",
-    "async def create_db(database_name, connector):\n",
+    "async def create_db(database_name, connector):    \n",
     "    pool = await init_connection_pool(connector, \"postgres\")\n",
     "    async with pool.connect() as conn:\n",
     "        try:\n",
@@ -382,8 +386,7 @@
     "        except exc.ProgrammingError as e:\n",
     "          print(e)\n",
     "\n",
-    "await create_db(database_name=database_name, connector=connector)\n",
-    "\"\"\""
+    "await create_db(database_name=database_name, connector=connector)"
    ]
   },
   {
@@ -774,7 +777,7 @@
     "\n",
     "**embed_text**\n",
     "\n",
-    "This function your text data and sends it to vertex AI, transforming the text in specific columns into embeddings.\n",
+    "This function your text data and sends it to Vertex AI, transforming the text in specific columns into embeddings.\n",
     "\n",
     "**embed_objects_concurrently**\n",
     "\n",
@@ -868,9 +871,7 @@
     "                logger.warning(f\"Error: {e}. Retrying in {delay} seconds...\")\n",
     "                await asyncio.sleep(delay)  # Wait before retrying\n",
     "            else:\n",
-    "                logger.error(f\"Failed to get embeddings after {retries} attempts.\")\n",
-    "                raise  # Raise the error if all retries fail\n",
-    "\n",
+    "                logger.error(f\"Failed to get embeddings for data: {batch_data} after {retries} attempts.\")\n",
     "    return []\n",
     "\n",
     "\n",

alloydb/notebooks/e2e_test.py …/embeddings_batch_processing_e2e_test.pyalloydb/notebooks/e2e_test.py renamed to alloydb/notebooks/embeddings_batch_processing_e2e_test.py alloydb/notebooks/e2e_test.py renamed to alloydb/notebooks/embeddings_batch_processing_e2e_test.py

@@ -18,9 +18,9 @@
 
 import asyncpg  # type: ignore
 import conftest as conftest  # python-docs-samples/alloydb/conftest.py
+from google.cloud.alloydb.connector import AsyncConnector, IPTypes
 import pytest
 import sqlalchemy
-from google.cloud.alloydb.connector import AsyncConnector, IPTypes
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
 
@@ -43,7 +43,11 @@ async def _init_connection_pool(
     region: str,
     password: str,
 ) -> AsyncEngine:
-    connection_string = f"projects/{project_id}/locations/{region}/clusters/{cluster_name}/instances/{instance_name}"
+    connection_string = (
+        f"projects/{project_id}/locations/"
+        f"{region}/clusters/{cluster_name}/"
+        f"instances/{instance_name}"
+    )
 
     async def getconn() -> asyncpg.Connection:
         conn: asyncpg.Connection = await connector.connect(
@@ -71,27 +75,34 @@ async def test_embeddings_batch_processing(
     instance_name: str,
     region: str,
     database_name: str,
-    username: str,
     password: str,
     table_name: str,
 ) -> None:
+    # TODO: Create new table
     # Populate the table with embeddings by running the notebook
     conftest.run_notebook(
         "embeddings_batch_processing.ipynb",
         variables={
             "project_id": project_id,
             "cluster_name": cluster_name,
             "database_name": database_name,
-            "username": username,
             "region": region,
             "instance_name": instance_name,
             "table_name": table_name,
         },
         preprocess=preprocess,
         skip_shell_commands=True,
         replace={
-            "password = input(\"Please provide a password to be used for 'postgres' database user: \")": f"password = '{password}'",
-            "await create_db(database_name=database_name, connector=connector)": "",
+            (
+                "password = input(\"Please provide "
+                "a password to be used for 'postgres' "
+                "database user: \")"
+            ): f"password = '{password}'",
+            (
+                "await create_db("
+                "database_name=database_name, "
+                "connector=connector)"
+            ): "",
         },
         until_end=True,
     )
@@ -111,25 +122,35 @@ async def test_embeddings_batch_processing(
             # Validate that embeddings are non-empty for all rows
             result = await conn.execute(
                 sqlalchemy.text(
-                    f"SELECT COUNT(*) FROM {table_name} WHERE analysis_embedding IS NULL"
+                    f"SELECT COUNT(*) FROM "
+                    f"{table_name} WHERE "
+                    f"analysis_embedding IS NULL"
                 )
             )
             row = result.fetchone()
             assert row[0] == 0
             result = await conn.execute(
                 sqlalchemy.text(
-                    f"SELECT COUNT(*) FROM {table_name} WHERE overview_embedding IS NULL"
+                    f"SELECT COUNT(*) FROM "
+                    f"{table_name} WHERE "
+                    f"overview_embedding IS NULL"
                 )
             )
             row = result.fetchone()
             assert row[0] == 0
 
             # Get the table back to the original state
             await conn.execute(
-                sqlalchemy.text(f"UPDATE {table_name} set analysis_embedding = NULL")
+                sqlalchemy.text(
+                    f"UPDATE {table_name} set "
+                    f"analysis_embedding = NULL"
+                )
             )
             await conn.execute(
-                sqlalchemy.text(f"UPDATE {table_name} set overview_embedding = NULL")
+                sqlalchemy.text(
+                    f"UPDATE {table_name} set "
+                    f"overview_embedding = NULL"
+                )
             )
             await conn.commit()
         await pool.dispose()

alloydb/notebooks/requirements.txt

@@ -1 +0,0 @@
-# empty