# Apache Beam with GCP support 
 
 pip  
install  
apache_beam [ 
gcp ] 
> = 
v2.64.0  
--quiet 
  # Huggingface sentence-transformers for embedding models 
 
 pip  
install  
sentence-transformers  
--quiet 

 pip  
install  
 "google-cloud-alloydb-connector[pg8000]" 
  
sqlalchemy 

  PROJECT_ID 
 = 
 "" 
 # @param {type:'string'} 
 INSTANCE_URI 
 = 
 "" 
 # @param {type:'string'} 
 DB_NAME 
 = 
 "postgres" 
 #  @param {type:'string'} 
 DB_USER 
 = 
 "postgres" 
 # @param {type:'string'} 
 DB_PASSWORD 
 = 
 "" 
 # @param {type:'string'} 
 

  from 
  
 google.colab 
  
 import 
 auth 
 auth 
 . 
 authenticate_user 
 ( 
 project_id 
 = 
 PROJECT_ID 
 ) 
 

  # @title SQLAlchemy + AlloyDB Connector helpers for creating tables and verifying data 
 import 
  
 sqlalchemy 
 from 
  
 sqlalchemy 
  
 import 
 text 
 # Import text construct explicitly 
 from 
  
 sqlalchemy.exc 
  
 import 
 SQLAlchemyError 
 # Import specific exception type 
 from 
  
 google.cloud.alloydb.connector 
  
 import 
 Connector 
 def 
  
 get_alloydb_engine 
 ( 
 instance_uri 
 : 
 str 
 , 
 user 
 : 
 str 
 , 
 password 
 : 
 str 
 , 
 db 
 : 
 str 
 , 
 ** 
 connect_kwargs 
 ) 
 - 
> sqlalchemy 
 . 
 engine 
 . 
 Engine 
 : 
  
 """Creates a SQLAlchemy engine configured for AlloyDB.""" 
 connector 
 = 
 Connector 
 () 
 connect_kwargs 
 . 
 setdefault 
 ( 
 'ip_type' 
 , 
 'PUBLIC' 
 ) 
 def 
  
 get_conn 
 () 
 - 
> sqlalchemy 
 . 
 engine 
 . 
 base 
 . 
 Connection 
 : 
 conn 
 = 
 connector 
 . 
 connect 
 ( 
 instance_uri 
 , 
 "pg8000" 
 , 
 user 
 = 
 user 
 , 
 password 
 = 
 password 
 , 
 db 
 = 
 db 
 , 
 ** 
 connect_kwargs 
 # Pass additional options like ip_type='PUBLIC' if needed 
 ) 
 return 
 conn 
 # Create the SQLAlchemy engine using the connection function 
 engine 
 = 
 sqlalchemy 
 . 
 create_engine 
 ( 
 "postgresql+pg8000://" 
 , 
 creator 
 = 
 get_conn 
 , 
 ) 
 engine 
 . 
 pool 
 . 
 dispose 
 = 
 lambda 
 : 
 connector 
 . 
 close 
 () 
 return 
 engine 
 def 
  
 setup_alloydb_table_sqlalchemy 
 ( 
 instance_uri 
 : 
 str 
 , 
 database 
 : 
 str 
 , 
 table_name 
 : 
 str 
 , 
 table_schema 
 : 
 str 
 , 
 user 
 : 
 str 
 , 
 password 
 : 
 str 
 , 
 ** 
 connect_kwargs 
 ): 
  
 """Set up AlloyDB table with vector extension and proper schema using SQLAlchemy. 
 Args: 
 instance_uri: AlloyDB instance URI (e.g., projects/.../locations/.../clusters/.../instances/...) 
 database: Database name 
 table_name: Name of the table to create. 
 table_schema: SQL string defining the table columns (e.g., "id SERIAL PRIMARY KEY, embedding VECTOR(768)") 
 user: Database user 
 password: Database password 
 connect_kwargs: Additional keyword arguments passed to connector.connect() (e.g., ip_type="PUBLIC") 
 """ 
 engine 
 = 
 None 
 try 
 : 
 engine 
 = 
 get_alloydb_engine 
 ( 
 instance_uri 
 , 
 user 
 , 
 password 
 , 
 database 
 , 
 ** 
 connect_kwargs 
 ) 
 # Use a connection from the pool 
 with 
 engine 
 . 
 connect 
 () 
 as 
 connection 
 : 
 # Use execution options for autocommit for DDL statements 
 # Alternatively, execute outside an explicit transaction block (begin()) 
 with 
 connection 
 . 
 execution_options 
 ( 
 isolation_level 
 = 
 "AUTOCOMMIT" 
 ): 
 print 
 ( 
 "Connected to AlloyDB successfully via SQLAlchemy!" 
 ) 
 # Create pgvector extension if it doesn't exist 
 print 
 ( 
 "Creating pgvector extension..." 
 ) 
 connection 
 . 
 execute 
 ( 
 text 
 ( 
 "CREATE EXTENSION IF NOT EXISTS vector;" 
 )) 
 # Drop the table if it exists 
 print 
 ( 
 f 
 "Dropping table 
 { 
 table_name 
 } 
 if exists..." 
 ) 
 # Use f-string for table name (generally okay for DDL if source is trusted) 
 connection 
 . 
 execute 
 ( 
 text 
 ( 
 f 
 "DROP TABLE IF EXISTS 
 { 
 table_name 
 } 
 ;" 
 )) 
 # Create the table 
 print 
 ( 
 f 
 "Creating table 
 { 
 table_name 
 } 
 ..." 
 ) 
 # Use f-string for table name and schema (validate input if necessary) 
 create_sql 
 = 
 f 
 """ 
 CREATE TABLE IF NOT EXISTS 
 { 
 table_name 
 } 
 ( 
  
 { 
 table_schema 
 } 
 ); 
 """ 
 connection 
 . 
 execute 
 ( 
 text 
 ( 
 create_sql 
 )) 
 # Optional: Commit if not using autocommit (SQLAlchemy >= 2.0 often commits implicitly) 
 # connection.commit() # Usually not needed with autocommit or implicit commit behavior 
 print 
 ( 
 "Setup completed successfully using SQLAlchemy!" 
 ) 
 except 
 SQLAlchemyError 
 as 
 e 
 : 
 print 
 ( 
 f 
 "An SQLAlchemy error occurred during setup: 
 { 
 e 
 } 
 " 
 ) 
 except 
 Exception 
 as 
 e 
 : 
 print 
 ( 
 f 
 "An unexpected error occurred during setup: 
 { 
 e 
 } 
 " 
 ) 
 finally 
 : 
 if 
 engine 
 : 
 engine 
 . 
 dispose 
 () 
 # Close connection pool and connector 
 def 
  
 test_alloydb_connection_sqlalchemy 
 ( 
 instance_uri 
 : 
 str 
 , 
 database 
 : 
 str 
 , 
 table_name 
 : 
 str 
 , 
 user 
 : 
 str 
 , 
 password 
 : 
 str 
 , 
 ** 
 connect_kwargs 
 ): 
  
 """Test the AlloyDB connection and verify table/extension using SQLAlchemy. 
 Args: 
 instance_uri: AlloyDB instance URI 
 database: Database name 
 table_name: Name of the table to check. 
 user: Database user 
 password: Database password 
 connect_kwargs: Additional keyword arguments passed to connector.connect() 
 """ 
 engine 
 = 
 None 
 try 
 : 
 engine 
 = 
 get_alloydb_engine 
 ( 
 instance_uri 
 , 
 user 
 , 
 password 
 , 
 database 
 , 
 ** 
 connect_kwargs 
 ) 
 with 
 engine 
 . 
 connect 
 () 
 as 
 connection 
 : 
 print 
 ( 
 "Testing connection..." 
 ) 
 # Simple query to confirm connection 
 connection 
 . 
 execute 
 ( 
 text 
 ( 
 "SELECT 1" 
 )) 
 print 
 ( 
 "✓ Connection successful" 
 ) 
 # Check if table exists using information_schema 
 # Use bind parameters (:tname) for safety, even though it's a table name here 
 table_exists_query 
 = 
 text 
 ( 
 """ 
 SELECT EXISTS ( 
 SELECT FROM information_schema.tables 
 WHERE table_schema = 'public' AND table_name = :tname 
 ); 
 """ 
 ) 
 # .scalar() fetches the first column of the first row 
 table_exists 
 = 
 connection 
 . 
 execute 
 ( 
 table_exists_query 
 , 
 { 
 "tname" 
 : 
 table_name 
 }) 
 . 
 scalar 
 () 
 if 
 table_exists 
 : 
 print 
 ( 
 f 
 "✓ ' 
 { 
 table_name 
 } 
 ' table exists" 
 ) 
 # Check if vector extension is installed 
 ext_exists_query 
 = 
 text 
 ( 
 """ 
 SELECT EXISTS ( 
 SELECT FROM pg_extension WHERE extname = 'vector' 
 ); 
 """ 
 ) 
 vector_installed 
 = 
 connection 
 . 
 execute 
 ( 
 ext_exists_query 
 ) 
 . 
 scalar 
 () 
 if 
 vector_installed 
 : 
 print 
 ( 
 "✓ pgvector extension is installed" 
 ) 
 else 
 : 
 print 
 ( 
 "✗ pgvector extension is NOT installed" 
 ) 
 else 
 : 
 print 
 ( 
 f 
 "✗ ' 
 { 
 table_name 
 } 
 ' table does NOT exist" 
 ) 
 except 
 SQLAlchemyError 
 as 
 e 
 : 
 print 
 ( 
 f 
 "Connection test failed (SQLAlchemy error): 
 { 
 e 
 } 
 " 
 ) 
 except 
 Exception 
 as 
 e 
 : 
 print 
 ( 
 f 
 "Connection test failed (Unexpected error): 
 { 
 e 
 } 
 " 
 ) 
 finally 
 : 
 if 
 engine 
 : 
 engine 
 . 
 dispose 
 () 
 def 
  
 verify_embeddings_sqlalchemy 
 ( 
 instance_uri 
 : 
 str 
 , 
 database 
 : 
 str 
 , 
 table_name 
 : 
 str 
 , 
 user 
 : 
 str 
 , 
 password 
 : 
 str 
 , 
 ** 
 connect_kwargs 
 ): 
  
 """Connect to AlloyDB using SQLAlchemy and print all rows from the table.""" 
 engine 
 = 
 None 
 try 
 : 
 engine 
 = 
 get_alloydb_engine 
 ( 
 instance_uri 
 , 
 user 
 , 
 password 
 , 
 database 
 , 
 ** 
 connect_kwargs 
 ) 
 with 
 engine 
 . 
 connect 
 () 
 as 
 connection 
 : 
 # Use f-string for table name in SELECT (ensure table_name is controlled) 
 select_query 
 = 
 text 
 ( 
 f 
 "SELECT * FROM 
 { 
 table_name 
 } 
 ;" 
 ) 
 result 
 = 
 connection 
 . 
 execute 
 ( 
 select_query 
 ) 
 # Get column names from the result keys 
 columns 
 = 
 result 
 . 
 keys 
 () 
 # Fetch all rows as mapping objects (dict-like) 
 rows 
 = 
 result 
 . 
 mappings 
 () 
 . 
 all 
 () 
 print 
 ( 
 f 
 " 
 \n 
 Found 
 { 
 len 
 ( 
 rows 
 ) 
 } 
 products in ' 
 { 
 table_name 
 } 
 ':" 
 ) 
 print 
 ( 
 "-" 
 * 
 80 
 ) 
 if 
 not 
 rows 
 : 
 print 
 ( 
 "Table is empty." 
 ) 
 print 
 ( 
 "-" 
 * 
 80 
 ) 
 else 
 : 
 # Print each row 
 for 
 row 
 in 
 rows 
 : 
 for 
 col 
 in 
 columns 
 : 
 print 
 ( 
 f 
 " 
 { 
 col 
 } 
 : 
 { 
 row 
 [ 
 col 
 ] 
 } 
 " 
 ) 
 print 
 ( 
 "-" 
 * 
 80 
 ) 
 except 
 SQLAlchemyError 
 as 
 e 
 : 
 print 
 ( 
 f 
 "Failed to verify embeddings (SQLAlchemy error): 
 { 
 e 
 } 
 " 
 ) 
 # You might want to check specifically for ProgrammingError if the table doesn't exist 
 # from sqlalchemy.exc import ProgrammingError 
 # except ProgrammingError as pe: 
 #    print(f"Failed to query table '{table_name}'. Does it exist? Error: {pe}") 
 except 
 Exception 
 as 
 e 
 : 
 print 
 ( 
 f 
 "Failed to verify embeddings (Unexpected error): 
 { 
 e 
 } 
 " 
 ) 
 finally 
 : 
 if 
 engine 
 : 
 engine 
 . 
 dispose 
 ()