Set up the PostgreSQL schema for vector search
(args)
| 64 | """ |
| 65 | |
| 66 | async def setup_postgres_schema(args): |
| 67 | """Set up the PostgreSQL schema for vector search""" |
| 68 | print("\n=== PostgreSQL Schema Setup ===\n") |
| 69 | |
| 70 | client = PgVectorClient(args.endpoint) |
| 71 | |
| 72 | # First test the connection |
| 73 | print("Testing PostgreSQL connection...") |
| 74 | try: |
| 75 | connection_info = await client.test_connection() |
| 76 | |
| 77 | if not connection_info.get("success"): |
| 78 | print(f"ERROR: Could not connect to PostgreSQL: {connection_info.get('error')}") |
| 79 | return False |
| 80 | except Exception as e: |
| 81 | print(f"ERROR: Could not connect to PostgreSQL: {e}") |
| 82 | return False |
| 83 | |
| 84 | print(f"Successfully connected to PostgreSQL {connection_info.get('database_version')}") |
| 85 | |
| 86 | # Check if pgvector is installed |
| 87 | if not connection_info.get("pgvector_installed"): |
| 88 | print("WARNING: pgvector extension is not installed in the database") |
| 89 | print("Please install the pgvector extension before continuing") |
| 90 | return False |
| 91 | |
| 92 | # Check if table exists and has correct schema |
| 93 | print(f"\nChecking schema for table '{client.table_name}'...") |
| 94 | schema_info = await client.check_table_schema() |
| 95 | |
| 96 | if schema_info.get("error"): |
| 97 | print(f"ERROR checking table schema: {schema_info.get('error')}") |
| 98 | return False |
| 99 | |
| 100 | if not schema_info.get("table_exists"): |
| 101 | print(f"Table '{client.table_name}' does not exist. Creating it...") |
| 102 | |
| 103 | # Create the table and indexes |
| 104 | async def _create_schema(conn): |
| 105 | async with conn.cursor() as cur: |
| 106 | await cur.execute(CREATE_TABLE_SQL) |
| 107 | await conn.commit() |
| 108 | return True |
| 109 | |
| 110 | try: |
| 111 | await client._execute_with_retry(_create_schema) |
| 112 | print(f"Successfully created table '{client.table_name}' and indexes") |
| 113 | except Exception as e: |
| 114 | print(f"ERROR creating schema: {e}") |
| 115 | await client.close() # Make sure to close the connection on error |
| 116 | return False |
| 117 | else: |
| 118 | print(f"Table '{client.table_name}' already exists") |
| 119 | |
| 120 | # Check for any schema issues |
| 121 | if schema_info.get("needs_corrections"): |
| 122 | print("\nThe following schema issues were detected:") |
| 123 | for issue in schema_info.get("needs_corrections"): |
no test coverage detected