💨

TiDB for AI : pytidb AI に特化したTiDB用 Python SDK その２

2025/08/02に公開

TiDB

tech

今日は引き続きpytidbを触っていきます。

前回の記事では埋め込まれたベクトルデータをテーブルに保存して検索を行うサンプルを作成しました。pytidbはVector Search用に見えるかもしれませんが、もちろん通常のSQL発行にも対応しています。今日はそれを触っていきます。

さっそくやってみる

通常のSQL発行はこちらに公式ドキュメントがあります。

raw(生）と記載されているように、通常のSQL発行は関数が定義されているわけではなく、生SQLをそのまま実行する形態をとっています。基本pytidbはAIと連携することが前提のライブラリであることがわかります。

以下の内容でrawoperation.pyを作成します。

from pytidb import TiDBClient

# データベース接続
client = TiDBClient.connect(
    host="gateway01.us-west-2.prod.aws.tidbcloud.com",
    port=4000,
    username="237VxcjMhxqE85K.root",
    password="mnRExxxx",
    database="test",
)

# 接続確認
result = client.execute("SELECT 1 as test")
print("接続成功")
print(f"結果: {result}")

# テーブル作成
print("\nテーブル作成...")
create_table_sql = """
CREATE TABLE IF NOT EXISTS documents (
    id INT AUTO_INCREMENT PRIMARY KEY,
    title VARCHAR(255),
    content TEXT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
"""
try:
    result = client.execute(create_table_sql)
    print("✅ テーブル作成成功")
    print(f"結果: {result}")
except Exception as e:
    print(f"❌ テーブル作成エラー: {e}")

# データ挿入
print("\nデータ挿入...")
insert_sql = """
INSERT INTO documents (title, content) VALUES 
('Python Programming', 'Python is a high-level programming language.'),
('Database Systems', 'A database is an organized collection of data.'),
('Machine Learning', 'Machine learning is a subset of artificial intelligence.'),
('Web Development', 'Web development involves building websites and web applications.')
"""
try:
    result = client.execute(insert_sql)
    print("✅ データ挿入成功")
    print(f"結果: {result}")
except Exception as e:
    print(f"❌ データ挿入エラー: {e}")from pytidb import TiDBClient

python rawoperation.py
接続成功
結果: rowcount=1 success=True message=None

テーブル作成...
✅ テーブル作成成功
結果: rowcount=0 success=True message=None

データ挿入...
✅ データ挿入成功
結果: rowcount=4 success=True message=None

ORM的にSQL Injectionを防ぐモードも用意されています。
rawoperation.pyを以下に置き換えます。

rawoperation.py

from pytidb import TiDBClient

# データベース接続
client = TiDBClient.connect(
    host="gateway01.us-west-2.prod.aws.tidbcloud.com",
    port=4000,
    username="237VxcjMhxqE85K.root",
    password="mnREoh8Egrtg8fDr",
    database="test",
)

# 接続確認
result = client.execute("SELECT 1 as test")
print("接続成功")
print(f"結果: {result}")

# テーブル作成
print("\nテーブル作成...")
create_table_sql = """
CREATE TABLE IF NOT EXISTS documents (
    id INT AUTO_INCREMENT PRIMARY KEY,
    title VARCHAR(255),
    content TEXT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
"""
try:
    result = client.execute(create_table_sql)
    print("✅ テーブル作成成功")
    print(f"結果: {result}")
except Exception as e:
    print(f"❌ テーブル作成エラー: {e}")

# データ挿入
print("\nデータ挿入...")
try:
    # 1件目
    result1 = client.execute(
        "INSERT INTO documents(title, content) VALUES (:title, :content)",
        {
            "title": "Python Programming",
            "content": "Python is a high-level programming language.",
        },
    )
    print("✅ データ挿入1成功")
    print(f"結果: {result1}")
    
    # 2件目
    result2 = client.execute(
        "INSERT INTO documents(title, content) VALUES (:title, :content)",
        {
            "title": "Database Systems",
            "content": "A database is an organized collection of data.",
        },
    )
    print("✅ データ挿入2成功")
    print(f"結果: {result2}")
    
    # 3件目
    result3 = client.execute(
        "INSERT INTO documents(title, content) VALUES (:title, :content)",
        {
            "title": "Machine Learning",
            "content": "Machine learning is a subset of artificial intelligence.",
        },
    )
    print("✅ データ挿入3成功")
    print(f"結果: {result3}")
    
    # 4件目
    result4 = client.execute(
        "INSERT INTO documents(title, content) VALUES (:title, :content)",
        {
            "title": "Web Development",
            "content": "Web development involves building websites and web applications.",
        },
    )
    print("✅ データ挿入4成功")
    print(f"結果: {result4}")
    
except Exception as e:
    print(f"❌ データ挿入エラー: {e}")

 python rawoperation.py
接続成功
結果: rowcount=1 success=True message=None

テーブル作成...
✅ テーブル作成成功
結果: rowcount=0 success=True message=None

データ挿入...
✅ データ挿入1成功
結果: rowcount=1 success=True message=None
✅ データ挿入2成功
結果: rowcount=1 success=True message=None
✅ データ挿入3成功
結果: rowcount=1 success=True message=None
✅ データ挿入4成功
結果: rowcount=1 success=True message=None

コンソールのSQL Editorでテーブルの中身を見てみるとちゃんと入っていることがわかります。

use test;
select * from documents;

result1 = client.execute(
        "INSERT INTO documents(title, content) VALUES (:title, :content)",
        {
            "title": "Python Programming",
            "content": "Python is a high-level programming language.",
        },
    )

の部分は本来SQLですと以下を意味しています。

INSERT INTO documents(title, content) VALUES ('Python Programming', 'Python is a high-level programming language.');

これはSQL Injection用攻撃耐性を持ちます。仮にtitleが変数扱いで外部インプットを受けるとします。

title = "Test'); DROP TABLE documents; --"

と入力された場合生SQLであれば

INSERT INTO documents(title, content) VALUES ('Test'); DROP TABLE documents; --', 'Python is a high-level programming language.');

となるためテーブルが丸ごと削除されてしまいます。
プレースホルダ（:title, :content）を使うことで、ライブラリ側が自動で文字列を安全にエスケープ・バインドします。

client.execute と client.query

Select SQLの実行を表示したり受け取って取り出すのはexecuteではなくqueryを使います。
executeだと結果: rowcount=8 success=True message=Noneという値が戻ります。

client.execute() は INSERT, UPDATE, DELETE, CREATE TABLE などの作用のある操作に使います。client.query() は SELECT など結果を返す操作に使います。

rawselect.pyとして以下を作成します。

rawselect.py

from pytidb import TiDBClient

# データベース接続
client = TiDBClient.connect(
    host="gateway01.us-west-2.prod.aws.tidbcloud.com",
    port=4000,
    username="237VxcjMhxqE85K.root",
    password="mnRExxxxx",
    database="test",
)

# データ取得
print("\nデータ取得...")
try:
    result = client.query("SELECT * FROM documents").to_rows()
    print("✅ データ取得成功")
    print(f"結果: {result}")
except Exception as e:
    print(f"❌ データ取得エラー: {e}")

python rawselect.py

データ取得...
✅ データ取得成功
結果: [(1, 'Python Programming', 'Python is a high-level programming language.', datetime.datetime(2025, 8, 2, 3, 32, 34)), (2, 'Database Systems', 'A database is an organized collection of data.', datetime.datetime(2025, 8, 2, 3, 32, 34)), (3, 'Machine Learning', 'Machine learning is a subset of artificial intelligence.', datetime.datetime(2025, 8, 2, 3, 32, 34)), (4, 'Web Development', 'Web development involves building websites and web applications.', datetime.datetime(2025, 8, 2, 3, 32, 34)), (5, 'Python Programming', 'Python is a high-level programming language.', datetime.datetime(2025, 8, 2, 3, 41, 11)), (6, 'Database Systems', 'A database is an organized collection of data.', datetime.datetime(2025, 8, 2, 3, 41, 11)), (7, 'Machine Learning', 'Machine learning is a subset of artificial intelligence.', datetime.datetime(2025, 8, 2, 3, 41, 12)), (8, 'Web Development', 'Web development involves building websites and web applications.', datetime.datetime(2025, 8, 2, 3, 41, 12))]

その他以下のオペレーションが用意されています。
to_pydantic() : Pydantic用データモデル型で出力
to_list() ：辞書型で出力
to_pandas() ： pandas用データモデル型で出力
scalar()　:　一番最初のレコードの１行目を出力

さっそくやってみる

client.execute と client.query

Discussion