286 lines
8.4 KiB
Python
286 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
百度VDB向量数据库连接测试脚本
|
|
测试数据库连接、基本操作和可用性
|
|
"""
|
|
|
|
import pymochow
|
|
from pymochow.configuration import Configuration
|
|
from pymochow.auth.bce_credentials import BceCredentials
|
|
from pymochow.model.schema import Schema, Field, VectorIndex, HNSWParams
|
|
from pymochow.model.enum import FieldType, IndexType, MetricType
|
|
from pymochow.model.table import Partition, Row
|
|
import traceback
|
|
import time
|
|
|
|
# 数据库连接配置
|
|
ACCOUNT = 'root'
|
|
API_KEY = 'vdb$yjr9ln3n0td' # 你提供的密码作为API密钥
|
|
ENDPOINT = 'http://180.76.96.191:5287' # 使用标准端口5287
|
|
|
|
def test_connection():
|
|
"""测试数据库连接"""
|
|
print("=" * 50)
|
|
print("测试1: 数据库连接")
|
|
print("=" * 50)
|
|
|
|
try:
|
|
# 创建配置和客户端
|
|
config = Configuration(
|
|
credentials=BceCredentials(ACCOUNT, API_KEY),
|
|
endpoint=ENDPOINT
|
|
)
|
|
client = pymochow.MochowClient(config)
|
|
|
|
print(f"✓ 成功创建客户端连接")
|
|
print(f" - 账户: {ACCOUNT}")
|
|
print(f" - 端点: {ENDPOINT}")
|
|
|
|
return client
|
|
|
|
except Exception as e:
|
|
print(f"✗ 连接失败: {str(e)}")
|
|
print(f"详细错误: {traceback.format_exc()}")
|
|
return None
|
|
|
|
def test_list_databases(client):
|
|
"""测试数据库列表查询"""
|
|
print("\n" + "=" * 50)
|
|
print("测试2: 查询数据库列表")
|
|
print("=" * 50)
|
|
|
|
try:
|
|
# 查询数据库列表
|
|
db_list = client.list_databases()
|
|
print(f"✓ 成功查询数据库列表")
|
|
print(f" - 数据库数量: {len(db_list)}")
|
|
|
|
if db_list:
|
|
print(" - 数据库列表:")
|
|
for i, db in enumerate(db_list, 1):
|
|
print(f" {i}. {db.database_name}")
|
|
else:
|
|
print(" - 当前没有数据库")
|
|
|
|
return db_list
|
|
|
|
except Exception as e:
|
|
print(f"✗ 查询数据库列表失败: {str(e)}")
|
|
print(f"详细错误: {traceback.format_exc()}")
|
|
return None
|
|
|
|
def test_create_database(client):
|
|
"""测试创建数据库"""
|
|
print("\n" + "=" * 50)
|
|
print("测试3: 创建测试数据库")
|
|
print("=" * 50)
|
|
|
|
test_db_name = "test_db_" + str(int(time.time()))
|
|
|
|
try:
|
|
# 创建测试数据库
|
|
db = client.create_database(test_db_name)
|
|
print(f"✓ 成功创建数据库: {test_db_name}")
|
|
|
|
return db, test_db_name
|
|
|
|
except Exception as e:
|
|
print(f"✗ 创建数据库失败: {str(e)}")
|
|
print(f"详细错误: {traceback.format_exc()}")
|
|
return None, test_db_name
|
|
|
|
def test_create_table(client, db, db_name):
|
|
"""测试创建表"""
|
|
print("\n" + "=" * 50)
|
|
print("测试4: 创建测试表")
|
|
print("=" * 50)
|
|
|
|
table_name = "test_table"
|
|
|
|
try:
|
|
# 定义表字段
|
|
fields = []
|
|
fields.append(Field("id", FieldType.STRING, primary_key=True,
|
|
partition_key=True, auto_increment=False, not_null=True))
|
|
fields.append(Field("text", FieldType.STRING, not_null=True))
|
|
fields.append(Field("vector", FieldType.FLOAT_VECTOR, not_null=True, dimension=3))
|
|
|
|
# 定义索引
|
|
indexes = []
|
|
indexes.append(
|
|
VectorIndex(
|
|
index_name="vector_idx",
|
|
index_type=IndexType.HNSW,
|
|
field="vector",
|
|
metric_type=MetricType.L2,
|
|
params=HNSWParams(m=16, efconstruction=100),
|
|
auto_build=True
|
|
)
|
|
)
|
|
|
|
# 创建表
|
|
table = db.create_table(
|
|
table_name=table_name,
|
|
replication=2, # 最小副本数为2
|
|
partition=Partition(partition_num=1), # 单分区测试
|
|
schema=Schema(fields=fields, indexes=indexes)
|
|
)
|
|
|
|
print(f"✓ 成功创建表: {table_name}")
|
|
print(f" - 字段数量: {len(fields)}")
|
|
print(f" - 索引数量: {len(indexes)}")
|
|
|
|
return table, table_name
|
|
|
|
except Exception as e:
|
|
print(f"✗ 创建表失败: {str(e)}")
|
|
print(f"详细错误: {traceback.format_exc()}")
|
|
return None, table_name
|
|
|
|
def test_insert_data(table):
|
|
"""测试插入数据"""
|
|
print("\n" + "=" * 50)
|
|
print("测试5: 插入测试数据")
|
|
print("=" * 50)
|
|
|
|
try:
|
|
# 准备测试数据
|
|
rows = [
|
|
Row(id='001', text='测试文本1', vector=[0.1, 0.2, 0.3]),
|
|
Row(id='002', text='测试文本2', vector=[0.4, 0.5, 0.6]),
|
|
Row(id='003', text='测试文本3', vector=[0.7, 0.8, 0.9])
|
|
]
|
|
|
|
# 插入数据
|
|
table.upsert(rows)
|
|
print(f"✓ 成功插入 {len(rows)} 条测试数据")
|
|
|
|
# 等待一下让数据生效
|
|
time.sleep(2)
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"✗ 插入数据失败: {str(e)}")
|
|
print(f"详细错误: {traceback.format_exc()}")
|
|
return False
|
|
|
|
def test_query_data(table):
|
|
"""测试查询数据"""
|
|
print("\n" + "=" * 50)
|
|
print("测试6: 查询测试数据")
|
|
print("=" * 50)
|
|
|
|
try:
|
|
# 标量查询
|
|
primary_key = {'id': '001'}
|
|
result = table.query(primary_key=primary_key, retrieve_vector=True)
|
|
|
|
if result:
|
|
print(f"✓ 成功查询数据")
|
|
print(f" - ID: {result.get('id')}")
|
|
print(f" - 文本: {result.get('text')}")
|
|
print(f" - 向量: {result.get('vector')}")
|
|
else:
|
|
print("✗ 查询结果为空")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"✗ 查询数据失败: {str(e)}")
|
|
print(f"详细错误: {traceback.format_exc()}")
|
|
return False
|
|
|
|
def test_cleanup(client, db_name, table_name):
|
|
"""清理测试数据"""
|
|
print("\n" + "=" * 50)
|
|
print("测试7: 清理测试数据")
|
|
print("=" * 50)
|
|
|
|
try:
|
|
# 获取数据库对象
|
|
db = client.database(db_name)
|
|
|
|
# 删除表
|
|
try:
|
|
db.drop_table(table_name)
|
|
print(f"✓ 成功删除表: {table_name}")
|
|
except Exception as e:
|
|
print(f"⚠ 删除表失败: {str(e)}")
|
|
|
|
# 删除数据库
|
|
try:
|
|
client.drop_database(db_name)
|
|
print(f"✓ 成功删除数据库: {db_name}")
|
|
except Exception as e:
|
|
print(f"⚠ 删除数据库失败: {str(e)}")
|
|
|
|
except Exception as e:
|
|
print(f"⚠ 清理过程出现错误: {str(e)}")
|
|
|
|
def main():
|
|
"""主测试函数"""
|
|
print("百度VDB向量数据库可用性测试")
|
|
print("测试配置:")
|
|
print(f" - 用户名: {ACCOUNT}")
|
|
print(f" - 服务器: {ENDPOINT}")
|
|
print(f" - 测试时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
|
|
|
client = None
|
|
db = None
|
|
db_name = None
|
|
table = None
|
|
table_name = None
|
|
|
|
try:
|
|
# 测试1: 连接
|
|
client = test_connection()
|
|
if not client:
|
|
print("\n❌ 数据库连接失败，无法继续测试")
|
|
return
|
|
|
|
# 测试2: 查询数据库列表
|
|
db_list = test_list_databases(client)
|
|
|
|
# 测试3: 创建数据库
|
|
db, db_name = test_create_database(client)
|
|
if not db:
|
|
print("\n❌ 无法创建数据库，跳过后续测试")
|
|
return
|
|
|
|
# 测试4: 创建表
|
|
table, table_name = test_create_table(client, db, db_name)
|
|
if not table:
|
|
print("\n❌ 无法创建表，跳过后续测试")
|
|
return
|
|
|
|
# 测试5: 插入数据
|
|
if test_insert_data(table):
|
|
# 测试6: 查询数据
|
|
test_query_data(table)
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ 测试过程中发生未预期错误: {str(e)}")
|
|
print(f"详细错误: {traceback.format_exc()}")
|
|
|
|
finally:
|
|
# 测试7: 清理
|
|
if client and db_name:
|
|
test_cleanup(client, db_name, table_name)
|
|
|
|
# 关闭连接
|
|
if client:
|
|
try:
|
|
client.close()
|
|
print(f"\n✓ 已关闭数据库连接")
|
|
except:
|
|
pass
|
|
|
|
print("\n" + "=" * 50)
|
|
print("测试完成")
|
|
print("=" * 50)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|