sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ac8f31065839703b4e70__mypyc.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/__init__.py +140 -0
- sqlspec/__main__.py +12 -0
- sqlspec/__metadata__.py +14 -0
- sqlspec/_serialization.py +315 -0
- sqlspec/_typing.py +700 -0
- sqlspec/adapters/__init__.py +0 -0
- sqlspec/adapters/adbc/__init__.py +5 -0
- sqlspec/adapters/adbc/_typing.py +82 -0
- sqlspec/adapters/adbc/adk/__init__.py +5 -0
- sqlspec/adapters/adbc/adk/store.py +1273 -0
- sqlspec/adapters/adbc/config.py +295 -0
- sqlspec/adapters/adbc/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/adbc/core.py +735 -0
- sqlspec/adapters/adbc/data_dictionary.py +334 -0
- sqlspec/adapters/adbc/driver.py +529 -0
- sqlspec/adapters/adbc/events/__init__.py +5 -0
- sqlspec/adapters/adbc/events/store.py +285 -0
- sqlspec/adapters/adbc/litestar/__init__.py +5 -0
- sqlspec/adapters/adbc/litestar/store.py +502 -0
- sqlspec/adapters/adbc/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/adbc/type_converter.py +140 -0
- sqlspec/adapters/aiosqlite/__init__.py +25 -0
- sqlspec/adapters/aiosqlite/_typing.py +82 -0
- sqlspec/adapters/aiosqlite/adk/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/adk/store.py +818 -0
- sqlspec/adapters/aiosqlite/config.py +334 -0
- sqlspec/adapters/aiosqlite/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/aiosqlite/core.py +315 -0
- sqlspec/adapters/aiosqlite/data_dictionary.py +208 -0
- sqlspec/adapters/aiosqlite/driver.py +313 -0
- sqlspec/adapters/aiosqlite/events/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/events/store.py +20 -0
- sqlspec/adapters/aiosqlite/litestar/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/litestar/store.py +279 -0
- sqlspec/adapters/aiosqlite/pool.py +533 -0
- sqlspec/adapters/asyncmy/__init__.py +21 -0
- sqlspec/adapters/asyncmy/_typing.py +87 -0
- sqlspec/adapters/asyncmy/adk/__init__.py +5 -0
- sqlspec/adapters/asyncmy/adk/store.py +703 -0
- sqlspec/adapters/asyncmy/config.py +302 -0
- sqlspec/adapters/asyncmy/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/asyncmy/core.py +360 -0
- sqlspec/adapters/asyncmy/data_dictionary.py +124 -0
- sqlspec/adapters/asyncmy/driver.py +383 -0
- sqlspec/adapters/asyncmy/events/__init__.py +5 -0
- sqlspec/adapters/asyncmy/events/store.py +104 -0
- sqlspec/adapters/asyncmy/litestar/__init__.py +5 -0
- sqlspec/adapters/asyncmy/litestar/store.py +296 -0
- sqlspec/adapters/asyncpg/__init__.py +19 -0
- sqlspec/adapters/asyncpg/_typing.py +88 -0
- sqlspec/adapters/asyncpg/adk/__init__.py +5 -0
- sqlspec/adapters/asyncpg/adk/store.py +748 -0
- sqlspec/adapters/asyncpg/config.py +569 -0
- sqlspec/adapters/asyncpg/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/asyncpg/core.py +367 -0
- sqlspec/adapters/asyncpg/data_dictionary.py +162 -0
- sqlspec/adapters/asyncpg/driver.py +487 -0
- sqlspec/adapters/asyncpg/events/__init__.py +6 -0
- sqlspec/adapters/asyncpg/events/backend.py +286 -0
- sqlspec/adapters/asyncpg/events/store.py +40 -0
- sqlspec/adapters/asyncpg/litestar/__init__.py +5 -0
- sqlspec/adapters/asyncpg/litestar/store.py +251 -0
- sqlspec/adapters/bigquery/__init__.py +14 -0
- sqlspec/adapters/bigquery/_typing.py +86 -0
- sqlspec/adapters/bigquery/adk/__init__.py +5 -0
- sqlspec/adapters/bigquery/adk/store.py +827 -0
- sqlspec/adapters/bigquery/config.py +353 -0
- sqlspec/adapters/bigquery/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/bigquery/core.py +715 -0
- sqlspec/adapters/bigquery/data_dictionary.py +128 -0
- sqlspec/adapters/bigquery/driver.py +548 -0
- sqlspec/adapters/bigquery/events/__init__.py +5 -0
- sqlspec/adapters/bigquery/events/store.py +139 -0
- sqlspec/adapters/bigquery/litestar/__init__.py +5 -0
- sqlspec/adapters/bigquery/litestar/store.py +325 -0
- sqlspec/adapters/bigquery/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/bigquery/type_converter.py +107 -0
- sqlspec/adapters/cockroach_asyncpg/__init__.py +24 -0
- sqlspec/adapters/cockroach_asyncpg/_typing.py +72 -0
- sqlspec/adapters/cockroach_asyncpg/adk/__init__.py +3 -0
- sqlspec/adapters/cockroach_asyncpg/adk/store.py +410 -0
- sqlspec/adapters/cockroach_asyncpg/config.py +238 -0
- sqlspec/adapters/cockroach_asyncpg/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/cockroach_asyncpg/core.py +55 -0
- sqlspec/adapters/cockroach_asyncpg/data_dictionary.py +107 -0
- sqlspec/adapters/cockroach_asyncpg/driver.py +144 -0
- sqlspec/adapters/cockroach_asyncpg/events/__init__.py +3 -0
- sqlspec/adapters/cockroach_asyncpg/events/store.py +20 -0
- sqlspec/adapters/cockroach_asyncpg/litestar/__init__.py +3 -0
- sqlspec/adapters/cockroach_asyncpg/litestar/store.py +142 -0
- sqlspec/adapters/cockroach_psycopg/__init__.py +38 -0
- sqlspec/adapters/cockroach_psycopg/_typing.py +129 -0
- sqlspec/adapters/cockroach_psycopg/adk/__init__.py +13 -0
- sqlspec/adapters/cockroach_psycopg/adk/store.py +868 -0
- sqlspec/adapters/cockroach_psycopg/config.py +484 -0
- sqlspec/adapters/cockroach_psycopg/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/cockroach_psycopg/core.py +63 -0
- sqlspec/adapters/cockroach_psycopg/data_dictionary.py +215 -0
- sqlspec/adapters/cockroach_psycopg/driver.py +284 -0
- sqlspec/adapters/cockroach_psycopg/events/__init__.py +6 -0
- sqlspec/adapters/cockroach_psycopg/events/store.py +34 -0
- sqlspec/adapters/cockroach_psycopg/litestar/__init__.py +3 -0
- sqlspec/adapters/cockroach_psycopg/litestar/store.py +325 -0
- sqlspec/adapters/duckdb/__init__.py +25 -0
- sqlspec/adapters/duckdb/_typing.py +81 -0
- sqlspec/adapters/duckdb/adk/__init__.py +14 -0
- sqlspec/adapters/duckdb/adk/store.py +850 -0
- sqlspec/adapters/duckdb/config.py +463 -0
- sqlspec/adapters/duckdb/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/duckdb/core.py +257 -0
- sqlspec/adapters/duckdb/data_dictionary.py +140 -0
- sqlspec/adapters/duckdb/driver.py +430 -0
- sqlspec/adapters/duckdb/events/__init__.py +5 -0
- sqlspec/adapters/duckdb/events/store.py +57 -0
- sqlspec/adapters/duckdb/litestar/__init__.py +5 -0
- sqlspec/adapters/duckdb/litestar/store.py +330 -0
- sqlspec/adapters/duckdb/pool.py +293 -0
- sqlspec/adapters/duckdb/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/duckdb/type_converter.py +118 -0
- sqlspec/adapters/mock/__init__.py +72 -0
- sqlspec/adapters/mock/_typing.py +147 -0
- sqlspec/adapters/mock/config.py +483 -0
- sqlspec/adapters/mock/core.py +319 -0
- sqlspec/adapters/mock/data_dictionary.py +366 -0
- sqlspec/adapters/mock/driver.py +721 -0
- sqlspec/adapters/mysqlconnector/__init__.py +36 -0
- sqlspec/adapters/mysqlconnector/_typing.py +141 -0
- sqlspec/adapters/mysqlconnector/adk/__init__.py +15 -0
- sqlspec/adapters/mysqlconnector/adk/store.py +1060 -0
- sqlspec/adapters/mysqlconnector/config.py +394 -0
- sqlspec/adapters/mysqlconnector/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/mysqlconnector/core.py +303 -0
- sqlspec/adapters/mysqlconnector/data_dictionary.py +235 -0
- sqlspec/adapters/mysqlconnector/driver.py +483 -0
- sqlspec/adapters/mysqlconnector/events/__init__.py +8 -0
- sqlspec/adapters/mysqlconnector/events/store.py +98 -0
- sqlspec/adapters/mysqlconnector/litestar/__init__.py +5 -0
- sqlspec/adapters/mysqlconnector/litestar/store.py +426 -0
- sqlspec/adapters/oracledb/__init__.py +60 -0
- sqlspec/adapters/oracledb/_numpy_handlers.py +141 -0
- sqlspec/adapters/oracledb/_typing.py +182 -0
- sqlspec/adapters/oracledb/_uuid_handlers.py +166 -0
- sqlspec/adapters/oracledb/adk/__init__.py +10 -0
- sqlspec/adapters/oracledb/adk/store.py +2369 -0
- sqlspec/adapters/oracledb/config.py +550 -0
- sqlspec/adapters/oracledb/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/oracledb/core.py +543 -0
- sqlspec/adapters/oracledb/data_dictionary.py +536 -0
- sqlspec/adapters/oracledb/driver.py +1229 -0
- sqlspec/adapters/oracledb/events/__init__.py +16 -0
- sqlspec/adapters/oracledb/events/backend.py +347 -0
- sqlspec/adapters/oracledb/events/store.py +420 -0
- sqlspec/adapters/oracledb/litestar/__init__.py +5 -0
- sqlspec/adapters/oracledb/litestar/store.py +781 -0
- sqlspec/adapters/oracledb/migrations.py +535 -0
- sqlspec/adapters/oracledb/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/oracledb/type_converter.py +211 -0
- sqlspec/adapters/psqlpy/__init__.py +17 -0
- sqlspec/adapters/psqlpy/_typing.py +79 -0
- sqlspec/adapters/psqlpy/adk/__init__.py +5 -0
- sqlspec/adapters/psqlpy/adk/store.py +766 -0
- sqlspec/adapters/psqlpy/config.py +304 -0
- sqlspec/adapters/psqlpy/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/psqlpy/core.py +480 -0
- sqlspec/adapters/psqlpy/data_dictionary.py +126 -0
- sqlspec/adapters/psqlpy/driver.py +438 -0
- sqlspec/adapters/psqlpy/events/__init__.py +6 -0
- sqlspec/adapters/psqlpy/events/backend.py +310 -0
- sqlspec/adapters/psqlpy/events/store.py +20 -0
- sqlspec/adapters/psqlpy/litestar/__init__.py +5 -0
- sqlspec/adapters/psqlpy/litestar/store.py +270 -0
- sqlspec/adapters/psqlpy/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/psqlpy/type_converter.py +113 -0
- sqlspec/adapters/psycopg/__init__.py +32 -0
- sqlspec/adapters/psycopg/_typing.py +164 -0
- sqlspec/adapters/psycopg/adk/__init__.py +10 -0
- sqlspec/adapters/psycopg/adk/store.py +1387 -0
- sqlspec/adapters/psycopg/config.py +576 -0
- sqlspec/adapters/psycopg/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/psycopg/core.py +450 -0
- sqlspec/adapters/psycopg/data_dictionary.py +289 -0
- sqlspec/adapters/psycopg/driver.py +975 -0
- sqlspec/adapters/psycopg/events/__init__.py +20 -0
- sqlspec/adapters/psycopg/events/backend.py +458 -0
- sqlspec/adapters/psycopg/events/store.py +42 -0
- sqlspec/adapters/psycopg/litestar/__init__.py +5 -0
- sqlspec/adapters/psycopg/litestar/store.py +552 -0
- sqlspec/adapters/psycopg/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/psycopg/type_converter.py +93 -0
- sqlspec/adapters/pymysql/__init__.py +21 -0
- sqlspec/adapters/pymysql/_typing.py +71 -0
- sqlspec/adapters/pymysql/adk/__init__.py +5 -0
- sqlspec/adapters/pymysql/adk/store.py +540 -0
- sqlspec/adapters/pymysql/config.py +195 -0
- sqlspec/adapters/pymysql/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/pymysql/core.py +299 -0
- sqlspec/adapters/pymysql/data_dictionary.py +122 -0
- sqlspec/adapters/pymysql/driver.py +259 -0
- sqlspec/adapters/pymysql/events/__init__.py +5 -0
- sqlspec/adapters/pymysql/events/store.py +50 -0
- sqlspec/adapters/pymysql/litestar/__init__.py +5 -0
- sqlspec/adapters/pymysql/litestar/store.py +232 -0
- sqlspec/adapters/pymysql/pool.py +137 -0
- sqlspec/adapters/spanner/__init__.py +40 -0
- sqlspec/adapters/spanner/_typing.py +86 -0
- sqlspec/adapters/spanner/adk/__init__.py +5 -0
- sqlspec/adapters/spanner/adk/store.py +732 -0
- sqlspec/adapters/spanner/config.py +352 -0
- sqlspec/adapters/spanner/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/spanner/core.py +188 -0
- sqlspec/adapters/spanner/data_dictionary.py +120 -0
- sqlspec/adapters/spanner/dialect/__init__.py +6 -0
- sqlspec/adapters/spanner/dialect/_spangres.py +57 -0
- sqlspec/adapters/spanner/dialect/_spanner.py +130 -0
- sqlspec/adapters/spanner/driver.py +373 -0
- sqlspec/adapters/spanner/events/__init__.py +5 -0
- sqlspec/adapters/spanner/events/store.py +187 -0
- sqlspec/adapters/spanner/litestar/__init__.py +5 -0
- sqlspec/adapters/spanner/litestar/store.py +291 -0
- sqlspec/adapters/spanner/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/spanner/type_converter.py +331 -0
- sqlspec/adapters/sqlite/__init__.py +19 -0
- sqlspec/adapters/sqlite/_typing.py +80 -0
- sqlspec/adapters/sqlite/adk/__init__.py +5 -0
- sqlspec/adapters/sqlite/adk/store.py +958 -0
- sqlspec/adapters/sqlite/config.py +280 -0
- sqlspec/adapters/sqlite/core.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/sqlite/core.py +312 -0
- sqlspec/adapters/sqlite/data_dictionary.py +202 -0
- sqlspec/adapters/sqlite/driver.py +359 -0
- sqlspec/adapters/sqlite/events/__init__.py +5 -0
- sqlspec/adapters/sqlite/events/store.py +20 -0
- sqlspec/adapters/sqlite/litestar/__init__.py +5 -0
- sqlspec/adapters/sqlite/litestar/store.py +316 -0
- sqlspec/adapters/sqlite/pool.py +198 -0
- sqlspec/adapters/sqlite/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/adapters/sqlite/type_converter.py +114 -0
- sqlspec/base.py +747 -0
- sqlspec/builder/__init__.py +179 -0
- sqlspec/builder/_base.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_base.py +1022 -0
- sqlspec/builder/_column.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_column.py +521 -0
- sqlspec/builder/_ddl.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_ddl.py +1642 -0
- sqlspec/builder/_delete.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_delete.py +95 -0
- sqlspec/builder/_dml.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_dml.py +365 -0
- sqlspec/builder/_explain.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_explain.py +579 -0
- sqlspec/builder/_expression_wrappers.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_expression_wrappers.py +46 -0
- sqlspec/builder/_factory.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_factory.py +1697 -0
- sqlspec/builder/_insert.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_insert.py +328 -0
- sqlspec/builder/_join.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_join.py +499 -0
- sqlspec/builder/_merge.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_merge.py +821 -0
- sqlspec/builder/_parsing_utils.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_parsing_utils.py +297 -0
- sqlspec/builder/_select.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_select.py +1660 -0
- sqlspec/builder/_temporal.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_temporal.py +139 -0
- sqlspec/builder/_update.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/builder/_update.py +173 -0
- sqlspec/builder/_vector_expressions.py +267 -0
- sqlspec/cli.py +911 -0
- sqlspec/config.py +1755 -0
- sqlspec/core/__init__.py +374 -0
- sqlspec/core/_correlation.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/_correlation.py +176 -0
- sqlspec/core/cache.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/cache.py +1069 -0
- sqlspec/core/compiler.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/compiler.py +954 -0
- sqlspec/core/explain.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/explain.py +275 -0
- sqlspec/core/filters.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/filters.py +952 -0
- sqlspec/core/hashing.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/hashing.py +262 -0
- sqlspec/core/metrics.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/metrics.py +83 -0
- sqlspec/core/parameters/__init__.py +71 -0
- sqlspec/core/parameters/_alignment.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/parameters/_alignment.py +270 -0
- sqlspec/core/parameters/_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/parameters/_converter.py +543 -0
- sqlspec/core/parameters/_processor.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/parameters/_processor.py +505 -0
- sqlspec/core/parameters/_registry.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/parameters/_registry.py +206 -0
- sqlspec/core/parameters/_transformers.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/parameters/_transformers.py +292 -0
- sqlspec/core/parameters/_types.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/parameters/_types.py +499 -0
- sqlspec/core/parameters/_validator.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/parameters/_validator.py +180 -0
- sqlspec/core/pipeline.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/pipeline.py +319 -0
- sqlspec/core/query_modifiers.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/query_modifiers.py +437 -0
- sqlspec/core/result/__init__.py +23 -0
- sqlspec/core/result/_base.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/result/_base.py +1121 -0
- sqlspec/core/result/_io.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/result/_io.py +28 -0
- sqlspec/core/splitter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/splitter.py +966 -0
- sqlspec/core/stack.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/stack.py +163 -0
- sqlspec/core/statement.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/statement.py +1503 -0
- sqlspec/core/type_converter.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/core/type_converter.py +339 -0
- sqlspec/data_dictionary/__init__.py +22 -0
- sqlspec/data_dictionary/_loader.py +123 -0
- sqlspec/data_dictionary/_registry.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/_registry.py +74 -0
- sqlspec/data_dictionary/_types.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/_types.py +121 -0
- sqlspec/data_dictionary/dialects/__init__.py +21 -0
- sqlspec/data_dictionary/dialects/bigquery.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/bigquery.py +49 -0
- sqlspec/data_dictionary/dialects/cockroachdb.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/cockroachdb.py +43 -0
- sqlspec/data_dictionary/dialects/duckdb.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/duckdb.py +47 -0
- sqlspec/data_dictionary/dialects/mysql.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/mysql.py +42 -0
- sqlspec/data_dictionary/dialects/oracle.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/oracle.py +34 -0
- sqlspec/data_dictionary/dialects/postgres.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/postgres.py +46 -0
- sqlspec/data_dictionary/dialects/spanner.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/spanner.py +37 -0
- sqlspec/data_dictionary/dialects/sqlite.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/data_dictionary/dialects/sqlite.py +42 -0
- sqlspec/data_dictionary/sql/.gitkeep +0 -0
- sqlspec/data_dictionary/sql/bigquery/columns.sql +23 -0
- sqlspec/data_dictionary/sql/bigquery/foreign_keys.sql +34 -0
- sqlspec/data_dictionary/sql/bigquery/indexes.sql +19 -0
- sqlspec/data_dictionary/sql/bigquery/tables.sql +33 -0
- sqlspec/data_dictionary/sql/bigquery/version.sql +3 -0
- sqlspec/data_dictionary/sql/cockroachdb/columns.sql +34 -0
- sqlspec/data_dictionary/sql/cockroachdb/foreign_keys.sql +40 -0
- sqlspec/data_dictionary/sql/cockroachdb/indexes.sql +32 -0
- sqlspec/data_dictionary/sql/cockroachdb/tables.sql +44 -0
- sqlspec/data_dictionary/sql/cockroachdb/version.sql +3 -0
- sqlspec/data_dictionary/sql/duckdb/columns.sql +23 -0
- sqlspec/data_dictionary/sql/duckdb/foreign_keys.sql +36 -0
- sqlspec/data_dictionary/sql/duckdb/indexes.sql +19 -0
- sqlspec/data_dictionary/sql/duckdb/tables.sql +38 -0
- sqlspec/data_dictionary/sql/duckdb/version.sql +3 -0
- sqlspec/data_dictionary/sql/mysql/columns.sql +23 -0
- sqlspec/data_dictionary/sql/mysql/foreign_keys.sql +28 -0
- sqlspec/data_dictionary/sql/mysql/indexes.sql +26 -0
- sqlspec/data_dictionary/sql/mysql/tables.sql +33 -0
- sqlspec/data_dictionary/sql/mysql/version.sql +3 -0
- sqlspec/data_dictionary/sql/oracle/columns.sql +23 -0
- sqlspec/data_dictionary/sql/oracle/foreign_keys.sql +48 -0
- sqlspec/data_dictionary/sql/oracle/indexes.sql +44 -0
- sqlspec/data_dictionary/sql/oracle/tables.sql +25 -0
- sqlspec/data_dictionary/sql/oracle/version.sql +20 -0
- sqlspec/data_dictionary/sql/postgres/columns.sql +34 -0
- sqlspec/data_dictionary/sql/postgres/foreign_keys.sql +40 -0
- sqlspec/data_dictionary/sql/postgres/indexes.sql +56 -0
- sqlspec/data_dictionary/sql/postgres/tables.sql +44 -0
- sqlspec/data_dictionary/sql/postgres/version.sql +3 -0
- sqlspec/data_dictionary/sql/spanner/columns.sql +23 -0
- sqlspec/data_dictionary/sql/spanner/foreign_keys.sql +70 -0
- sqlspec/data_dictionary/sql/spanner/indexes.sql +30 -0
- sqlspec/data_dictionary/sql/spanner/tables.sql +9 -0
- sqlspec/data_dictionary/sql/spanner/version.sql +3 -0
- sqlspec/data_dictionary/sql/sqlite/columns.sql +23 -0
- sqlspec/data_dictionary/sql/sqlite/foreign_keys.sql +22 -0
- sqlspec/data_dictionary/sql/sqlite/indexes.sql +7 -0
- sqlspec/data_dictionary/sql/sqlite/tables.sql +28 -0
- sqlspec/data_dictionary/sql/sqlite/version.sql +3 -0
- sqlspec/driver/__init__.py +32 -0
- sqlspec/driver/_async.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/driver/_async.py +1737 -0
- sqlspec/driver/_common.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/driver/_common.py +1478 -0
- sqlspec/driver/_sql_helpers.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/driver/_sql_helpers.py +148 -0
- sqlspec/driver/_storage_helpers.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/driver/_storage_helpers.py +144 -0
- sqlspec/driver/_sync.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/driver/_sync.py +1710 -0
- sqlspec/exceptions.py +338 -0
- sqlspec/extensions/__init__.py +0 -0
- sqlspec/extensions/adk/__init__.py +70 -0
- sqlspec/extensions/adk/_types.py +51 -0
- sqlspec/extensions/adk/converters.py +172 -0
- sqlspec/extensions/adk/memory/__init__.py +69 -0
- sqlspec/extensions/adk/memory/_types.py +30 -0
- sqlspec/extensions/adk/memory/converters.py +149 -0
- sqlspec/extensions/adk/memory/service.py +217 -0
- sqlspec/extensions/adk/memory/store.py +569 -0
- sqlspec/extensions/adk/migrations/0001_create_adk_tables.py +246 -0
- sqlspec/extensions/adk/migrations/__init__.py +0 -0
- sqlspec/extensions/adk/service.py +225 -0
- sqlspec/extensions/adk/store.py +567 -0
- sqlspec/extensions/events/__init__.py +51 -0
- sqlspec/extensions/events/_channel.py +703 -0
- sqlspec/extensions/events/_hints.py +45 -0
- sqlspec/extensions/events/_models.py +23 -0
- sqlspec/extensions/events/_payload.py +69 -0
- sqlspec/extensions/events/_protocols.py +134 -0
- sqlspec/extensions/events/_queue.py +461 -0
- sqlspec/extensions/events/_store.py +209 -0
- sqlspec/extensions/events/migrations/0001_create_event_queue.py +59 -0
- sqlspec/extensions/events/migrations/__init__.py +3 -0
- sqlspec/extensions/fastapi/__init__.py +19 -0
- sqlspec/extensions/fastapi/extension.py +351 -0
- sqlspec/extensions/fastapi/providers.py +607 -0
- sqlspec/extensions/flask/__init__.py +37 -0
- sqlspec/extensions/flask/_state.py +76 -0
- sqlspec/extensions/flask/_utils.py +71 -0
- sqlspec/extensions/flask/extension.py +519 -0
- sqlspec/extensions/litestar/__init__.py +28 -0
- sqlspec/extensions/litestar/_utils.py +52 -0
- sqlspec/extensions/litestar/channels.py +165 -0
- sqlspec/extensions/litestar/cli.py +102 -0
- sqlspec/extensions/litestar/config.py +90 -0
- sqlspec/extensions/litestar/handlers.py +316 -0
- sqlspec/extensions/litestar/migrations/0001_create_session_table.py +137 -0
- sqlspec/extensions/litestar/migrations/__init__.py +3 -0
- sqlspec/extensions/litestar/plugin.py +671 -0
- sqlspec/extensions/litestar/providers.py +526 -0
- sqlspec/extensions/litestar/store.py +296 -0
- sqlspec/extensions/otel/__init__.py +58 -0
- sqlspec/extensions/prometheus/__init__.py +113 -0
- sqlspec/extensions/starlette/__init__.py +19 -0
- sqlspec/extensions/starlette/_state.py +30 -0
- sqlspec/extensions/starlette/_utils.py +96 -0
- sqlspec/extensions/starlette/extension.py +346 -0
- sqlspec/extensions/starlette/middleware.py +235 -0
- sqlspec/loader.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/loader.py +702 -0
- sqlspec/migrations/__init__.py +36 -0
- sqlspec/migrations/base.py +731 -0
- sqlspec/migrations/commands.py +1232 -0
- sqlspec/migrations/context.py +157 -0
- sqlspec/migrations/fix.py +204 -0
- sqlspec/migrations/loaders.py +443 -0
- sqlspec/migrations/runner.py +1172 -0
- sqlspec/migrations/templates.py +234 -0
- sqlspec/migrations/tracker.py +611 -0
- sqlspec/migrations/utils.py +256 -0
- sqlspec/migrations/validation.py +207 -0
- sqlspec/migrations/version.py +446 -0
- sqlspec/observability/__init__.py +55 -0
- sqlspec/observability/_common.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_common.py +77 -0
- sqlspec/observability/_config.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_config.py +348 -0
- sqlspec/observability/_diagnostics.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_diagnostics.py +74 -0
- sqlspec/observability/_dispatcher.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_dispatcher.py +152 -0
- sqlspec/observability/_formatters/__init__.py +13 -0
- sqlspec/observability/_formatters/_aws.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_formatters/_aws.py +102 -0
- sqlspec/observability/_formatters/_azure.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_formatters/_azure.py +96 -0
- sqlspec/observability/_formatters/_base.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_formatters/_base.py +57 -0
- sqlspec/observability/_formatters/_gcp.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_formatters/_gcp.py +131 -0
- sqlspec/observability/_formatting.py +58 -0
- sqlspec/observability/_observer.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_observer.py +357 -0
- sqlspec/observability/_runtime.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_runtime.py +420 -0
- sqlspec/observability/_sampling.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_sampling.py +188 -0
- sqlspec/observability/_spans.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/observability/_spans.py +161 -0
- sqlspec/protocols.py +916 -0
- sqlspec/py.typed +0 -0
- sqlspec/storage/__init__.py +48 -0
- sqlspec/storage/_utils.py +104 -0
- sqlspec/storage/backends/__init__.py +1 -0
- sqlspec/storage/backends/base.py +253 -0
- sqlspec/storage/backends/fsspec.py +529 -0
- sqlspec/storage/backends/local.py +441 -0
- sqlspec/storage/backends/obstore.py +916 -0
- sqlspec/storage/errors.py +104 -0
- sqlspec/storage/pipeline.py +582 -0
- sqlspec/storage/registry.py +301 -0
- sqlspec/typing.py +395 -0
- sqlspec/utils/__init__.py +7 -0
- sqlspec/utils/arrow_helpers.py +318 -0
- sqlspec/utils/config_tools.py +332 -0
- sqlspec/utils/correlation.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/correlation.py +134 -0
- sqlspec/utils/deprecation.py +190 -0
- sqlspec/utils/fixtures.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/fixtures.py +258 -0
- sqlspec/utils/logging.py +222 -0
- sqlspec/utils/module_loader.py +306 -0
- sqlspec/utils/portal.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/portal.py +375 -0
- sqlspec/utils/schema.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/schema.py +485 -0
- sqlspec/utils/serializers.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/serializers.py +408 -0
- sqlspec/utils/singleton.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/singleton.py +41 -0
- sqlspec/utils/sync_tools.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/sync_tools.py +311 -0
- sqlspec/utils/text.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/text.py +108 -0
- sqlspec/utils/type_converters.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/type_converters.py +128 -0
- sqlspec/utils/type_guards.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/type_guards.py +1360 -0
- sqlspec/utils/uuids.cpython-310-aarch64-linux-gnu.so +0 -0
- sqlspec/utils/uuids.py +225 -0
- sqlspec-0.36.0.dist-info/METADATA +205 -0
- sqlspec-0.36.0.dist-info/RECORD +531 -0
- sqlspec-0.36.0.dist-info/WHEEL +7 -0
- sqlspec-0.36.0.dist-info/entry_points.txt +2 -0
- sqlspec-0.36.0.dist-info/licenses/LICENSE +21 -0
|
@@ -0,0 +1,827 @@
|
|
|
1
|
+
"""BigQuery ADK store for Google Agent Development Kit session/event storage."""
|
|
2
|
+
|
|
3
|
+
from collections.abc import Mapping
|
|
4
|
+
from datetime import datetime, timezone
|
|
5
|
+
from typing import TYPE_CHECKING, Any, cast
|
|
6
|
+
|
|
7
|
+
from google.api_core.exceptions import NotFound
|
|
8
|
+
from google.cloud.bigquery import QueryJobConfig, ScalarQueryParameter
|
|
9
|
+
|
|
10
|
+
from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
|
|
11
|
+
from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore
|
|
12
|
+
from sqlspec.utils.serializers import from_json, to_json
|
|
13
|
+
from sqlspec.utils.sync_tools import async_, run_
|
|
14
|
+
|
|
15
|
+
if TYPE_CHECKING:
|
|
16
|
+
from sqlspec.adapters.bigquery.config import BigQueryConfig
|
|
17
|
+
from sqlspec.extensions.adk import MemoryRecord
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
__all__ = ("BigQueryADKMemoryStore", "BigQueryADKStore")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class BigQueryADKStore(BaseAsyncADKStore["BigQueryConfig"]):
|
|
24
|
+
"""BigQuery ADK store using synchronous BigQuery client with async wrapper.
|
|
25
|
+
|
|
26
|
+
Implements session and event storage for Google Agent Development Kit
|
|
27
|
+
using Google Cloud BigQuery. Uses BigQuery's native JSON type for state/metadata
|
|
28
|
+
storage and async_() wrapper to provide async interface.
|
|
29
|
+
|
|
30
|
+
Provides:
|
|
31
|
+
- Serverless, scalable session state management with JSON storage
|
|
32
|
+
- Event history tracking optimized for analytics
|
|
33
|
+
- Microsecond-precision timestamps with TIMESTAMP type
|
|
34
|
+
- Cost-optimized queries with partitioning and clustering
|
|
35
|
+
- Efficient JSON handling with BigQuery's JSON type
|
|
36
|
+
- Manual cascade delete pattern (no foreign key support)
|
|
37
|
+
|
|
38
|
+
Args:
|
|
39
|
+
config: BigQueryConfig with extension_config["adk"] settings.
|
|
40
|
+
|
|
41
|
+
Example:
|
|
42
|
+
from sqlspec.adapters.bigquery import BigQueryConfig
|
|
43
|
+
from sqlspec.adapters.bigquery.adk import BigQueryADKStore
|
|
44
|
+
|
|
45
|
+
config = BigQueryConfig(
|
|
46
|
+
connection_config={
|
|
47
|
+
"project": "my-project",
|
|
48
|
+
"dataset_id": "my_dataset",
|
|
49
|
+
},
|
|
50
|
+
extension_config={
|
|
51
|
+
"adk": {
|
|
52
|
+
"session_table": "my_sessions",
|
|
53
|
+
"events_table": "my_events",
|
|
54
|
+
"owner_id_column": "tenant_id INT64 NOT NULL"
|
|
55
|
+
}
|
|
56
|
+
}
|
|
57
|
+
)
|
|
58
|
+
store = BigQueryADKStore(config)
|
|
59
|
+
await store.ensure_tables()
|
|
60
|
+
|
|
61
|
+
Notes:
|
|
62
|
+
- JSON type for state, content, and metadata (native BigQuery JSON)
|
|
63
|
+
- BYTES for pre-serialized actions from Google ADK
|
|
64
|
+
- TIMESTAMP for timezone-aware microsecond precision
|
|
65
|
+
- Partitioned by DATE(create_time) for cost optimization
|
|
66
|
+
- Clustered by app_name, user_id for query performance
|
|
67
|
+
- Uses to_json/from_json for serialization to JSON columns
|
|
68
|
+
- BigQuery has eventual consistency - handle appropriately
|
|
69
|
+
- No true foreign keys but implements cascade delete pattern
|
|
70
|
+
- Configuration is read from config.extension_config["adk"]
|
|
71
|
+
"""
|
|
72
|
+
|
|
73
|
+
__slots__ = ("_dataset_id",)
|
|
74
|
+
|
|
75
|
+
def __init__(self, config: "BigQueryConfig") -> None:
|
|
76
|
+
"""Initialize BigQuery ADK store.
|
|
77
|
+
|
|
78
|
+
Args:
|
|
79
|
+
config: BigQueryConfig instance.
|
|
80
|
+
|
|
81
|
+
Notes:
|
|
82
|
+
Configuration is read from config.extension_config["adk"]:
|
|
83
|
+
- session_table: Sessions table name (default: "adk_sessions")
|
|
84
|
+
- events_table: Events table name (default: "adk_events")
|
|
85
|
+
- owner_id_column: Optional owner FK column DDL (default: None)
|
|
86
|
+
"""
|
|
87
|
+
super().__init__(config)
|
|
88
|
+
self._dataset_id = config.connection_config.get("dataset_id")
|
|
89
|
+
|
|
90
|
+
def _get_full_table_name(self, table_name: str) -> str:
|
|
91
|
+
"""Get fully qualified table name for BigQuery.
|
|
92
|
+
|
|
93
|
+
Args:
|
|
94
|
+
table_name: Base table name.
|
|
95
|
+
|
|
96
|
+
Returns:
|
|
97
|
+
Fully qualified table name with backticks.
|
|
98
|
+
|
|
99
|
+
Notes:
|
|
100
|
+
BigQuery requires backtick-quoted identifiers for table names.
|
|
101
|
+
Format: `project.dataset.table` or `dataset.table`
|
|
102
|
+
"""
|
|
103
|
+
if self._dataset_id:
|
|
104
|
+
return f"`{self._dataset_id}.{table_name}`"
|
|
105
|
+
return f"`{table_name}`"
|
|
106
|
+
|
|
107
|
+
async def _get_create_sessions_table_sql(self) -> str:
|
|
108
|
+
"""Get BigQuery CREATE TABLE SQL for sessions.
|
|
109
|
+
|
|
110
|
+
Returns:
|
|
111
|
+
SQL statement to create adk_sessions table.
|
|
112
|
+
|
|
113
|
+
Notes:
|
|
114
|
+
- STRING for IDs and names
|
|
115
|
+
- JSON type for state storage (native BigQuery JSON)
|
|
116
|
+
- TIMESTAMP for timezone-aware microsecond precision
|
|
117
|
+
- Partitioned by DATE(create_time) for cost optimization
|
|
118
|
+
- Clustered by app_name, user_id for query performance
|
|
119
|
+
- No indexes needed (BigQuery auto-optimizes)
|
|
120
|
+
- Optional owner ID column for multi-tenant scenarios
|
|
121
|
+
- Note: BigQuery doesn't enforce FK constraints
|
|
122
|
+
"""
|
|
123
|
+
owner_id_line = ""
|
|
124
|
+
if self._owner_id_column_ddl:
|
|
125
|
+
owner_id_line = f",\n {self._owner_id_column_ddl}"
|
|
126
|
+
|
|
127
|
+
table_name = self._get_full_table_name(self._session_table)
|
|
128
|
+
return f"""
|
|
129
|
+
CREATE TABLE IF NOT EXISTS {table_name} (
|
|
130
|
+
id STRING NOT NULL,
|
|
131
|
+
app_name STRING NOT NULL,
|
|
132
|
+
user_id STRING NOT NULL{owner_id_line},
|
|
133
|
+
state JSON NOT NULL,
|
|
134
|
+
create_time TIMESTAMP NOT NULL,
|
|
135
|
+
update_time TIMESTAMP NOT NULL
|
|
136
|
+
)
|
|
137
|
+
PARTITION BY DATE(create_time)
|
|
138
|
+
CLUSTER BY app_name, user_id
|
|
139
|
+
"""
|
|
140
|
+
|
|
141
|
+
async def _get_create_events_table_sql(self) -> str:
|
|
142
|
+
"""Get BigQuery CREATE TABLE SQL for events.
|
|
143
|
+
|
|
144
|
+
Returns:
|
|
145
|
+
SQL statement to create adk_events table.
|
|
146
|
+
|
|
147
|
+
Notes:
|
|
148
|
+
- STRING for IDs and text fields
|
|
149
|
+
- BYTES for pickled actions
|
|
150
|
+
- JSON for content, grounding_metadata, custom_metadata, long_running_tool_ids_json
|
|
151
|
+
- BOOL for boolean flags
|
|
152
|
+
- TIMESTAMP for timezone-aware timestamps
|
|
153
|
+
- Partitioned by DATE(timestamp) for cost optimization
|
|
154
|
+
- Clustered by session_id, timestamp for ordered retrieval
|
|
155
|
+
"""
|
|
156
|
+
table_name = self._get_full_table_name(self._events_table)
|
|
157
|
+
return f"""
|
|
158
|
+
CREATE TABLE IF NOT EXISTS {table_name} (
|
|
159
|
+
id STRING NOT NULL,
|
|
160
|
+
session_id STRING NOT NULL,
|
|
161
|
+
app_name STRING NOT NULL,
|
|
162
|
+
user_id STRING NOT NULL,
|
|
163
|
+
invocation_id STRING,
|
|
164
|
+
author STRING,
|
|
165
|
+
actions BYTES,
|
|
166
|
+
long_running_tool_ids_json JSON,
|
|
167
|
+
branch STRING,
|
|
168
|
+
timestamp TIMESTAMP NOT NULL,
|
|
169
|
+
content JSON,
|
|
170
|
+
grounding_metadata JSON,
|
|
171
|
+
custom_metadata JSON,
|
|
172
|
+
partial BOOL,
|
|
173
|
+
turn_complete BOOL,
|
|
174
|
+
interrupted BOOL,
|
|
175
|
+
error_code STRING,
|
|
176
|
+
error_message STRING
|
|
177
|
+
)
|
|
178
|
+
PARTITION BY DATE(timestamp)
|
|
179
|
+
CLUSTER BY session_id, timestamp
|
|
180
|
+
"""
|
|
181
|
+
|
|
182
|
+
def _get_drop_tables_sql(self) -> "list[str]":
|
|
183
|
+
"""Get BigQuery DROP TABLE SQL statements.
|
|
184
|
+
|
|
185
|
+
Returns:
|
|
186
|
+
List of SQL statements to drop tables.
|
|
187
|
+
|
|
188
|
+
Notes:
|
|
189
|
+
Order matters: drop events table before sessions table.
|
|
190
|
+
BigQuery uses IF EXISTS for idempotent drops.
|
|
191
|
+
"""
|
|
192
|
+
events_table = self._get_full_table_name(self._events_table)
|
|
193
|
+
sessions_table = self._get_full_table_name(self._session_table)
|
|
194
|
+
return [f"DROP TABLE IF EXISTS {events_table}", f"DROP TABLE IF EXISTS {sessions_table}"]
|
|
195
|
+
|
|
196
|
+
def _create_tables(self) -> None:
|
|
197
|
+
"""Synchronous implementation of create_tables."""
|
|
198
|
+
with self._config.provide_session() as driver:
|
|
199
|
+
driver.execute_script(run_(self._get_create_sessions_table_sql)())
|
|
200
|
+
driver.execute_script(run_(self._get_create_events_table_sql)())
|
|
201
|
+
|
|
202
|
+
async def create_tables(self) -> None:
|
|
203
|
+
"""Create both sessions and events tables if they don't exist."""
|
|
204
|
+
await async_(self._create_tables)()
|
|
205
|
+
|
|
206
|
+
def _create_session(
|
|
207
|
+
self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
|
|
208
|
+
) -> SessionRecord:
|
|
209
|
+
"""Synchronous implementation of create_session."""
|
|
210
|
+
now = datetime.now(timezone.utc)
|
|
211
|
+
state_json = to_json(state) if state else "{}"
|
|
212
|
+
|
|
213
|
+
table_name = self._get_full_table_name(self._session_table)
|
|
214
|
+
|
|
215
|
+
if self._owner_id_column_name:
|
|
216
|
+
sql = f"""
|
|
217
|
+
INSERT INTO {table_name} (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
|
|
218
|
+
VALUES (@id, @app_name, @user_id, @owner_id, JSON(@state), @create_time, @update_time)
|
|
219
|
+
"""
|
|
220
|
+
|
|
221
|
+
params = [
|
|
222
|
+
ScalarQueryParameter("id", "STRING", session_id),
|
|
223
|
+
ScalarQueryParameter("app_name", "STRING", app_name),
|
|
224
|
+
ScalarQueryParameter("user_id", "STRING", user_id),
|
|
225
|
+
ScalarQueryParameter("owner_id", "STRING", str(owner_id) if owner_id is not None else None),
|
|
226
|
+
ScalarQueryParameter("state", "STRING", state_json),
|
|
227
|
+
ScalarQueryParameter("create_time", "TIMESTAMP", now),
|
|
228
|
+
ScalarQueryParameter("update_time", "TIMESTAMP", now),
|
|
229
|
+
]
|
|
230
|
+
else:
|
|
231
|
+
sql = f"""
|
|
232
|
+
INSERT INTO {table_name} (id, app_name, user_id, state, create_time, update_time)
|
|
233
|
+
VALUES (@id, @app_name, @user_id, JSON(@state), @create_time, @update_time)
|
|
234
|
+
"""
|
|
235
|
+
|
|
236
|
+
params = [
|
|
237
|
+
ScalarQueryParameter("id", "STRING", session_id),
|
|
238
|
+
ScalarQueryParameter("app_name", "STRING", app_name),
|
|
239
|
+
ScalarQueryParameter("user_id", "STRING", user_id),
|
|
240
|
+
ScalarQueryParameter("state", "STRING", state_json),
|
|
241
|
+
ScalarQueryParameter("create_time", "TIMESTAMP", now),
|
|
242
|
+
ScalarQueryParameter("update_time", "TIMESTAMP", now),
|
|
243
|
+
]
|
|
244
|
+
|
|
245
|
+
with self._config.provide_connection() as conn:
|
|
246
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
247
|
+
conn.query(sql, job_config=job_config).result()
|
|
248
|
+
|
|
249
|
+
return SessionRecord(
|
|
250
|
+
id=session_id, app_name=app_name, user_id=user_id, state=state, create_time=now, update_time=now
|
|
251
|
+
)
|
|
252
|
+
|
|
253
|
+
async def create_session(
|
|
254
|
+
self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
|
|
255
|
+
) -> SessionRecord:
|
|
256
|
+
"""Create a new session.
|
|
257
|
+
|
|
258
|
+
Args:
|
|
259
|
+
session_id: Unique session identifier.
|
|
260
|
+
app_name: Application name.
|
|
261
|
+
user_id: User identifier.
|
|
262
|
+
state: Initial session state.
|
|
263
|
+
owner_id: Optional owner ID value for owner_id_column (if configured).
|
|
264
|
+
|
|
265
|
+
Returns:
|
|
266
|
+
Created session record.
|
|
267
|
+
|
|
268
|
+
Notes:
|
|
269
|
+
Uses CURRENT_TIMESTAMP() for timestamps.
|
|
270
|
+
State is JSON-serialized then stored in JSON column.
|
|
271
|
+
If owner_id_column is configured, owner_id value must be provided.
|
|
272
|
+
BigQuery doesn't enforce FK constraints, but column is useful for JOINs.
|
|
273
|
+
"""
|
|
274
|
+
return await async_(self._create_session)(session_id, app_name, user_id, state, owner_id)
|
|
275
|
+
|
|
276
|
+
def _get_session(self, session_id: str) -> "SessionRecord | None":
|
|
277
|
+
"""Synchronous implementation of get_session."""
|
|
278
|
+
table_name = self._get_full_table_name(self._session_table)
|
|
279
|
+
sql = f"""
|
|
280
|
+
SELECT id, app_name, user_id, JSON_VALUE(state) as state, create_time, update_time
|
|
281
|
+
FROM {table_name}
|
|
282
|
+
WHERE id = @session_id
|
|
283
|
+
"""
|
|
284
|
+
|
|
285
|
+
params = [ScalarQueryParameter("session_id", "STRING", session_id)]
|
|
286
|
+
|
|
287
|
+
with self._config.provide_connection() as conn:
|
|
288
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
289
|
+
query_job = conn.query(sql, job_config=job_config)
|
|
290
|
+
results = list(query_job.result())
|
|
291
|
+
|
|
292
|
+
if not results:
|
|
293
|
+
return None
|
|
294
|
+
|
|
295
|
+
row = results[0]
|
|
296
|
+
return SessionRecord(
|
|
297
|
+
id=row.id,
|
|
298
|
+
app_name=row.app_name,
|
|
299
|
+
user_id=row.user_id,
|
|
300
|
+
state=from_json(row.state) if row.state else {},
|
|
301
|
+
create_time=row.create_time,
|
|
302
|
+
update_time=row.update_time,
|
|
303
|
+
)
|
|
304
|
+
|
|
305
|
+
async def get_session(self, session_id: str) -> "SessionRecord | None":
|
|
306
|
+
"""Get session by ID.
|
|
307
|
+
|
|
308
|
+
Args:
|
|
309
|
+
session_id: Session identifier.
|
|
310
|
+
|
|
311
|
+
Returns:
|
|
312
|
+
Session record or None if not found.
|
|
313
|
+
|
|
314
|
+
Notes:
|
|
315
|
+
BigQuery returns datetime objects for TIMESTAMP columns.
|
|
316
|
+
JSON_VALUE extracts string representation for parsing.
|
|
317
|
+
"""
|
|
318
|
+
return await async_(self._get_session)(session_id)
|
|
319
|
+
|
|
320
|
+
def _update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
|
|
321
|
+
"""Synchronous implementation of update_session_state."""
|
|
322
|
+
now = datetime.now(timezone.utc)
|
|
323
|
+
state_json = to_json(state) if state else "{}"
|
|
324
|
+
|
|
325
|
+
table_name = self._get_full_table_name(self._session_table)
|
|
326
|
+
sql = f"""
|
|
327
|
+
UPDATE {table_name}
|
|
328
|
+
SET state = JSON(@state), update_time = @update_time
|
|
329
|
+
WHERE id = @session_id
|
|
330
|
+
"""
|
|
331
|
+
|
|
332
|
+
params = [
|
|
333
|
+
ScalarQueryParameter("state", "STRING", state_json),
|
|
334
|
+
ScalarQueryParameter("update_time", "TIMESTAMP", now),
|
|
335
|
+
ScalarQueryParameter("session_id", "STRING", session_id),
|
|
336
|
+
]
|
|
337
|
+
|
|
338
|
+
with self._config.provide_connection() as conn:
|
|
339
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
340
|
+
conn.query(sql, job_config=job_config).result()
|
|
341
|
+
|
|
342
|
+
async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
|
|
343
|
+
"""Update session state.
|
|
344
|
+
|
|
345
|
+
Args:
|
|
346
|
+
session_id: Session identifier.
|
|
347
|
+
state: New state dictionary (replaces existing state).
|
|
348
|
+
|
|
349
|
+
Notes:
|
|
350
|
+
Replaces entire state dictionary.
|
|
351
|
+
Updates update_time to CURRENT_TIMESTAMP().
|
|
352
|
+
"""
|
|
353
|
+
await async_(self._update_session_state)(session_id, state)
|
|
354
|
+
|
|
355
|
+
def _list_sessions(self, app_name: str, user_id: "str | None") -> "list[SessionRecord]":
|
|
356
|
+
"""Synchronous implementation of list_sessions."""
|
|
357
|
+
table_name = self._get_full_table_name(self._session_table)
|
|
358
|
+
|
|
359
|
+
if user_id is None:
|
|
360
|
+
sql = f"""
|
|
361
|
+
SELECT id, app_name, user_id, JSON_VALUE(state) as state, create_time, update_time
|
|
362
|
+
FROM {table_name}
|
|
363
|
+
WHERE app_name = @app_name
|
|
364
|
+
ORDER BY update_time DESC
|
|
365
|
+
"""
|
|
366
|
+
params = [ScalarQueryParameter("app_name", "STRING", app_name)]
|
|
367
|
+
else:
|
|
368
|
+
sql = f"""
|
|
369
|
+
SELECT id, app_name, user_id, JSON_VALUE(state) as state, create_time, update_time
|
|
370
|
+
FROM {table_name}
|
|
371
|
+
WHERE app_name = @app_name AND user_id = @user_id
|
|
372
|
+
ORDER BY update_time DESC
|
|
373
|
+
"""
|
|
374
|
+
params = [
|
|
375
|
+
ScalarQueryParameter("app_name", "STRING", app_name),
|
|
376
|
+
ScalarQueryParameter("user_id", "STRING", user_id),
|
|
377
|
+
]
|
|
378
|
+
|
|
379
|
+
with self._config.provide_connection() as conn:
|
|
380
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
381
|
+
query_job = conn.query(sql, job_config=job_config)
|
|
382
|
+
results = list(query_job.result())
|
|
383
|
+
|
|
384
|
+
return [
|
|
385
|
+
SessionRecord(
|
|
386
|
+
id=row.id,
|
|
387
|
+
app_name=row.app_name,
|
|
388
|
+
user_id=row.user_id,
|
|
389
|
+
state=from_json(row.state) if row.state else {},
|
|
390
|
+
create_time=row.create_time,
|
|
391
|
+
update_time=row.update_time,
|
|
392
|
+
)
|
|
393
|
+
for row in results
|
|
394
|
+
]
|
|
395
|
+
|
|
396
|
+
async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
|
|
397
|
+
"""List sessions for an app, optionally filtered by user.
|
|
398
|
+
|
|
399
|
+
Args:
|
|
400
|
+
app_name: Application name.
|
|
401
|
+
user_id: User identifier. If None, lists all sessions for the app.
|
|
402
|
+
|
|
403
|
+
Returns:
|
|
404
|
+
List of session records ordered by update_time DESC.
|
|
405
|
+
|
|
406
|
+
Notes:
|
|
407
|
+
Uses clustering on (app_name, user_id) when user_id is provided for efficiency.
|
|
408
|
+
"""
|
|
409
|
+
return await async_(self._list_sessions)(app_name, user_id)
|
|
410
|
+
|
|
411
|
+
def _delete_session(self, session_id: str) -> None:
|
|
412
|
+
"""Synchronous implementation of delete_session."""
|
|
413
|
+
events_table = self._get_full_table_name(self._events_table)
|
|
414
|
+
sessions_table = self._get_full_table_name(self._session_table)
|
|
415
|
+
|
|
416
|
+
params = [ScalarQueryParameter("session_id", "STRING", session_id)]
|
|
417
|
+
|
|
418
|
+
with self._config.provide_connection() as conn:
|
|
419
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
420
|
+
conn.query(f"DELETE FROM {events_table} WHERE session_id = @session_id", job_config=job_config).result()
|
|
421
|
+
conn.query(f"DELETE FROM {sessions_table} WHERE id = @session_id", job_config=job_config).result()
|
|
422
|
+
|
|
423
|
+
async def delete_session(self, session_id: str) -> None:
|
|
424
|
+
"""Delete session and all associated events.
|
|
425
|
+
|
|
426
|
+
Args:
|
|
427
|
+
session_id: Session identifier.
|
|
428
|
+
|
|
429
|
+
Notes:
|
|
430
|
+
BigQuery doesn't support foreign keys, so we manually delete events first.
|
|
431
|
+
Uses two separate DELETE statements in sequence.
|
|
432
|
+
"""
|
|
433
|
+
await async_(self._delete_session)(session_id)
|
|
434
|
+
|
|
435
|
+
def _append_event(self, event_record: EventRecord) -> None:
|
|
436
|
+
"""Synchronous implementation of append_event."""
|
|
437
|
+
table_name = self._get_full_table_name(self._events_table)
|
|
438
|
+
|
|
439
|
+
content_json = to_json(event_record.get("content")) if event_record.get("content") else None
|
|
440
|
+
grounding_metadata_json = (
|
|
441
|
+
to_json(event_record.get("grounding_metadata")) if event_record.get("grounding_metadata") else None
|
|
442
|
+
)
|
|
443
|
+
custom_metadata_json = (
|
|
444
|
+
to_json(event_record.get("custom_metadata")) if event_record.get("custom_metadata") else None
|
|
445
|
+
)
|
|
446
|
+
|
|
447
|
+
sql = f"""
|
|
448
|
+
INSERT INTO {table_name} (
|
|
449
|
+
id, session_id, app_name, user_id, invocation_id, author, actions,
|
|
450
|
+
long_running_tool_ids_json, branch, timestamp, content,
|
|
451
|
+
grounding_metadata, custom_metadata, partial, turn_complete,
|
|
452
|
+
interrupted, error_code, error_message
|
|
453
|
+
) VALUES (
|
|
454
|
+
@id, @session_id, @app_name, @user_id, @invocation_id, @author, @actions,
|
|
455
|
+
@long_running_tool_ids_json, @branch, @timestamp,
|
|
456
|
+
{"JSON(@content)" if content_json else "NULL"},
|
|
457
|
+
{"JSON(@grounding_metadata)" if grounding_metadata_json else "NULL"},
|
|
458
|
+
{"JSON(@custom_metadata)" if custom_metadata_json else "NULL"},
|
|
459
|
+
@partial, @turn_complete, @interrupted, @error_code, @error_message
|
|
460
|
+
)
|
|
461
|
+
"""
|
|
462
|
+
|
|
463
|
+
actions_value = event_record.get("actions")
|
|
464
|
+
params = [
|
|
465
|
+
ScalarQueryParameter("id", "STRING", event_record["id"]),
|
|
466
|
+
ScalarQueryParameter("session_id", "STRING", event_record["session_id"]),
|
|
467
|
+
ScalarQueryParameter("app_name", "STRING", event_record["app_name"]),
|
|
468
|
+
ScalarQueryParameter("user_id", "STRING", event_record["user_id"]),
|
|
469
|
+
ScalarQueryParameter("invocation_id", "STRING", event_record.get("invocation_id")),
|
|
470
|
+
ScalarQueryParameter("author", "STRING", event_record.get("author")),
|
|
471
|
+
ScalarQueryParameter(
|
|
472
|
+
"actions",
|
|
473
|
+
"BYTES",
|
|
474
|
+
actions_value.decode("latin1") if isinstance(actions_value, bytes) else actions_value,
|
|
475
|
+
),
|
|
476
|
+
ScalarQueryParameter(
|
|
477
|
+
"long_running_tool_ids_json", "STRING", event_record.get("long_running_tool_ids_json")
|
|
478
|
+
),
|
|
479
|
+
ScalarQueryParameter("branch", "STRING", event_record.get("branch")),
|
|
480
|
+
ScalarQueryParameter("timestamp", "TIMESTAMP", event_record["timestamp"]),
|
|
481
|
+
ScalarQueryParameter("partial", "BOOL", event_record.get("partial")),
|
|
482
|
+
ScalarQueryParameter("turn_complete", "BOOL", event_record.get("turn_complete")),
|
|
483
|
+
ScalarQueryParameter("interrupted", "BOOL", event_record.get("interrupted")),
|
|
484
|
+
ScalarQueryParameter("error_code", "STRING", event_record.get("error_code")),
|
|
485
|
+
ScalarQueryParameter("error_message", "STRING", event_record.get("error_message")),
|
|
486
|
+
]
|
|
487
|
+
|
|
488
|
+
if content_json:
|
|
489
|
+
params.append(ScalarQueryParameter("content", "STRING", content_json))
|
|
490
|
+
if grounding_metadata_json:
|
|
491
|
+
params.append(ScalarQueryParameter("grounding_metadata", "STRING", grounding_metadata_json))
|
|
492
|
+
if custom_metadata_json:
|
|
493
|
+
params.append(ScalarQueryParameter("custom_metadata", "STRING", custom_metadata_json))
|
|
494
|
+
|
|
495
|
+
with self._config.provide_connection() as conn:
|
|
496
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
497
|
+
conn.query(sql, job_config=job_config).result()
|
|
498
|
+
|
|
499
|
+
async def append_event(self, event_record: EventRecord) -> None:
|
|
500
|
+
"""Append an event to a session.
|
|
501
|
+
|
|
502
|
+
Args:
|
|
503
|
+
event_record: Event record to store.
|
|
504
|
+
|
|
505
|
+
Notes:
|
|
506
|
+
Uses BigQuery TIMESTAMP for timezone-aware timestamps.
|
|
507
|
+
JSON fields are serialized to STRING then cast to JSON.
|
|
508
|
+
Boolean fields stored natively as BOOL.
|
|
509
|
+
"""
|
|
510
|
+
await async_(self._append_event)(event_record)
|
|
511
|
+
|
|
512
|
+
def _get_events(
|
|
513
|
+
self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
|
|
514
|
+
) -> "list[EventRecord]":
|
|
515
|
+
"""Synchronous implementation of get_events."""
|
|
516
|
+
table_name = self._get_full_table_name(self._events_table)
|
|
517
|
+
|
|
518
|
+
where_clauses = ["session_id = @session_id"]
|
|
519
|
+
params: list[ScalarQueryParameter] = [ScalarQueryParameter("session_id", "STRING", session_id)]
|
|
520
|
+
|
|
521
|
+
if after_timestamp is not None:
|
|
522
|
+
where_clauses.append("timestamp > @after_timestamp")
|
|
523
|
+
params.append(ScalarQueryParameter("after_timestamp", "TIMESTAMP", after_timestamp))
|
|
524
|
+
|
|
525
|
+
where_clause = " AND ".join(where_clauses)
|
|
526
|
+
limit_clause = f" LIMIT {limit}" if limit else ""
|
|
527
|
+
|
|
528
|
+
sql = f"""
|
|
529
|
+
SELECT id, session_id, app_name, user_id, invocation_id, author, actions,
|
|
530
|
+
long_running_tool_ids_json, branch, timestamp,
|
|
531
|
+
JSON_VALUE(content) as content,
|
|
532
|
+
JSON_VALUE(grounding_metadata) as grounding_metadata,
|
|
533
|
+
JSON_VALUE(custom_metadata) as custom_metadata,
|
|
534
|
+
partial, turn_complete, interrupted, error_code, error_message
|
|
535
|
+
FROM {table_name}
|
|
536
|
+
WHERE {where_clause}
|
|
537
|
+
ORDER BY timestamp ASC{limit_clause}
|
|
538
|
+
"""
|
|
539
|
+
|
|
540
|
+
with self._config.provide_connection() as conn:
|
|
541
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
542
|
+
query_job = conn.query(sql, job_config=job_config)
|
|
543
|
+
results = list(query_job.result())
|
|
544
|
+
|
|
545
|
+
return [
|
|
546
|
+
EventRecord(
|
|
547
|
+
id=row.id,
|
|
548
|
+
session_id=row.session_id,
|
|
549
|
+
app_name=row.app_name,
|
|
550
|
+
user_id=row.user_id,
|
|
551
|
+
invocation_id=row.invocation_id,
|
|
552
|
+
author=row.author,
|
|
553
|
+
actions=bytes(row.actions) if row.actions else b"",
|
|
554
|
+
long_running_tool_ids_json=row.long_running_tool_ids_json,
|
|
555
|
+
branch=row.branch,
|
|
556
|
+
timestamp=row.timestamp,
|
|
557
|
+
content=from_json(row.content) if row.content else None,
|
|
558
|
+
grounding_metadata=from_json(row.grounding_metadata) if row.grounding_metadata else None,
|
|
559
|
+
custom_metadata=from_json(row.custom_metadata) if row.custom_metadata else None,
|
|
560
|
+
partial=row.partial,
|
|
561
|
+
turn_complete=row.turn_complete,
|
|
562
|
+
interrupted=row.interrupted,
|
|
563
|
+
error_code=row.error_code,
|
|
564
|
+
error_message=row.error_message,
|
|
565
|
+
)
|
|
566
|
+
for row in results
|
|
567
|
+
]
|
|
568
|
+
|
|
569
|
+
async def get_events(
|
|
570
|
+
self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
|
|
571
|
+
) -> "list[EventRecord]":
|
|
572
|
+
"""Get events for a session.
|
|
573
|
+
|
|
574
|
+
Args:
|
|
575
|
+
session_id: Session identifier.
|
|
576
|
+
after_timestamp: Only return events after this time.
|
|
577
|
+
limit: Maximum number of events to return.
|
|
578
|
+
|
|
579
|
+
Returns:
|
|
580
|
+
List of event records ordered by timestamp ASC.
|
|
581
|
+
|
|
582
|
+
Notes:
|
|
583
|
+
Uses clustering on (session_id, timestamp) for efficient retrieval.
|
|
584
|
+
Parses JSON fields and converts BYTES actions to bytes.
|
|
585
|
+
"""
|
|
586
|
+
return await async_(self._get_events)(session_id, after_timestamp, limit)
|
|
587
|
+
|
|
588
|
+
|
|
589
|
+
class BigQueryADKMemoryStore(BaseAsyncADKMemoryStore["BigQueryConfig"]):
|
|
590
|
+
"""BigQuery ADK memory store using synchronous BigQuery client with async wrapper."""
|
|
591
|
+
|
|
592
|
+
__slots__ = ("_dataset_id",)
|
|
593
|
+
|
|
594
|
+
def __init__(self, config: "BigQueryConfig") -> None:
|
|
595
|
+
"""Initialize BigQuery ADK memory store."""
|
|
596
|
+
super().__init__(config)
|
|
597
|
+
self._dataset_id = config.connection_config.get("dataset_id")
|
|
598
|
+
|
|
599
|
+
def _get_full_table_name(self, table_name: str) -> str:
|
|
600
|
+
"""Get fully qualified table name for BigQuery."""
|
|
601
|
+
if self._dataset_id:
|
|
602
|
+
return f"`{self._dataset_id}.{table_name}`"
|
|
603
|
+
return f"`{table_name}`"
|
|
604
|
+
|
|
605
|
+
async def _get_create_memory_table_sql(self) -> str:
|
|
606
|
+
"""Get BigQuery CREATE TABLE SQL for memory entries."""
|
|
607
|
+
owner_id_line = ""
|
|
608
|
+
if self._owner_id_column_ddl:
|
|
609
|
+
owner_id_line = f",\n {self._owner_id_column_ddl}"
|
|
610
|
+
|
|
611
|
+
table_name = self._get_full_table_name(self._memory_table)
|
|
612
|
+
fts_index = ""
|
|
613
|
+
if self._use_fts:
|
|
614
|
+
fts_index = f"""
|
|
615
|
+
CREATE SEARCH INDEX idx_{self._memory_table}_fts
|
|
616
|
+
ON {table_name}(content_text)
|
|
617
|
+
"""
|
|
618
|
+
|
|
619
|
+
return f"""
|
|
620
|
+
CREATE TABLE IF NOT EXISTS {table_name} (
|
|
621
|
+
id STRING NOT NULL,
|
|
622
|
+
session_id STRING NOT NULL,
|
|
623
|
+
app_name STRING NOT NULL,
|
|
624
|
+
user_id STRING NOT NULL,
|
|
625
|
+
event_id STRING NOT NULL,
|
|
626
|
+
author STRING{owner_id_line},
|
|
627
|
+
timestamp TIMESTAMP NOT NULL,
|
|
628
|
+
content_json JSON NOT NULL,
|
|
629
|
+
content_text STRING NOT NULL,
|
|
630
|
+
metadata_json JSON,
|
|
631
|
+
inserted_at TIMESTAMP NOT NULL
|
|
632
|
+
)
|
|
633
|
+
PARTITION BY DATE(timestamp)
|
|
634
|
+
CLUSTER BY app_name, user_id;
|
|
635
|
+
{fts_index}
|
|
636
|
+
"""
|
|
637
|
+
|
|
638
|
+
def _get_drop_memory_table_sql(self) -> "list[str]":
|
|
639
|
+
"""Get BigQuery DROP TABLE SQL statements."""
|
|
640
|
+
table_name = self._get_full_table_name(self._memory_table)
|
|
641
|
+
return [f"DROP TABLE IF EXISTS {table_name}"]
|
|
642
|
+
|
|
643
|
+
def _create_tables(self) -> None:
|
|
644
|
+
"""Synchronous implementation of create_tables."""
|
|
645
|
+
with self._config.provide_session() as driver:
|
|
646
|
+
driver.execute_script(run_(self._get_create_memory_table_sql)())
|
|
647
|
+
|
|
648
|
+
async def create_tables(self) -> None:
|
|
649
|
+
"""Create the memory table if it doesn't exist."""
|
|
650
|
+
if not self._enabled:
|
|
651
|
+
return
|
|
652
|
+
await async_(self._create_tables)()
|
|
653
|
+
|
|
654
|
+
def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
|
|
655
|
+
"""Synchronous implementation of insert_memory_entries."""
|
|
656
|
+
table_name = self._get_full_table_name(self._memory_table)
|
|
657
|
+
inserted_count = 0
|
|
658
|
+
|
|
659
|
+
with self._config.provide_connection() as conn:
|
|
660
|
+
for entry in entries:
|
|
661
|
+
content_json = to_json(entry["content_json"])
|
|
662
|
+
metadata_json = to_json(entry["metadata_json"]) if entry["metadata_json"] is not None else None
|
|
663
|
+
metadata_expr = "JSON(@metadata_json)" if metadata_json is not None else "NULL"
|
|
664
|
+
|
|
665
|
+
owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else ""
|
|
666
|
+
owner_value = ", @owner_id" if self._owner_id_column_name else ""
|
|
667
|
+
|
|
668
|
+
sql = f"""
|
|
669
|
+
MERGE {table_name} T
|
|
670
|
+
USING (SELECT @event_id AS event_id) S
|
|
671
|
+
ON T.event_id = S.event_id
|
|
672
|
+
WHEN NOT MATCHED THEN
|
|
673
|
+
INSERT (id, session_id, app_name, user_id, event_id, author{owner_column},
|
|
674
|
+
timestamp, content_json, content_text, metadata_json, inserted_at)
|
|
675
|
+
VALUES (@id, @session_id, @app_name, @user_id, @event_id, @author{owner_value},
|
|
676
|
+
@timestamp, JSON(@content_json), @content_text, {metadata_expr}, @inserted_at)
|
|
677
|
+
"""
|
|
678
|
+
|
|
679
|
+
params = [
|
|
680
|
+
ScalarQueryParameter("id", "STRING", entry["id"]),
|
|
681
|
+
ScalarQueryParameter("session_id", "STRING", entry["session_id"]),
|
|
682
|
+
ScalarQueryParameter("app_name", "STRING", entry["app_name"]),
|
|
683
|
+
ScalarQueryParameter("user_id", "STRING", entry["user_id"]),
|
|
684
|
+
ScalarQueryParameter("event_id", "STRING", entry["event_id"]),
|
|
685
|
+
ScalarQueryParameter("author", "STRING", entry["author"]),
|
|
686
|
+
ScalarQueryParameter("timestamp", "TIMESTAMP", entry["timestamp"]),
|
|
687
|
+
ScalarQueryParameter("content_json", "STRING", content_json),
|
|
688
|
+
ScalarQueryParameter("content_text", "STRING", entry["content_text"]),
|
|
689
|
+
ScalarQueryParameter("inserted_at", "TIMESTAMP", entry["inserted_at"]),
|
|
690
|
+
]
|
|
691
|
+
|
|
692
|
+
if self._owner_id_column_name:
|
|
693
|
+
params.append(ScalarQueryParameter("owner_id", "STRING", str(owner_id) if owner_id else None))
|
|
694
|
+
if metadata_json is not None:
|
|
695
|
+
params.append(ScalarQueryParameter("metadata_json", "STRING", metadata_json))
|
|
696
|
+
|
|
697
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
698
|
+
job = conn.query(sql, job_config=job_config)
|
|
699
|
+
job.result()
|
|
700
|
+
if job.num_dml_affected_rows:
|
|
701
|
+
inserted_count += int(job.num_dml_affected_rows)
|
|
702
|
+
|
|
703
|
+
return inserted_count
|
|
704
|
+
|
|
705
|
+
async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
|
|
706
|
+
"""Bulk insert memory entries with deduplication."""
|
|
707
|
+
if not self._enabled:
|
|
708
|
+
msg = "Memory store is disabled"
|
|
709
|
+
raise RuntimeError(msg)
|
|
710
|
+
|
|
711
|
+
if not entries:
|
|
712
|
+
return 0
|
|
713
|
+
|
|
714
|
+
return await async_(self._insert_memory_entries)(entries, owner_id)
|
|
715
|
+
|
|
716
|
+
def _search_entries(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]":
|
|
717
|
+
"""Synchronous implementation of search_entries."""
|
|
718
|
+
table_name = self._get_full_table_name(self._memory_table)
|
|
719
|
+
base_params = [
|
|
720
|
+
ScalarQueryParameter("app_name", "STRING", app_name),
|
|
721
|
+
ScalarQueryParameter("user_id", "STRING", user_id),
|
|
722
|
+
ScalarQueryParameter("limit", "INT64", limit),
|
|
723
|
+
]
|
|
724
|
+
|
|
725
|
+
if self._use_fts:
|
|
726
|
+
sql = f"""
|
|
727
|
+
SELECT id, session_id, app_name, user_id, event_id, author,
|
|
728
|
+
timestamp, content_json, content_text, metadata_json, inserted_at
|
|
729
|
+
FROM {table_name}
|
|
730
|
+
WHERE app_name = @app_name
|
|
731
|
+
AND user_id = @user_id
|
|
732
|
+
AND SEARCH(content_text, @query)
|
|
733
|
+
ORDER BY timestamp DESC
|
|
734
|
+
LIMIT @limit
|
|
735
|
+
"""
|
|
736
|
+
params = [*base_params, ScalarQueryParameter("query", "STRING", query)]
|
|
737
|
+
else:
|
|
738
|
+
sql = f"""
|
|
739
|
+
SELECT id, session_id, app_name, user_id, event_id, author,
|
|
740
|
+
timestamp, content_json, content_text, metadata_json, inserted_at
|
|
741
|
+
FROM {table_name}
|
|
742
|
+
WHERE app_name = @app_name
|
|
743
|
+
AND user_id = @user_id
|
|
744
|
+
AND LOWER(content_text) LIKE LOWER(@pattern)
|
|
745
|
+
ORDER BY timestamp DESC
|
|
746
|
+
LIMIT @limit
|
|
747
|
+
"""
|
|
748
|
+
pattern = f"%{query}%"
|
|
749
|
+
params = [*base_params, ScalarQueryParameter("pattern", "STRING", pattern)]
|
|
750
|
+
|
|
751
|
+
with self._config.provide_connection() as conn:
|
|
752
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
753
|
+
rows = conn.query(sql, job_config=job_config).result()
|
|
754
|
+
return _rows_to_records(rows)
|
|
755
|
+
|
|
756
|
+
async def search_entries(
|
|
757
|
+
self, query: str, app_name: str, user_id: str, limit: "int | None" = None
|
|
758
|
+
) -> "list[MemoryRecord]":
|
|
759
|
+
"""Search memory entries by text query."""
|
|
760
|
+
if not self._enabled:
|
|
761
|
+
msg = "Memory store is disabled"
|
|
762
|
+
raise RuntimeError(msg)
|
|
763
|
+
|
|
764
|
+
effective_limit = limit if limit is not None else self._max_results
|
|
765
|
+
|
|
766
|
+
try:
|
|
767
|
+
return await async_(self._search_entries)(query, app_name, user_id, effective_limit)
|
|
768
|
+
except NotFound:
|
|
769
|
+
return []
|
|
770
|
+
|
|
771
|
+
def _delete_entries_by_session(self, session_id: str) -> int:
|
|
772
|
+
table_name = self._get_full_table_name(self._memory_table)
|
|
773
|
+
sql = f"DELETE FROM {table_name} WHERE session_id = @session_id"
|
|
774
|
+
params = [ScalarQueryParameter("session_id", "STRING", session_id)]
|
|
775
|
+
with self._config.provide_connection() as conn:
|
|
776
|
+
job_config = QueryJobConfig(query_parameters=params)
|
|
777
|
+
job = conn.query(sql, job_config=job_config)
|
|
778
|
+
job.result()
|
|
779
|
+
return int(job.num_dml_affected_rows or 0)
|
|
780
|
+
|
|
781
|
+
async def delete_entries_by_session(self, session_id: str) -> int:
|
|
782
|
+
"""Delete all memory entries for a specific session."""
|
|
783
|
+
return await async_(self._delete_entries_by_session)(session_id)
|
|
784
|
+
|
|
785
|
+
def _delete_entries_older_than(self, days: int) -> int:
|
|
786
|
+
table_name = self._get_full_table_name(self._memory_table)
|
|
787
|
+
sql = f"""
|
|
788
|
+
DELETE FROM {table_name}
|
|
789
|
+
WHERE inserted_at < TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {days} DAY)
|
|
790
|
+
"""
|
|
791
|
+
with self._config.provide_connection() as conn:
|
|
792
|
+
job = conn.query(sql)
|
|
793
|
+
job.result()
|
|
794
|
+
return int(job.num_dml_affected_rows or 0)
|
|
795
|
+
|
|
796
|
+
async def delete_entries_older_than(self, days: int) -> int:
|
|
797
|
+
"""Delete memory entries older than specified days."""
|
|
798
|
+
return await async_(self._delete_entries_older_than)(days)
|
|
799
|
+
|
|
800
|
+
|
|
801
|
+
def _decode_json_field(value: Any) -> "dict[str, Any] | None":
|
|
802
|
+
if value is None:
|
|
803
|
+
return None
|
|
804
|
+
if isinstance(value, str):
|
|
805
|
+
return cast("dict[str, Any]", from_json(value))
|
|
806
|
+
if isinstance(value, Mapping):
|
|
807
|
+
return dict(value)
|
|
808
|
+
return None
|
|
809
|
+
|
|
810
|
+
|
|
811
|
+
def _rows_to_records(rows: Any) -> "list[MemoryRecord]":
|
|
812
|
+
return [
|
|
813
|
+
{
|
|
814
|
+
"id": row["id"],
|
|
815
|
+
"session_id": row["session_id"],
|
|
816
|
+
"app_name": row["app_name"],
|
|
817
|
+
"user_id": row["user_id"],
|
|
818
|
+
"event_id": row["event_id"],
|
|
819
|
+
"author": row["author"],
|
|
820
|
+
"timestamp": row["timestamp"],
|
|
821
|
+
"content_json": _decode_json_field(row["content_json"]) or {},
|
|
822
|
+
"content_text": row["content_text"],
|
|
823
|
+
"metadata_json": _decode_json_field(row["metadata_json"]),
|
|
824
|
+
"inserted_at": row["inserted_at"],
|
|
825
|
+
}
|
|
826
|
+
for row in rows
|
|
827
|
+
]
|