sqlspec 0.47.0__cp314-cp314-win_amd64.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.
- f68e0789eb443ecb1c2c__mypyc.cp314-win_amd64.pyd +0 -0
- sqlspec/__init__.py +167 -0
- sqlspec/__main__.py +12 -0
- sqlspec/__metadata__.py +14 -0
- sqlspec/_typing.py +714 -0
- sqlspec/adapters/__init__.py +0 -0
- sqlspec/adapters/adbc/__init__.py +13 -0
- sqlspec/adapters/adbc/_typing.py +106 -0
- sqlspec/adapters/adbc/adk/__init__.py +5 -0
- sqlspec/adapters/adbc/adk/store.py +1280 -0
- sqlspec/adapters/adbc/config.py +378 -0
- sqlspec/adapters/adbc/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/adbc/core.py +922 -0
- sqlspec/adapters/adbc/data_dictionary.py +339 -0
- sqlspec/adapters/adbc/driver.py +534 -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 +534 -0
- sqlspec/adapters/adbc/type_converter.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/adbc/type_converter.py +142 -0
- sqlspec/adapters/aiomysql/__init__.py +21 -0
- sqlspec/adapters/aiomysql/_typing.py +137 -0
- sqlspec/adapters/aiomysql/adk/__init__.py +5 -0
- sqlspec/adapters/aiomysql/adk/store.py +678 -0
- sqlspec/adapters/aiomysql/config.py +305 -0
- sqlspec/adapters/aiomysql/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/aiomysql/core.py +536 -0
- sqlspec/adapters/aiomysql/data_dictionary.py +121 -0
- sqlspec/adapters/aiomysql/driver.py +386 -0
- sqlspec/adapters/aiomysql/events/__init__.py +5 -0
- sqlspec/adapters/aiomysql/events/store.py +104 -0
- sqlspec/adapters/aiomysql/litestar/__init__.py +5 -0
- sqlspec/adapters/aiomysql/litestar/store.py +314 -0
- sqlspec/adapters/aiosqlite/__init__.py +26 -0
- sqlspec/adapters/aiosqlite/_typing.py +109 -0
- sqlspec/adapters/aiosqlite/adk/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/adk/store.py +829 -0
- sqlspec/adapters/aiosqlite/config.py +315 -0
- sqlspec/adapters/aiosqlite/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/aiosqlite/core.py +315 -0
- sqlspec/adapters/aiosqlite/data_dictionary.py +202 -0
- sqlspec/adapters/aiosqlite/driver.py +311 -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.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/aiosqlite/pool.py +734 -0
- sqlspec/adapters/asyncmy/__init__.py +21 -0
- sqlspec/adapters/asyncmy/_typing.py +113 -0
- sqlspec/adapters/asyncmy/adk/__init__.py +5 -0
- sqlspec/adapters/asyncmy/adk/store.py +644 -0
- sqlspec/adapters/asyncmy/config.py +307 -0
- sqlspec/adapters/asyncmy/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/asyncmy/core.py +538 -0
- sqlspec/adapters/asyncmy/data_dictionary.py +122 -0
- sqlspec/adapters/asyncmy/driver.py +391 -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 +26 -0
- sqlspec/adapters/asyncpg/_typing.py +103 -0
- sqlspec/adapters/asyncpg/adk/__init__.py +5 -0
- sqlspec/adapters/asyncpg/adk/store.py +483 -0
- sqlspec/adapters/asyncpg/config.py +575 -0
- sqlspec/adapters/asyncpg/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/asyncpg/core.py +480 -0
- sqlspec/adapters/asyncpg/data_dictionary.py +157 -0
- sqlspec/adapters/asyncpg/driver.py +487 -0
- sqlspec/adapters/asyncpg/events/__init__.py +6 -0
- sqlspec/adapters/asyncpg/events/_hub.py +181 -0
- sqlspec/adapters/asyncpg/events/backend.py +210 -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 +15 -0
- sqlspec/adapters/bigquery/_typing.py +108 -0
- sqlspec/adapters/bigquery/config.py +362 -0
- sqlspec/adapters/bigquery/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/bigquery/core.py +768 -0
- sqlspec/adapters/bigquery/data_dictionary.py +120 -0
- sqlspec/adapters/bigquery/driver.py +542 -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.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/bigquery/type_converter.py +107 -0
- sqlspec/adapters/cockroach_asyncpg/__init__.py +26 -0
- sqlspec/adapters/cockroach_asyncpg/_typing.py +73 -0
- sqlspec/adapters/cockroach_asyncpg/adk/__init__.py +3 -0
- sqlspec/adapters/cockroach_asyncpg/adk/store.py +465 -0
- sqlspec/adapters/cockroach_asyncpg/config.py +248 -0
- sqlspec/adapters/cockroach_asyncpg/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/cockroach_asyncpg/core.py +55 -0
- sqlspec/adapters/cockroach_asyncpg/data_dictionary.py +110 -0
- sqlspec/adapters/cockroach_asyncpg/driver.py +142 -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 +39 -0
- sqlspec/adapters/cockroach_psycopg/_typing.py +137 -0
- sqlspec/adapters/cockroach_psycopg/adk/__init__.py +13 -0
- sqlspec/adapters/cockroach_psycopg/adk/store.py +1039 -0
- sqlspec/adapters/cockroach_psycopg/config.py +511 -0
- sqlspec/adapters/cockroach_psycopg/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/cockroach_psycopg/core.py +63 -0
- sqlspec/adapters/cockroach_psycopg/data_dictionary.py +220 -0
- sqlspec/adapters/cockroach_psycopg/driver.py +273 -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 +327 -0
- sqlspec/adapters/duckdb/__init__.py +29 -0
- sqlspec/adapters/duckdb/_typing.py +104 -0
- sqlspec/adapters/duckdb/adk/__init__.py +14 -0
- sqlspec/adapters/duckdb/adk/store.py +935 -0
- sqlspec/adapters/duckdb/config.py +386 -0
- sqlspec/adapters/duckdb/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/duckdb/core.py +332 -0
- sqlspec/adapters/duckdb/data_dictionary.py +140 -0
- sqlspec/adapters/duckdb/driver.py +426 -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.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/duckdb/pool.py +350 -0
- sqlspec/adapters/duckdb/type_converter.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/duckdb/type_converter.py +118 -0
- sqlspec/adapters/mysqlconnector/__init__.py +39 -0
- sqlspec/adapters/mysqlconnector/_typing.py +186 -0
- sqlspec/adapters/mysqlconnector/adk/__init__.py +15 -0
- sqlspec/adapters/mysqlconnector/adk/store.py +1183 -0
- sqlspec/adapters/mysqlconnector/config.py +421 -0
- sqlspec/adapters/mysqlconnector/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/mysqlconnector/core.py +472 -0
- sqlspec/adapters/mysqlconnector/data_dictionary.py +230 -0
- sqlspec/adapters/mysqlconnector/driver.py +516 -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 +39 -0
- sqlspec/adapters/oracledb/_json_handlers.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/oracledb/_json_handlers.py +196 -0
- sqlspec/adapters/oracledb/_param_types.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/oracledb/_param_types.py +46 -0
- sqlspec/adapters/oracledb/_typing.py +258 -0
- sqlspec/adapters/oracledb/_uuid_handlers.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/oracledb/_uuid_handlers.py +163 -0
- sqlspec/adapters/oracledb/_vector_handlers.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/oracledb/_vector_handlers.py +228 -0
- sqlspec/adapters/oracledb/adk/__init__.py +21 -0
- sqlspec/adapters/oracledb/adk/store.py +2453 -0
- sqlspec/adapters/oracledb/config.py +575 -0
- sqlspec/adapters/oracledb/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/oracledb/core.py +820 -0
- sqlspec/adapters/oracledb/data_dictionary.py +404 -0
- sqlspec/adapters/oracledb/driver.py +1277 -0
- sqlspec/adapters/oracledb/events/__init__.py +16 -0
- sqlspec/adapters/oracledb/events/_hub.py +345 -0
- sqlspec/adapters/oracledb/events/backend.py +300 -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 +539 -0
- sqlspec/adapters/oracledb/type_converter.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/oracledb/type_converter.py +211 -0
- sqlspec/adapters/psqlpy/__init__.py +18 -0
- sqlspec/adapters/psqlpy/_typing.py +121 -0
- sqlspec/adapters/psqlpy/adk/__init__.py +5 -0
- sqlspec/adapters/psqlpy/adk/store.py +591 -0
- sqlspec/adapters/psqlpy/config.py +376 -0
- sqlspec/adapters/psqlpy/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/psqlpy/core.py +694 -0
- sqlspec/adapters/psqlpy/data_dictionary.py +121 -0
- sqlspec/adapters/psqlpy/driver.py +411 -0
- sqlspec/adapters/psqlpy/events/__init__.py +6 -0
- sqlspec/adapters/psqlpy/events/_hub.py +204 -0
- sqlspec/adapters/psqlpy/events/backend.py +210 -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.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/psqlpy/type_converter.py +113 -0
- sqlspec/adapters/psycopg/__init__.py +38 -0
- sqlspec/adapters/psycopg/_typing.py +218 -0
- sqlspec/adapters/psycopg/adk/__init__.py +10 -0
- sqlspec/adapters/psycopg/adk/store.py +1106 -0
- sqlspec/adapters/psycopg/config.py +695 -0
- sqlspec/adapters/psycopg/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/psycopg/core.py +520 -0
- sqlspec/adapters/psycopg/data_dictionary.py +278 -0
- sqlspec/adapters/psycopg/driver.py +1033 -0
- sqlspec/adapters/psycopg/events/__init__.py +20 -0
- sqlspec/adapters/psycopg/events/_hub.py +388 -0
- sqlspec/adapters/psycopg/events/backend.py +398 -0
- sqlspec/adapters/psycopg/events/store.py +42 -0
- sqlspec/adapters/psycopg/litestar/__init__.py +5 -0
- sqlspec/adapters/psycopg/litestar/store.py +554 -0
- sqlspec/adapters/psycopg/type_converter.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/psycopg/type_converter.py +93 -0
- sqlspec/adapters/pymysql/__init__.py +21 -0
- sqlspec/adapters/pymysql/_typing.py +92 -0
- sqlspec/adapters/pymysql/adk/__init__.py +5 -0
- sqlspec/adapters/pymysql/adk/store.py +657 -0
- sqlspec/adapters/pymysql/config.py +176 -0
- sqlspec/adapters/pymysql/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/pymysql/core.py +469 -0
- sqlspec/adapters/pymysql/data_dictionary.py +120 -0
- sqlspec/adapters/pymysql/driver.py +271 -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.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/pymysql/pool.py +184 -0
- sqlspec/adapters/spanner/__init__.py +33 -0
- sqlspec/adapters/spanner/_typing.py +102 -0
- sqlspec/adapters/spanner/adk/__init__.py +5 -0
- sqlspec/adapters/spanner/adk/store.py +758 -0
- sqlspec/adapters/spanner/config.py +355 -0
- sqlspec/adapters/spanner/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/spanner/core.py +263 -0
- sqlspec/adapters/spanner/data_dictionary.py +120 -0
- sqlspec/adapters/spanner/driver.py +407 -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.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/spanner/type_converter.py +342 -0
- sqlspec/adapters/sqlite/__init__.py +19 -0
- sqlspec/adapters/sqlite/_typing.py +123 -0
- sqlspec/adapters/sqlite/adk/__init__.py +5 -0
- sqlspec/adapters/sqlite/adk/store.py +992 -0
- sqlspec/adapters/sqlite/config.py +240 -0
- sqlspec/adapters/sqlite/core.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/sqlite/core.py +357 -0
- sqlspec/adapters/sqlite/data_dictionary.py +198 -0
- sqlspec/adapters/sqlite/driver.py +527 -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.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/sqlite/pool.py +237 -0
- sqlspec/adapters/sqlite/type_converter.cp314-win_amd64.pyd +0 -0
- sqlspec/adapters/sqlite/type_converter.py +114 -0
- sqlspec/base.py +832 -0
- sqlspec/builder/__init__.py +181 -0
- sqlspec/builder/_base.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_base.py +1071 -0
- sqlspec/builder/_column.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_column.py +521 -0
- sqlspec/builder/_ddl.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_ddl.py +1691 -0
- sqlspec/builder/_delete.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_delete.py +95 -0
- sqlspec/builder/_dml.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_dml.py +386 -0
- sqlspec/builder/_explain.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_explain.py +579 -0
- sqlspec/builder/_expression_wrappers.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_expression_wrappers.py +46 -0
- sqlspec/builder/_factory.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_factory.py +1884 -0
- sqlspec/builder/_insert.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_insert.py +405 -0
- sqlspec/builder/_join.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_join.py +489 -0
- sqlspec/builder/_merge.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_merge.py +823 -0
- sqlspec/builder/_parsing_utils.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_parsing_utils.py +295 -0
- sqlspec/builder/_select.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_select.py +1666 -0
- sqlspec/builder/_temporal.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_temporal.py +167 -0
- sqlspec/builder/_update.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_update.py +173 -0
- sqlspec/builder/_vector_distance.cp314-win_amd64.pyd +0 -0
- sqlspec/builder/_vector_distance.py +330 -0
- sqlspec/cli.py +1095 -0
- sqlspec/config.py +2383 -0
- sqlspec/core/__init__.py +372 -0
- sqlspec/core/_correlation.cp314-win_amd64.pyd +0 -0
- sqlspec/core/_correlation.py +176 -0
- sqlspec/core/_pagination.py +42 -0
- sqlspec/core/_pool.cp314-win_amd64.pyd +0 -0
- sqlspec/core/_pool.py +76 -0
- sqlspec/core/cache.cp314-win_amd64.pyd +0 -0
- sqlspec/core/cache.py +1085 -0
- sqlspec/core/compiler.cp314-win_amd64.pyd +0 -0
- sqlspec/core/compiler.py +1090 -0
- sqlspec/core/config_runtime.cp314-win_amd64.pyd +0 -0
- sqlspec/core/config_runtime.py +174 -0
- sqlspec/core/explain.cp314-win_amd64.pyd +0 -0
- sqlspec/core/explain.py +275 -0
- sqlspec/core/filters.cp314-win_amd64.pyd +0 -0
- sqlspec/core/filters.py +969 -0
- sqlspec/core/hashing.cp314-win_amd64.pyd +0 -0
- sqlspec/core/hashing.py +266 -0
- sqlspec/core/metrics.cp314-win_amd64.pyd +0 -0
- sqlspec/core/metrics.py +83 -0
- sqlspec/core/parameters/__init__.py +72 -0
- sqlspec/core/parameters/_alignment.cp314-win_amd64.pyd +0 -0
- sqlspec/core/parameters/_alignment.py +283 -0
- sqlspec/core/parameters/_converter.cp314-win_amd64.pyd +0 -0
- sqlspec/core/parameters/_converter.py +554 -0
- sqlspec/core/parameters/_processor.cp314-win_amd64.pyd +0 -0
- sqlspec/core/parameters/_processor.py +1182 -0
- sqlspec/core/parameters/_registry.cp314-win_amd64.pyd +0 -0
- sqlspec/core/parameters/_registry.py +206 -0
- sqlspec/core/parameters/_transformers.cp314-win_amd64.pyd +0 -0
- sqlspec/core/parameters/_transformers.py +324 -0
- sqlspec/core/parameters/_types.cp314-win_amd64.pyd +0 -0
- sqlspec/core/parameters/_types.py +536 -0
- sqlspec/core/parameters/_validator.cp314-win_amd64.pyd +0 -0
- sqlspec/core/parameters/_validator.py +171 -0
- sqlspec/core/pipeline.cp314-win_amd64.pyd +0 -0
- sqlspec/core/pipeline.py +333 -0
- sqlspec/core/query_modifiers.cp314-win_amd64.pyd +0 -0
- sqlspec/core/query_modifiers.py +508 -0
- sqlspec/core/result/__init__.py +25 -0
- sqlspec/core/result/_base.cp314-win_amd64.pyd +0 -0
- sqlspec/core/result/_base.py +1232 -0
- sqlspec/core/result/_io.cp314-win_amd64.pyd +0 -0
- sqlspec/core/result/_io.py +28 -0
- sqlspec/core/splitter.cp314-win_amd64.pyd +0 -0
- sqlspec/core/splitter.py +1021 -0
- sqlspec/core/sqlcommenter.cp314-win_amd64.pyd +0 -0
- sqlspec/core/sqlcommenter.py +249 -0
- sqlspec/core/stack.cp314-win_amd64.pyd +0 -0
- sqlspec/core/stack.py +163 -0
- sqlspec/core/statement.cp314-win_amd64.pyd +0 -0
- sqlspec/core/statement.py +1865 -0
- sqlspec/core/type_converter.cp314-win_amd64.pyd +0 -0
- sqlspec/core/type_converter.py +340 -0
- sqlspec/data_dictionary/__init__.py +22 -0
- sqlspec/data_dictionary/_loader.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/_loader.py +138 -0
- sqlspec/data_dictionary/_registry.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/_registry.py +74 -0
- sqlspec/data_dictionary/_types.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/_types.py +121 -0
- sqlspec/data_dictionary/dialects/__init__.py +21 -0
- sqlspec/data_dictionary/dialects/bigquery.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/bigquery.py +81 -0
- sqlspec/data_dictionary/dialects/cockroachdb.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/cockroachdb.py +54 -0
- sqlspec/data_dictionary/dialects/duckdb.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/duckdb.py +47 -0
- sqlspec/data_dictionary/dialects/mysql.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/mysql.py +53 -0
- sqlspec/data_dictionary/dialects/oracle.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/oracle.py +197 -0
- sqlspec/data_dictionary/dialects/postgres.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/postgres.py +69 -0
- sqlspec/data_dictionary/dialects/spanner.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/spanner.py +37 -0
- sqlspec/data_dictionary/dialects/sqlite.cp314-win_amd64.pyd +0 -0
- sqlspec/data_dictionary/dialects/sqlite.py +59 -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/dialects/__init__.py +22 -0
- sqlspec/dialects/_compat.cp314-win_amd64.pyd +0 -0
- sqlspec/dialects/_compat.py +14 -0
- sqlspec/dialects/postgres/__init__.py +9 -0
- sqlspec/dialects/postgres/_generators.cp314-win_amd64.pyd +0 -0
- sqlspec/dialects/postgres/_generators.py +57 -0
- sqlspec/dialects/postgres/_operators.cp314-win_amd64.pyd +0 -0
- sqlspec/dialects/postgres/_operators.py +81 -0
- sqlspec/dialects/postgres/_paradedb.py +50 -0
- sqlspec/dialects/postgres/_pgvector.py +36 -0
- sqlspec/dialects/spanner/__init__.py +6 -0
- sqlspec/dialects/spanner/_generators.cp314-win_amd64.pyd +0 -0
- sqlspec/dialects/spanner/_generators.py +206 -0
- sqlspec/dialects/spanner/_spangres.py +77 -0
- sqlspec/dialects/spanner/_spanner.py +179 -0
- sqlspec/driver/__init__.py +49 -0
- sqlspec/driver/_async.cp314-win_amd64.pyd +0 -0
- sqlspec/driver/_async.py +1830 -0
- sqlspec/driver/_common.cp314-win_amd64.pyd +0 -0
- sqlspec/driver/_common.py +2292 -0
- sqlspec/driver/_exception_handler.cp314-win_amd64.pyd +0 -0
- sqlspec/driver/_exception_handler.py +108 -0
- sqlspec/driver/_query_cache.cp314-win_amd64.pyd +0 -0
- sqlspec/driver/_query_cache.py +96 -0
- sqlspec/driver/_sql_helpers.cp314-win_amd64.pyd +0 -0
- sqlspec/driver/_sql_helpers.py +139 -0
- sqlspec/driver/_storage_helpers.cp314-win_amd64.pyd +0 -0
- sqlspec/driver/_storage_helpers.py +153 -0
- sqlspec/driver/_sync.cp314-win_amd64.pyd +0 -0
- sqlspec/driver/_sync.py +1817 -0
- sqlspec/exceptions.cp314-win_amd64.pyd +0 -0
- sqlspec/exceptions.py +480 -0
- sqlspec/extensions/__init__.py +0 -0
- sqlspec/extensions/adk/__init__.py +84 -0
- sqlspec/extensions/adk/_config_utils.py +199 -0
- sqlspec/extensions/adk/_types.cp314-win_amd64.pyd +0 -0
- sqlspec/extensions/adk/_types.py +41 -0
- sqlspec/extensions/adk/artifact/__init__.py +57 -0
- sqlspec/extensions/adk/artifact/_types.cp314-win_amd64.pyd +0 -0
- sqlspec/extensions/adk/artifact/_types.py +32 -0
- sqlspec/extensions/adk/artifact/service.py +508 -0
- sqlspec/extensions/adk/artifact/store.py +361 -0
- sqlspec/extensions/adk/converters.py +212 -0
- sqlspec/extensions/adk/memory/__init__.py +69 -0
- sqlspec/extensions/adk/memory/_types.cp314-win_amd64.pyd +0 -0
- sqlspec/extensions/adk/memory/_types.py +30 -0
- sqlspec/extensions/adk/memory/converters.py +225 -0
- sqlspec/extensions/adk/memory/service.py +316 -0
- sqlspec/extensions/adk/memory/store.py +525 -0
- sqlspec/extensions/adk/migrations/0001_create_adk_tables.py +184 -0
- sqlspec/extensions/adk/migrations/__init__.py +0 -0
- sqlspec/extensions/adk/service.py +279 -0
- sqlspec/extensions/adk/store.py +590 -0
- sqlspec/extensions/events/__init__.py +51 -0
- sqlspec/extensions/events/_channel.py +703 -0
- sqlspec/extensions/events/_hints.cp314-win_amd64.pyd +0 -0
- sqlspec/extensions/events/_hints.py +45 -0
- sqlspec/extensions/events/_models.py +23 -0
- sqlspec/extensions/events/_payload.cp314-win_amd64.pyd +0 -0
- sqlspec/extensions/events/_payload.py +69 -0
- sqlspec/extensions/events/_protocols.py +134 -0
- sqlspec/extensions/events/_queue.py +462 -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 +22 -0
- sqlspec/extensions/fastapi/extension.py +391 -0
- sqlspec/extensions/fastapi/providers.cp314-win_amd64.pyd +0 -0
- sqlspec/extensions/fastapi/providers.py +712 -0
- sqlspec/extensions/flask/__init__.py +38 -0
- sqlspec/extensions/flask/_state.py +87 -0
- sqlspec/extensions/flask/_utils.py +71 -0
- sqlspec/extensions/flask/extension.py +539 -0
- sqlspec/extensions/litestar/__init__.py +31 -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 +1066 -0
- sqlspec/extensions/litestar/providers.cp314-win_amd64.pyd +0 -0
- sqlspec/extensions/litestar/providers.py +784 -0
- sqlspec/extensions/litestar/store.py +298 -0
- sqlspec/extensions/otel/__init__.py +58 -0
- sqlspec/extensions/prometheus/__init__.py +113 -0
- sqlspec/extensions/sanic/__init__.py +19 -0
- sqlspec/extensions/sanic/_state.py +43 -0
- sqlspec/extensions/sanic/_utils.py +127 -0
- sqlspec/extensions/sanic/extension.py +647 -0
- sqlspec/extensions/starlette/__init__.py +22 -0
- sqlspec/extensions/starlette/_state.py +42 -0
- sqlspec/extensions/starlette/_utils.py +96 -0
- sqlspec/extensions/starlette/extension.py +374 -0
- sqlspec/extensions/starlette/middleware.py +281 -0
- sqlspec/loader.cp314-win_amd64.pyd +0 -0
- sqlspec/loader.py +727 -0
- sqlspec/migrations/__init__.py +39 -0
- sqlspec/migrations/base.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/base.py +862 -0
- sqlspec/migrations/commands.py +2151 -0
- sqlspec/migrations/context.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/context.py +157 -0
- sqlspec/migrations/fix.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/fix.py +204 -0
- sqlspec/migrations/loaders.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/loaders.py +443 -0
- sqlspec/migrations/runner.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/runner.py +1195 -0
- sqlspec/migrations/squash.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/squash.py +490 -0
- sqlspec/migrations/templates.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/templates.py +234 -0
- sqlspec/migrations/tracker.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/tracker.py +792 -0
- sqlspec/migrations/utils.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/utils.py +256 -0
- sqlspec/migrations/validation.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/validation.py +359 -0
- sqlspec/migrations/version.cp314-win_amd64.pyd +0 -0
- sqlspec/migrations/version.py +446 -0
- sqlspec/observability/__init__.py +57 -0
- sqlspec/observability/_common.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_common.py +77 -0
- sqlspec/observability/_config.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_config.py +364 -0
- sqlspec/observability/_diagnostics.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_diagnostics.py +74 -0
- sqlspec/observability/_dispatcher.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_dispatcher.py +200 -0
- sqlspec/observability/_formatters/__init__.py +13 -0
- sqlspec/observability/_formatters/_aws.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_formatters/_aws.py +102 -0
- sqlspec/observability/_formatters/_azure.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_formatters/_azure.py +96 -0
- sqlspec/observability/_formatters/_base.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_formatters/_base.py +57 -0
- sqlspec/observability/_formatters/_gcp.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_formatters/_gcp.py +131 -0
- sqlspec/observability/_formatting.py +58 -0
- sqlspec/observability/_observer.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_observer.py +361 -0
- sqlspec/observability/_runtime.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_runtime.py +461 -0
- sqlspec/observability/_sampling.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_sampling.py +188 -0
- sqlspec/observability/_spans.cp314-win_amd64.pyd +0 -0
- sqlspec/observability/_spans.py +161 -0
- sqlspec/protocols.py +955 -0
- sqlspec/py.typed +0 -0
- sqlspec/service.py +433 -0
- sqlspec/storage/__init__.py +48 -0
- sqlspec/storage/_arrow_payload.py +68 -0
- sqlspec/storage/_paths.cp314-win_amd64.pyd +0 -0
- sqlspec/storage/_paths.py +58 -0
- sqlspec/storage/_utils.py +46 -0
- sqlspec/storage/backends/__init__.py +1 -0
- sqlspec/storage/backends/base.cp314-win_amd64.pyd +0 -0
- sqlspec/storage/backends/base.py +374 -0
- sqlspec/storage/backends/fsspec.py +574 -0
- sqlspec/storage/backends/local.py +468 -0
- sqlspec/storage/backends/obstore.py +956 -0
- sqlspec/storage/errors.cp314-win_amd64.pyd +0 -0
- sqlspec/storage/errors.py +102 -0
- sqlspec/storage/pipeline.cp314-win_amd64.pyd +0 -0
- sqlspec/storage/pipeline.py +628 -0
- sqlspec/storage/registry.cp314-win_amd64.pyd +0 -0
- sqlspec/storage/registry.py +329 -0
- sqlspec/typing.py +405 -0
- sqlspec/utils/__init__.py +7 -0
- sqlspec/utils/arrow_helpers.py +384 -0
- sqlspec/utils/config_tools.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/config_tools.py +314 -0
- sqlspec/utils/correlation.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/correlation.py +134 -0
- sqlspec/utils/deprecation.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/deprecation.py +157 -0
- sqlspec/utils/dispatch.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/dispatch.py +101 -0
- sqlspec/utils/fixtures.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/fixtures.py +260 -0
- sqlspec/utils/logging.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/logging.py +251 -0
- sqlspec/utils/module_loader.py +306 -0
- sqlspec/utils/portal.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/portal.py +377 -0
- sqlspec/utils/schema.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/schema.py +1040 -0
- sqlspec/utils/serializers/__init__.py +30 -0
- sqlspec/utils/serializers/_json.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/serializers/_json.py +415 -0
- sqlspec/utils/serializers/_numpy.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/serializers/_numpy.py +65 -0
- sqlspec/utils/serializers/_schema.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/serializers/_schema.py +285 -0
- sqlspec/utils/singleton.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/singleton.py +41 -0
- sqlspec/utils/sync_tools.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/sync_tools.py +316 -0
- sqlspec/utils/text.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/text.py +109 -0
- sqlspec/utils/type_converters.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/type_converters.py +216 -0
- sqlspec/utils/type_guards.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/type_guards.py +1508 -0
- sqlspec/utils/uuids.cp314-win_amd64.pyd +0 -0
- sqlspec/utils/uuids.py +241 -0
- sqlspec-0.47.0.dist-info/METADATA +202 -0
- sqlspec-0.47.0.dist-info/RECORD +621 -0
- sqlspec-0.47.0.dist-info/WHEEL +4 -0
- sqlspec-0.47.0.dist-info/entry_points.txt +6 -0
- sqlspec-0.47.0.dist-info/licenses/LICENSE +21 -0
|
@@ -0,0 +1,1039 @@
|
|
|
1
|
+
"""CockroachDB ADK store for Google Agent Development Kit session/event storage (psycopg)."""
|
|
2
|
+
|
|
3
|
+
from typing import TYPE_CHECKING, Any, cast
|
|
4
|
+
|
|
5
|
+
from psycopg import errors
|
|
6
|
+
from psycopg import sql as pg_sql
|
|
7
|
+
from psycopg.types.json import Jsonb
|
|
8
|
+
|
|
9
|
+
from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
|
|
10
|
+
from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore
|
|
11
|
+
from sqlspec.utils.logging import get_logger
|
|
12
|
+
from sqlspec.utils.sync_tools import async_, run_
|
|
13
|
+
|
|
14
|
+
if TYPE_CHECKING:
|
|
15
|
+
from datetime import datetime
|
|
16
|
+
|
|
17
|
+
from sqlspec.adapters.cockroach_psycopg.config import CockroachPsycopgAsyncConfig, CockroachPsycopgSyncConfig
|
|
18
|
+
from sqlspec.extensions.adk import MemoryRecord
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
__all__ = (
|
|
22
|
+
"CockroachPsycopgAsyncADKMemoryStore",
|
|
23
|
+
"CockroachPsycopgAsyncADKStore",
|
|
24
|
+
"CockroachPsycopgSyncADKMemoryStore",
|
|
25
|
+
"CockroachPsycopgSyncADKStore",
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
logger = get_logger("sqlspec.adapters.cockroach_psycopg.adk.store")
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def _build_insert_params(entry: "MemoryRecord") -> "tuple[object, ...]":
|
|
32
|
+
return (
|
|
33
|
+
entry["id"],
|
|
34
|
+
entry["session_id"],
|
|
35
|
+
entry["app_name"],
|
|
36
|
+
entry["user_id"],
|
|
37
|
+
entry["event_id"],
|
|
38
|
+
entry["author"],
|
|
39
|
+
entry["timestamp"],
|
|
40
|
+
Jsonb(entry["content_json"]),
|
|
41
|
+
entry["content_text"],
|
|
42
|
+
Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None,
|
|
43
|
+
entry["inserted_at"],
|
|
44
|
+
)
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def _build_insert_params_with_owner(entry: "MemoryRecord", owner_id: "object | None") -> "tuple[object, ...]":
|
|
48
|
+
return (
|
|
49
|
+
entry["id"],
|
|
50
|
+
entry["session_id"],
|
|
51
|
+
entry["app_name"],
|
|
52
|
+
entry["user_id"],
|
|
53
|
+
entry["event_id"],
|
|
54
|
+
entry["author"],
|
|
55
|
+
owner_id,
|
|
56
|
+
entry["timestamp"],
|
|
57
|
+
Jsonb(entry["content_json"]),
|
|
58
|
+
entry["content_text"],
|
|
59
|
+
Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None,
|
|
60
|
+
entry["inserted_at"],
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
class CockroachPsycopgAsyncADKStore(BaseAsyncADKStore["CockroachPsycopgAsyncConfig"]):
|
|
65
|
+
"""CockroachDB ADK store using psycopg async driver.
|
|
66
|
+
|
|
67
|
+
Implements session and event storage for Google Agent Development Kit
|
|
68
|
+
using CockroachDB via psycopg in PostgreSQL compatibility mode.
|
|
69
|
+
Events are stored as a single JSONB blob (``event_json``) alongside
|
|
70
|
+
indexed scalar columns for efficient querying.
|
|
71
|
+
|
|
72
|
+
CockroachDB-specific differences from native PostgreSQL:
|
|
73
|
+
- No FILLFACTOR (CockroachDB uses different storage engine)
|
|
74
|
+
- SQL strings require ``.encode()`` for cockroach-psycopg driver
|
|
75
|
+
- GIN/Inverted indexes on JSONB are fully supported (v23.1+)
|
|
76
|
+
- Native tsvector/tsquery FTS with GIN is supported (v23.1+)
|
|
77
|
+
"""
|
|
78
|
+
|
|
79
|
+
__slots__ = ()
|
|
80
|
+
|
|
81
|
+
def __init__(self, config: "CockroachPsycopgAsyncConfig") -> None:
|
|
82
|
+
super().__init__(config)
|
|
83
|
+
|
|
84
|
+
async def _get_create_sessions_table_sql(self) -> str:
|
|
85
|
+
owner_id_line = ""
|
|
86
|
+
if self._owner_id_column_ddl:
|
|
87
|
+
owner_id_line = f",\n {self._owner_id_column_ddl}"
|
|
88
|
+
|
|
89
|
+
return f"""
|
|
90
|
+
CREATE TABLE IF NOT EXISTS {self._session_table} (
|
|
91
|
+
id VARCHAR(128) PRIMARY KEY,
|
|
92
|
+
app_name VARCHAR(128) NOT NULL,
|
|
93
|
+
user_id VARCHAR(128) NOT NULL{owner_id_line},
|
|
94
|
+
state JSONB NOT NULL DEFAULT '{{}}'::jsonb,
|
|
95
|
+
create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
96
|
+
update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
97
|
+
);
|
|
98
|
+
|
|
99
|
+
CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user
|
|
100
|
+
ON {self._session_table}(app_name, user_id);
|
|
101
|
+
|
|
102
|
+
CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time
|
|
103
|
+
ON {self._session_table}(update_time DESC);
|
|
104
|
+
|
|
105
|
+
CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state
|
|
106
|
+
ON {self._session_table} USING GIN (state)
|
|
107
|
+
WHERE state != '{{}}'::jsonb;
|
|
108
|
+
"""
|
|
109
|
+
|
|
110
|
+
async def _get_create_events_table_sql(self) -> str:
|
|
111
|
+
return f"""
|
|
112
|
+
CREATE TABLE IF NOT EXISTS {self._events_table} (
|
|
113
|
+
session_id VARCHAR(128) NOT NULL,
|
|
114
|
+
invocation_id VARCHAR(256) NOT NULL,
|
|
115
|
+
author VARCHAR(256) NOT NULL,
|
|
116
|
+
timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
117
|
+
event_json JSONB NOT NULL,
|
|
118
|
+
FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
|
|
119
|
+
);
|
|
120
|
+
|
|
121
|
+
CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session
|
|
122
|
+
ON {self._events_table}(session_id, timestamp ASC);
|
|
123
|
+
|
|
124
|
+
CREATE INDEX IF NOT EXISTS idx_{self._events_table}_event_json
|
|
125
|
+
ON {self._events_table} USING GIN (event_json);
|
|
126
|
+
"""
|
|
127
|
+
|
|
128
|
+
def _get_drop_tables_sql(self) -> "list[str]":
|
|
129
|
+
return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
|
|
130
|
+
|
|
131
|
+
async def create_tables(self) -> None:
|
|
132
|
+
async with self._config.provide_session() as driver:
|
|
133
|
+
await driver.execute_script(await self._get_create_sessions_table_sql())
|
|
134
|
+
await driver.execute_script(await self._get_create_events_table_sql())
|
|
135
|
+
|
|
136
|
+
async def create_session(
|
|
137
|
+
self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
|
|
138
|
+
) -> SessionRecord:
|
|
139
|
+
state_json = Jsonb(state)
|
|
140
|
+
|
|
141
|
+
params: tuple[Any, ...]
|
|
142
|
+
if self._owner_id_column_name:
|
|
143
|
+
sql = f"""
|
|
144
|
+
INSERT INTO {self._session_table} (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
|
|
145
|
+
VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
|
146
|
+
"""
|
|
147
|
+
params = (session_id, app_name, user_id, owner_id, state_json)
|
|
148
|
+
else:
|
|
149
|
+
sql = f"""
|
|
150
|
+
INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
|
|
151
|
+
VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
|
152
|
+
"""
|
|
153
|
+
params = (session_id, app_name, user_id, state_json)
|
|
154
|
+
|
|
155
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
156
|
+
await cur.execute(sql.encode(), params)
|
|
157
|
+
await conn.commit()
|
|
158
|
+
|
|
159
|
+
result = await self.get_session(session_id)
|
|
160
|
+
if result is None:
|
|
161
|
+
msg = "Session creation failed"
|
|
162
|
+
raise RuntimeError(msg)
|
|
163
|
+
return result
|
|
164
|
+
|
|
165
|
+
async def get_session(self, session_id: str) -> "SessionRecord | None":
|
|
166
|
+
sql = f"""
|
|
167
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
168
|
+
FROM {self._session_table}
|
|
169
|
+
WHERE id = %s
|
|
170
|
+
"""
|
|
171
|
+
|
|
172
|
+
try:
|
|
173
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
174
|
+
await cur.execute(sql.encode(), (session_id,))
|
|
175
|
+
row = await cur.fetchone()
|
|
176
|
+
|
|
177
|
+
if row is None:
|
|
178
|
+
return None
|
|
179
|
+
|
|
180
|
+
return SessionRecord(
|
|
181
|
+
id=row["id"],
|
|
182
|
+
app_name=row["app_name"],
|
|
183
|
+
user_id=row["user_id"],
|
|
184
|
+
state=row["state"],
|
|
185
|
+
create_time=row["create_time"],
|
|
186
|
+
update_time=row["update_time"],
|
|
187
|
+
)
|
|
188
|
+
except errors.UndefinedTable:
|
|
189
|
+
return None
|
|
190
|
+
|
|
191
|
+
async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
|
|
192
|
+
sql = f"""
|
|
193
|
+
UPDATE {self._session_table}
|
|
194
|
+
SET state = %s, update_time = CURRENT_TIMESTAMP
|
|
195
|
+
WHERE id = %s
|
|
196
|
+
"""
|
|
197
|
+
|
|
198
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
199
|
+
await cur.execute(sql.encode(), (Jsonb(state), session_id))
|
|
200
|
+
await conn.commit()
|
|
201
|
+
|
|
202
|
+
async def delete_session(self, session_id: str) -> None:
|
|
203
|
+
sql = f"DELETE FROM {self._session_table} WHERE id = %s"
|
|
204
|
+
|
|
205
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
206
|
+
await cur.execute(sql.encode(), (session_id,))
|
|
207
|
+
await conn.commit()
|
|
208
|
+
|
|
209
|
+
async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
|
|
210
|
+
if user_id is None:
|
|
211
|
+
sql = f"""
|
|
212
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
213
|
+
FROM {self._session_table}
|
|
214
|
+
WHERE app_name = %s
|
|
215
|
+
ORDER BY update_time DESC
|
|
216
|
+
"""
|
|
217
|
+
params: tuple[str, ...] = (app_name,)
|
|
218
|
+
else:
|
|
219
|
+
sql = f"""
|
|
220
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
221
|
+
FROM {self._session_table}
|
|
222
|
+
WHERE app_name = %s AND user_id = %s
|
|
223
|
+
ORDER BY update_time DESC
|
|
224
|
+
"""
|
|
225
|
+
params = (app_name, user_id)
|
|
226
|
+
|
|
227
|
+
try:
|
|
228
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
229
|
+
await cur.execute(sql.encode(), params)
|
|
230
|
+
rows = await cur.fetchall()
|
|
231
|
+
|
|
232
|
+
return [
|
|
233
|
+
SessionRecord(
|
|
234
|
+
id=row["id"],
|
|
235
|
+
app_name=row["app_name"],
|
|
236
|
+
user_id=row["user_id"],
|
|
237
|
+
state=row["state"],
|
|
238
|
+
create_time=row["create_time"],
|
|
239
|
+
update_time=row["update_time"],
|
|
240
|
+
)
|
|
241
|
+
for row in rows
|
|
242
|
+
]
|
|
243
|
+
except errors.UndefinedTable:
|
|
244
|
+
return []
|
|
245
|
+
|
|
246
|
+
async def append_event(self, event_record: EventRecord) -> None:
|
|
247
|
+
sql = f"""
|
|
248
|
+
INSERT INTO {self._events_table} (
|
|
249
|
+
session_id, invocation_id, author, timestamp, event_json
|
|
250
|
+
) VALUES (%s, %s, %s, %s, %s)
|
|
251
|
+
"""
|
|
252
|
+
|
|
253
|
+
event_json_value = event_record["event_json"]
|
|
254
|
+
jsonb_value = Jsonb(event_json_value) if isinstance(event_json_value, dict) else event_json_value
|
|
255
|
+
|
|
256
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
257
|
+
await cur.execute(
|
|
258
|
+
sql.encode(),
|
|
259
|
+
(
|
|
260
|
+
event_record["session_id"],
|
|
261
|
+
event_record["invocation_id"],
|
|
262
|
+
event_record["author"],
|
|
263
|
+
event_record["timestamp"],
|
|
264
|
+
jsonb_value,
|
|
265
|
+
),
|
|
266
|
+
)
|
|
267
|
+
await conn.commit()
|
|
268
|
+
|
|
269
|
+
async def append_event_and_update_state(
|
|
270
|
+
self, event_record: EventRecord, session_id: str, state: "dict[str, Any]"
|
|
271
|
+
) -> SessionRecord:
|
|
272
|
+
insert_sql = f"""
|
|
273
|
+
INSERT INTO {self._events_table} (
|
|
274
|
+
session_id, invocation_id, author, timestamp, event_json
|
|
275
|
+
) VALUES (%s, %s, %s, %s, %s)
|
|
276
|
+
"""
|
|
277
|
+
update_sql = f"""
|
|
278
|
+
UPDATE {self._session_table}
|
|
279
|
+
SET state = %s, update_time = CURRENT_TIMESTAMP
|
|
280
|
+
WHERE id = %s
|
|
281
|
+
RETURNING id, app_name, user_id, state, create_time, update_time
|
|
282
|
+
"""
|
|
283
|
+
|
|
284
|
+
event_json_value = event_record["event_json"]
|
|
285
|
+
jsonb_value = Jsonb(event_json_value) if isinstance(event_json_value, dict) else event_json_value
|
|
286
|
+
|
|
287
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
288
|
+
await cur.execute(
|
|
289
|
+
insert_sql.encode(),
|
|
290
|
+
(
|
|
291
|
+
event_record["session_id"],
|
|
292
|
+
event_record["invocation_id"],
|
|
293
|
+
event_record["author"],
|
|
294
|
+
event_record["timestamp"],
|
|
295
|
+
jsonb_value,
|
|
296
|
+
),
|
|
297
|
+
)
|
|
298
|
+
await cur.execute(update_sql.encode(), (Jsonb(state), session_id))
|
|
299
|
+
row = await cur.fetchone()
|
|
300
|
+
await conn.commit()
|
|
301
|
+
|
|
302
|
+
if row is None:
|
|
303
|
+
msg = f"Session {session_id} not found during append_event_and_update_state."
|
|
304
|
+
raise ValueError(msg)
|
|
305
|
+
|
|
306
|
+
return SessionRecord(
|
|
307
|
+
id=row["id"],
|
|
308
|
+
app_name=row["app_name"],
|
|
309
|
+
user_id=row["user_id"],
|
|
310
|
+
state=row["state"],
|
|
311
|
+
create_time=row["create_time"],
|
|
312
|
+
update_time=row["update_time"],
|
|
313
|
+
)
|
|
314
|
+
|
|
315
|
+
async def get_events(
|
|
316
|
+
self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
|
|
317
|
+
) -> "list[EventRecord]":
|
|
318
|
+
where_clauses = ["session_id = %s"]
|
|
319
|
+
params: list[Any] = [session_id]
|
|
320
|
+
|
|
321
|
+
if after_timestamp is not None:
|
|
322
|
+
where_clauses.append("timestamp > %s")
|
|
323
|
+
params.append(after_timestamp)
|
|
324
|
+
|
|
325
|
+
where_clause = " AND ".join(where_clauses)
|
|
326
|
+
limit_clause = " LIMIT %s" if limit else ""
|
|
327
|
+
if limit:
|
|
328
|
+
params.append(limit)
|
|
329
|
+
|
|
330
|
+
sql = f"""
|
|
331
|
+
SELECT session_id, invocation_id, author, timestamp, event_json
|
|
332
|
+
FROM {self._events_table}
|
|
333
|
+
WHERE {where_clause}
|
|
334
|
+
ORDER BY timestamp ASC{limit_clause}
|
|
335
|
+
"""
|
|
336
|
+
|
|
337
|
+
try:
|
|
338
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
339
|
+
await cur.execute(sql.encode(), tuple(params))
|
|
340
|
+
rows = await cur.fetchall()
|
|
341
|
+
|
|
342
|
+
return [
|
|
343
|
+
EventRecord(
|
|
344
|
+
session_id=row["session_id"],
|
|
345
|
+
invocation_id=row["invocation_id"],
|
|
346
|
+
author=row["author"],
|
|
347
|
+
timestamp=row["timestamp"],
|
|
348
|
+
event_json=row["event_json"],
|
|
349
|
+
)
|
|
350
|
+
for row in rows
|
|
351
|
+
]
|
|
352
|
+
except errors.UndefinedTable:
|
|
353
|
+
return []
|
|
354
|
+
|
|
355
|
+
|
|
356
|
+
class CockroachPsycopgSyncADKStore(BaseAsyncADKStore["CockroachPsycopgSyncConfig"]):
|
|
357
|
+
"""CockroachDB ADK store using psycopg sync driver.
|
|
358
|
+
|
|
359
|
+
Implements session and event storage for Google Agent Development Kit
|
|
360
|
+
using CockroachDB via psycopg in PostgreSQL compatibility mode (sync).
|
|
361
|
+
Events are stored as a single JSONB blob (``event_json``) alongside
|
|
362
|
+
indexed scalar columns for efficient querying.
|
|
363
|
+
|
|
364
|
+
CockroachDB-specific differences from native PostgreSQL:
|
|
365
|
+
- No FILLFACTOR (CockroachDB uses different storage engine)
|
|
366
|
+
- SQL strings require ``.encode()`` for cockroach-psycopg driver
|
|
367
|
+
- GIN/Inverted indexes on JSONB are fully supported (v23.1+)
|
|
368
|
+
"""
|
|
369
|
+
|
|
370
|
+
__slots__ = ()
|
|
371
|
+
|
|
372
|
+
def __init__(self, config: "CockroachPsycopgSyncConfig") -> None:
|
|
373
|
+
super().__init__(config)
|
|
374
|
+
|
|
375
|
+
async def _get_create_sessions_table_sql(self) -> str:
|
|
376
|
+
owner_id_line = ""
|
|
377
|
+
if self._owner_id_column_ddl:
|
|
378
|
+
owner_id_line = f",\n {self._owner_id_column_ddl}"
|
|
379
|
+
|
|
380
|
+
return f"""
|
|
381
|
+
CREATE TABLE IF NOT EXISTS {self._session_table} (
|
|
382
|
+
id VARCHAR(128) PRIMARY KEY,
|
|
383
|
+
app_name VARCHAR(128) NOT NULL,
|
|
384
|
+
user_id VARCHAR(128) NOT NULL{owner_id_line},
|
|
385
|
+
state JSONB NOT NULL DEFAULT '{{}}'::jsonb,
|
|
386
|
+
create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
387
|
+
update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
388
|
+
);
|
|
389
|
+
|
|
390
|
+
CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user
|
|
391
|
+
ON {self._session_table}(app_name, user_id);
|
|
392
|
+
|
|
393
|
+
CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time
|
|
394
|
+
ON {self._session_table}(update_time DESC);
|
|
395
|
+
|
|
396
|
+
CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state
|
|
397
|
+
ON {self._session_table} USING GIN (state)
|
|
398
|
+
WHERE state != '{{}}'::jsonb;
|
|
399
|
+
"""
|
|
400
|
+
|
|
401
|
+
async def _get_create_events_table_sql(self) -> str:
|
|
402
|
+
return f"""
|
|
403
|
+
CREATE TABLE IF NOT EXISTS {self._events_table} (
|
|
404
|
+
session_id VARCHAR(128) NOT NULL,
|
|
405
|
+
invocation_id VARCHAR(256) NOT NULL,
|
|
406
|
+
author VARCHAR(256) NOT NULL,
|
|
407
|
+
timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
408
|
+
event_json JSONB NOT NULL,
|
|
409
|
+
FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
|
|
410
|
+
);
|
|
411
|
+
|
|
412
|
+
CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session
|
|
413
|
+
ON {self._events_table}(session_id, timestamp ASC);
|
|
414
|
+
|
|
415
|
+
CREATE INDEX IF NOT EXISTS idx_{self._events_table}_event_json
|
|
416
|
+
ON {self._events_table} USING GIN (event_json);
|
|
417
|
+
"""
|
|
418
|
+
|
|
419
|
+
def _get_drop_tables_sql(self) -> "list[str]":
|
|
420
|
+
return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
|
|
421
|
+
|
|
422
|
+
def _create_tables(self) -> None:
|
|
423
|
+
with self._config.provide_session() as driver:
|
|
424
|
+
driver.execute_script(run_(self._get_create_sessions_table_sql)())
|
|
425
|
+
driver.execute_script(run_(self._get_create_events_table_sql)())
|
|
426
|
+
|
|
427
|
+
async def create_tables(self) -> None:
|
|
428
|
+
"""Create tables if they don't exist."""
|
|
429
|
+
await async_(self._create_tables)()
|
|
430
|
+
|
|
431
|
+
def _create_session(
|
|
432
|
+
self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
|
|
433
|
+
) -> SessionRecord:
|
|
434
|
+
state_json = Jsonb(state)
|
|
435
|
+
|
|
436
|
+
params: tuple[Any, ...]
|
|
437
|
+
if self._owner_id_column_name:
|
|
438
|
+
sql = f"""
|
|
439
|
+
INSERT INTO {self._session_table} (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
|
|
440
|
+
VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
|
441
|
+
"""
|
|
442
|
+
params = (session_id, app_name, user_id, owner_id, state_json)
|
|
443
|
+
else:
|
|
444
|
+
sql = f"""
|
|
445
|
+
INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
|
|
446
|
+
VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
|
447
|
+
"""
|
|
448
|
+
params = (session_id, app_name, user_id, state_json)
|
|
449
|
+
|
|
450
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
451
|
+
cur.execute(sql.encode(), params)
|
|
452
|
+
conn.commit()
|
|
453
|
+
|
|
454
|
+
result = self._get_session(session_id)
|
|
455
|
+
if result is None:
|
|
456
|
+
msg = "Session creation failed"
|
|
457
|
+
raise RuntimeError(msg)
|
|
458
|
+
return result
|
|
459
|
+
|
|
460
|
+
async def create_session(
|
|
461
|
+
self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
|
|
462
|
+
) -> SessionRecord:
|
|
463
|
+
"""Create a new session."""
|
|
464
|
+
return await async_(self._create_session)(session_id, app_name, user_id, state, owner_id)
|
|
465
|
+
|
|
466
|
+
def _get_session(self, session_id: str) -> "SessionRecord | None":
|
|
467
|
+
sql = f"""
|
|
468
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
469
|
+
FROM {self._session_table}
|
|
470
|
+
WHERE id = %s
|
|
471
|
+
"""
|
|
472
|
+
|
|
473
|
+
try:
|
|
474
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
475
|
+
cur.execute(sql.encode(), (session_id,))
|
|
476
|
+
row = cur.fetchone()
|
|
477
|
+
|
|
478
|
+
if row is None:
|
|
479
|
+
return None
|
|
480
|
+
|
|
481
|
+
return SessionRecord(
|
|
482
|
+
id=row["id"],
|
|
483
|
+
app_name=row["app_name"],
|
|
484
|
+
user_id=row["user_id"],
|
|
485
|
+
state=row["state"],
|
|
486
|
+
create_time=row["create_time"],
|
|
487
|
+
update_time=row["update_time"],
|
|
488
|
+
)
|
|
489
|
+
except errors.UndefinedTable:
|
|
490
|
+
return None
|
|
491
|
+
|
|
492
|
+
async def get_session(self, session_id: str) -> "SessionRecord | None":
|
|
493
|
+
"""Get session by ID."""
|
|
494
|
+
return await async_(self._get_session)(session_id)
|
|
495
|
+
|
|
496
|
+
def _update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
|
|
497
|
+
sql = f"""
|
|
498
|
+
UPDATE {self._session_table}
|
|
499
|
+
SET state = %s, update_time = CURRENT_TIMESTAMP
|
|
500
|
+
WHERE id = %s
|
|
501
|
+
"""
|
|
502
|
+
|
|
503
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
504
|
+
cur.execute(sql.encode(), (Jsonb(state), session_id))
|
|
505
|
+
conn.commit()
|
|
506
|
+
|
|
507
|
+
async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
|
|
508
|
+
"""Update session state."""
|
|
509
|
+
await async_(self._update_session_state)(session_id, state)
|
|
510
|
+
|
|
511
|
+
def _delete_session(self, session_id: str) -> None:
|
|
512
|
+
sql = f"DELETE FROM {self._session_table} WHERE id = %s"
|
|
513
|
+
|
|
514
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
515
|
+
cur.execute(sql.encode(), (session_id,))
|
|
516
|
+
conn.commit()
|
|
517
|
+
|
|
518
|
+
async def delete_session(self, session_id: str) -> None:
|
|
519
|
+
"""Delete session and associated events."""
|
|
520
|
+
await async_(self._delete_session)(session_id)
|
|
521
|
+
|
|
522
|
+
def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
|
|
523
|
+
if user_id is None:
|
|
524
|
+
sql = f"""
|
|
525
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
526
|
+
FROM {self._session_table}
|
|
527
|
+
WHERE app_name = %s
|
|
528
|
+
ORDER BY update_time DESC
|
|
529
|
+
"""
|
|
530
|
+
params: tuple[str, ...] = (app_name,)
|
|
531
|
+
else:
|
|
532
|
+
sql = f"""
|
|
533
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
534
|
+
FROM {self._session_table}
|
|
535
|
+
WHERE app_name = %s AND user_id = %s
|
|
536
|
+
ORDER BY update_time DESC
|
|
537
|
+
"""
|
|
538
|
+
params = (app_name, user_id)
|
|
539
|
+
|
|
540
|
+
try:
|
|
541
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
542
|
+
cur.execute(sql.encode(), params)
|
|
543
|
+
rows = cur.fetchall()
|
|
544
|
+
|
|
545
|
+
return [
|
|
546
|
+
SessionRecord(
|
|
547
|
+
id=row["id"],
|
|
548
|
+
app_name=row["app_name"],
|
|
549
|
+
user_id=row["user_id"],
|
|
550
|
+
state=row["state"],
|
|
551
|
+
create_time=row["create_time"],
|
|
552
|
+
update_time=row["update_time"],
|
|
553
|
+
)
|
|
554
|
+
for row in rows
|
|
555
|
+
]
|
|
556
|
+
except errors.UndefinedTable:
|
|
557
|
+
return []
|
|
558
|
+
|
|
559
|
+
async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
|
|
560
|
+
"""List sessions for an app."""
|
|
561
|
+
return await async_(self._list_sessions)(app_name, user_id)
|
|
562
|
+
|
|
563
|
+
def _append_event_and_update_state(
|
|
564
|
+
self, event_record: EventRecord, session_id: str, state: "dict[str, Any]"
|
|
565
|
+
) -> SessionRecord:
|
|
566
|
+
insert_sql = f"""
|
|
567
|
+
INSERT INTO {self._events_table} (
|
|
568
|
+
session_id, invocation_id, author, timestamp, event_json
|
|
569
|
+
) VALUES (%s, %s, %s, %s, %s)
|
|
570
|
+
"""
|
|
571
|
+
update_sql = f"""
|
|
572
|
+
UPDATE {self._session_table}
|
|
573
|
+
SET state = %s, update_time = CURRENT_TIMESTAMP
|
|
574
|
+
WHERE id = %s
|
|
575
|
+
RETURNING id, app_name, user_id, state, create_time, update_time
|
|
576
|
+
"""
|
|
577
|
+
|
|
578
|
+
event_json_value = event_record["event_json"]
|
|
579
|
+
jsonb_value = Jsonb(event_json_value) if isinstance(event_json_value, dict) else event_json_value
|
|
580
|
+
|
|
581
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
582
|
+
cur.execute(
|
|
583
|
+
insert_sql.encode(),
|
|
584
|
+
(
|
|
585
|
+
event_record["session_id"],
|
|
586
|
+
event_record["invocation_id"],
|
|
587
|
+
event_record["author"],
|
|
588
|
+
event_record["timestamp"],
|
|
589
|
+
jsonb_value,
|
|
590
|
+
),
|
|
591
|
+
)
|
|
592
|
+
cur.execute(update_sql.encode(), (Jsonb(state), session_id))
|
|
593
|
+
row = cur.fetchone()
|
|
594
|
+
conn.commit()
|
|
595
|
+
|
|
596
|
+
if row is None:
|
|
597
|
+
msg = f"Session {session_id} not found during append_event_and_update_state."
|
|
598
|
+
raise ValueError(msg)
|
|
599
|
+
|
|
600
|
+
return SessionRecord(
|
|
601
|
+
id=row["id"],
|
|
602
|
+
app_name=row["app_name"],
|
|
603
|
+
user_id=row["user_id"],
|
|
604
|
+
state=row["state"],
|
|
605
|
+
create_time=row["create_time"],
|
|
606
|
+
update_time=row["update_time"],
|
|
607
|
+
)
|
|
608
|
+
|
|
609
|
+
async def append_event_and_update_state(
|
|
610
|
+
self, event_record: EventRecord, session_id: str, state: "dict[str, Any]"
|
|
611
|
+
) -> SessionRecord:
|
|
612
|
+
"""Atomically append an event and update the session's durable state."""
|
|
613
|
+
return await async_(self._append_event_and_update_state)(event_record, session_id, state)
|
|
614
|
+
|
|
615
|
+
def _insert_event(self, event_record: EventRecord) -> None:
|
|
616
|
+
sql = f"""
|
|
617
|
+
INSERT INTO {self._events_table} (
|
|
618
|
+
session_id, invocation_id, author, timestamp, event_json
|
|
619
|
+
) VALUES (%s, %s, %s, %s, %s)
|
|
620
|
+
"""
|
|
621
|
+
|
|
622
|
+
event_json_value = event_record["event_json"]
|
|
623
|
+
jsonb_value = Jsonb(event_json_value) if isinstance(event_json_value, dict) else event_json_value
|
|
624
|
+
|
|
625
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
626
|
+
cur.execute(
|
|
627
|
+
sql.encode(),
|
|
628
|
+
(
|
|
629
|
+
event_record["session_id"],
|
|
630
|
+
event_record["invocation_id"],
|
|
631
|
+
event_record["author"],
|
|
632
|
+
event_record["timestamp"],
|
|
633
|
+
jsonb_value,
|
|
634
|
+
),
|
|
635
|
+
)
|
|
636
|
+
conn.commit()
|
|
637
|
+
|
|
638
|
+
def _get_events(
|
|
639
|
+
self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
|
|
640
|
+
) -> "list[EventRecord]":
|
|
641
|
+
where_clauses = ["session_id = %s"]
|
|
642
|
+
params: list[Any] = [session_id]
|
|
643
|
+
|
|
644
|
+
if after_timestamp is not None:
|
|
645
|
+
where_clauses.append("timestamp > %s")
|
|
646
|
+
params.append(after_timestamp)
|
|
647
|
+
|
|
648
|
+
where_clause = " AND ".join(where_clauses)
|
|
649
|
+
limit_clause = " LIMIT %s" if limit else ""
|
|
650
|
+
sql = f"""
|
|
651
|
+
SELECT session_id, invocation_id, author, timestamp, event_json
|
|
652
|
+
FROM {self._events_table}
|
|
653
|
+
WHERE {where_clause}
|
|
654
|
+
ORDER BY timestamp ASC{limit_clause}
|
|
655
|
+
"""
|
|
656
|
+
if limit:
|
|
657
|
+
params.append(limit)
|
|
658
|
+
|
|
659
|
+
try:
|
|
660
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
661
|
+
cur.execute(sql.encode(), tuple(params))
|
|
662
|
+
rows = cur.fetchall()
|
|
663
|
+
|
|
664
|
+
return [
|
|
665
|
+
EventRecord(
|
|
666
|
+
session_id=row["session_id"],
|
|
667
|
+
invocation_id=row["invocation_id"],
|
|
668
|
+
author=row["author"],
|
|
669
|
+
timestamp=row["timestamp"],
|
|
670
|
+
event_json=row["event_json"],
|
|
671
|
+
)
|
|
672
|
+
for row in rows
|
|
673
|
+
]
|
|
674
|
+
except errors.UndefinedTable:
|
|
675
|
+
return []
|
|
676
|
+
|
|
677
|
+
async def get_events(
|
|
678
|
+
self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
|
|
679
|
+
) -> "list[EventRecord]":
|
|
680
|
+
"""Get events for a session."""
|
|
681
|
+
return await async_(self._get_events)(session_id, after_timestamp, limit)
|
|
682
|
+
|
|
683
|
+
def _append_event(self, event_record: EventRecord) -> None:
|
|
684
|
+
"""Synchronous implementation of append_event."""
|
|
685
|
+
self._insert_event(event_record)
|
|
686
|
+
|
|
687
|
+
async def append_event(self, event_record: EventRecord) -> None:
|
|
688
|
+
"""Append an event to a session."""
|
|
689
|
+
await async_(self._append_event)(event_record)
|
|
690
|
+
|
|
691
|
+
|
|
692
|
+
class CockroachPsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["CockroachPsycopgAsyncConfig"]):
|
|
693
|
+
"""CockroachDB ADK memory store using psycopg async driver."""
|
|
694
|
+
|
|
695
|
+
__slots__ = ()
|
|
696
|
+
|
|
697
|
+
def __init__(self, config: "CockroachPsycopgAsyncConfig") -> None:
|
|
698
|
+
super().__init__(config)
|
|
699
|
+
|
|
700
|
+
async def _get_create_memory_table_sql(self) -> str:
|
|
701
|
+
owner_id_line = ""
|
|
702
|
+
if self._owner_id_column_ddl:
|
|
703
|
+
owner_id_line = f",\n {self._owner_id_column_ddl}"
|
|
704
|
+
|
|
705
|
+
fts_index = ""
|
|
706
|
+
if self._use_fts:
|
|
707
|
+
fts_index = f"""
|
|
708
|
+
CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts
|
|
709
|
+
ON {self._memory_table} USING GIN (to_tsvector('english', content_text));
|
|
710
|
+
"""
|
|
711
|
+
|
|
712
|
+
return f"""
|
|
713
|
+
CREATE TABLE IF NOT EXISTS {self._memory_table} (
|
|
714
|
+
id VARCHAR(128) PRIMARY KEY,
|
|
715
|
+
session_id VARCHAR(128) NOT NULL,
|
|
716
|
+
app_name VARCHAR(128) NOT NULL,
|
|
717
|
+
user_id VARCHAR(128) NOT NULL,
|
|
718
|
+
event_id VARCHAR(128) NOT NULL UNIQUE,
|
|
719
|
+
author VARCHAR(256){owner_id_line},
|
|
720
|
+
timestamp TIMESTAMPTZ NOT NULL,
|
|
721
|
+
content_json JSONB NOT NULL,
|
|
722
|
+
content_text TEXT NOT NULL,
|
|
723
|
+
metadata_json JSONB,
|
|
724
|
+
inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
725
|
+
);
|
|
726
|
+
|
|
727
|
+
CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time
|
|
728
|
+
ON {self._memory_table}(app_name, user_id, timestamp DESC);
|
|
729
|
+
|
|
730
|
+
CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session
|
|
731
|
+
ON {self._memory_table}(session_id);
|
|
732
|
+
{fts_index}
|
|
733
|
+
"""
|
|
734
|
+
|
|
735
|
+
def _get_drop_memory_table_sql(self) -> "list[str]":
|
|
736
|
+
return [f"DROP TABLE IF EXISTS {self._memory_table}"]
|
|
737
|
+
|
|
738
|
+
async def create_tables(self) -> None:
|
|
739
|
+
if not self._enabled:
|
|
740
|
+
return
|
|
741
|
+
|
|
742
|
+
async with self._config.provide_session() as driver:
|
|
743
|
+
await driver.execute_script(await self._get_create_memory_table_sql())
|
|
744
|
+
|
|
745
|
+
async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
|
|
746
|
+
if not self._enabled:
|
|
747
|
+
msg = "Memory store is disabled"
|
|
748
|
+
raise RuntimeError(msg)
|
|
749
|
+
|
|
750
|
+
if not entries:
|
|
751
|
+
return 0
|
|
752
|
+
|
|
753
|
+
inserted_count = 0
|
|
754
|
+
if self._owner_id_column_name:
|
|
755
|
+
query = pg_sql.SQL("""
|
|
756
|
+
INSERT INTO {table} (
|
|
757
|
+
id, session_id, app_name, user_id, event_id, author,
|
|
758
|
+
{owner_id_col}, timestamp, content_json, content_text,
|
|
759
|
+
metadata_json, inserted_at
|
|
760
|
+
) VALUES (
|
|
761
|
+
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
|
|
762
|
+
)
|
|
763
|
+
ON CONFLICT (event_id) DO NOTHING
|
|
764
|
+
""").format(
|
|
765
|
+
table=pg_sql.Identifier(self._memory_table), owner_id_col=pg_sql.Identifier(self._owner_id_column_name)
|
|
766
|
+
)
|
|
767
|
+
else:
|
|
768
|
+
query = pg_sql.SQL("""
|
|
769
|
+
INSERT INTO {table} (
|
|
770
|
+
id, session_id, app_name, user_id, event_id, author,
|
|
771
|
+
timestamp, content_json, content_text, metadata_json, inserted_at
|
|
772
|
+
) VALUES (
|
|
773
|
+
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
|
|
774
|
+
)
|
|
775
|
+
ON CONFLICT (event_id) DO NOTHING
|
|
776
|
+
""").format(table=pg_sql.Identifier(self._memory_table))
|
|
777
|
+
|
|
778
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
779
|
+
for entry in entries:
|
|
780
|
+
if self._owner_id_column_name:
|
|
781
|
+
await cur.execute(query, _build_insert_params_with_owner(entry, owner_id))
|
|
782
|
+
else:
|
|
783
|
+
await cur.execute(query, _build_insert_params(entry))
|
|
784
|
+
if cur.rowcount and cur.rowcount > 0:
|
|
785
|
+
inserted_count += cur.rowcount
|
|
786
|
+
await conn.commit()
|
|
787
|
+
|
|
788
|
+
return inserted_count
|
|
789
|
+
|
|
790
|
+
async def search_entries(
|
|
791
|
+
self, query: str, app_name: str, user_id: str, limit: "int | None" = None
|
|
792
|
+
) -> "list[MemoryRecord]":
|
|
793
|
+
if not self._enabled:
|
|
794
|
+
msg = "Memory store is disabled"
|
|
795
|
+
raise RuntimeError(msg)
|
|
796
|
+
|
|
797
|
+
if not query:
|
|
798
|
+
return []
|
|
799
|
+
|
|
800
|
+
effective_limit = limit if limit is not None else self._max_results
|
|
801
|
+
|
|
802
|
+
if self._use_fts:
|
|
803
|
+
sql = f"""
|
|
804
|
+
SELECT * FROM {self._memory_table}
|
|
805
|
+
WHERE app_name = %s AND user_id = %s
|
|
806
|
+
AND to_tsvector('english', content_text) @@ plainto_tsquery('english', %s)
|
|
807
|
+
ORDER BY timestamp DESC
|
|
808
|
+
LIMIT %s
|
|
809
|
+
"""
|
|
810
|
+
else:
|
|
811
|
+
sql = f"""
|
|
812
|
+
SELECT * FROM {self._memory_table}
|
|
813
|
+
WHERE app_name = %s AND user_id = %s AND content_text ILIKE %s
|
|
814
|
+
ORDER BY timestamp DESC
|
|
815
|
+
LIMIT %s
|
|
816
|
+
"""
|
|
817
|
+
|
|
818
|
+
search_param = query if self._use_fts else f"%{query}%"
|
|
819
|
+
params = (app_name, user_id, search_param, effective_limit)
|
|
820
|
+
|
|
821
|
+
try:
|
|
822
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
823
|
+
await cur.execute(sql.encode(), params)
|
|
824
|
+
rows = await cur.fetchall()
|
|
825
|
+
columns = [col[0] for col in cur.description or []]
|
|
826
|
+
except errors.UndefinedTable:
|
|
827
|
+
return []
|
|
828
|
+
|
|
829
|
+
return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows]
|
|
830
|
+
|
|
831
|
+
async def delete_entries_by_session(self, session_id: str) -> int:
|
|
832
|
+
if not self._enabled:
|
|
833
|
+
msg = "Memory store is disabled"
|
|
834
|
+
raise RuntimeError(msg)
|
|
835
|
+
|
|
836
|
+
sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s"
|
|
837
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
838
|
+
await cur.execute(sql.encode(), (session_id,))
|
|
839
|
+
await conn.commit()
|
|
840
|
+
return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0
|
|
841
|
+
|
|
842
|
+
async def delete_entries_older_than(self, days: int) -> int:
|
|
843
|
+
if not self._enabled:
|
|
844
|
+
msg = "Memory store is disabled"
|
|
845
|
+
raise RuntimeError(msg)
|
|
846
|
+
|
|
847
|
+
sql = f"""
|
|
848
|
+
DELETE FROM {self._memory_table}
|
|
849
|
+
WHERE inserted_at < CURRENT_TIMESTAMP - INTERVAL '{days} days'
|
|
850
|
+
"""
|
|
851
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
852
|
+
await cur.execute(sql.encode())
|
|
853
|
+
await conn.commit()
|
|
854
|
+
return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0
|
|
855
|
+
|
|
856
|
+
|
|
857
|
+
class CockroachPsycopgSyncADKMemoryStore(BaseAsyncADKMemoryStore["CockroachPsycopgSyncConfig"]):
|
|
858
|
+
"""CockroachDB ADK memory store using psycopg sync driver."""
|
|
859
|
+
|
|
860
|
+
__slots__ = ()
|
|
861
|
+
|
|
862
|
+
def __init__(self, config: "CockroachPsycopgSyncConfig") -> None:
|
|
863
|
+
super().__init__(config)
|
|
864
|
+
|
|
865
|
+
async def _get_create_memory_table_sql(self) -> str:
|
|
866
|
+
owner_id_line = ""
|
|
867
|
+
if self._owner_id_column_ddl:
|
|
868
|
+
owner_id_line = f",\n {self._owner_id_column_ddl}"
|
|
869
|
+
|
|
870
|
+
fts_index = ""
|
|
871
|
+
if self._use_fts:
|
|
872
|
+
fts_index = f"""
|
|
873
|
+
CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts
|
|
874
|
+
ON {self._memory_table} USING GIN (to_tsvector('english', content_text));
|
|
875
|
+
"""
|
|
876
|
+
|
|
877
|
+
return f"""
|
|
878
|
+
CREATE TABLE IF NOT EXISTS {self._memory_table} (
|
|
879
|
+
id VARCHAR(128) PRIMARY KEY,
|
|
880
|
+
session_id VARCHAR(128) NOT NULL,
|
|
881
|
+
app_name VARCHAR(128) NOT NULL,
|
|
882
|
+
user_id VARCHAR(128) NOT NULL,
|
|
883
|
+
event_id VARCHAR(128) NOT NULL UNIQUE,
|
|
884
|
+
author VARCHAR(256){owner_id_line},
|
|
885
|
+
timestamp TIMESTAMPTZ NOT NULL,
|
|
886
|
+
content_json JSONB NOT NULL,
|
|
887
|
+
content_text TEXT NOT NULL,
|
|
888
|
+
metadata_json JSONB,
|
|
889
|
+
inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
890
|
+
);
|
|
891
|
+
|
|
892
|
+
CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time
|
|
893
|
+
ON {self._memory_table}(app_name, user_id, timestamp DESC);
|
|
894
|
+
|
|
895
|
+
CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session
|
|
896
|
+
ON {self._memory_table}(session_id);
|
|
897
|
+
{fts_index}
|
|
898
|
+
"""
|
|
899
|
+
|
|
900
|
+
def _get_drop_memory_table_sql(self) -> "list[str]":
|
|
901
|
+
return [f"DROP TABLE IF EXISTS {self._memory_table}"]
|
|
902
|
+
|
|
903
|
+
def _create_tables(self) -> None:
|
|
904
|
+
if not self._enabled:
|
|
905
|
+
return
|
|
906
|
+
|
|
907
|
+
with self._config.provide_session() as driver:
|
|
908
|
+
driver.execute_script(run_(self._get_create_memory_table_sql)())
|
|
909
|
+
|
|
910
|
+
async def create_tables(self) -> None:
|
|
911
|
+
"""Create tables if they don't exist."""
|
|
912
|
+
await async_(self._create_tables)()
|
|
913
|
+
|
|
914
|
+
def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
|
|
915
|
+
if not self._enabled:
|
|
916
|
+
msg = "Memory store is disabled"
|
|
917
|
+
raise RuntimeError(msg)
|
|
918
|
+
|
|
919
|
+
if not entries:
|
|
920
|
+
return 0
|
|
921
|
+
|
|
922
|
+
inserted_count = 0
|
|
923
|
+
if self._owner_id_column_name:
|
|
924
|
+
query = pg_sql.SQL("""
|
|
925
|
+
INSERT INTO {table} (
|
|
926
|
+
id, session_id, app_name, user_id, event_id, author,
|
|
927
|
+
{owner_id_col}, timestamp, content_json, content_text,
|
|
928
|
+
metadata_json, inserted_at
|
|
929
|
+
) VALUES (
|
|
930
|
+
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
|
|
931
|
+
)
|
|
932
|
+
ON CONFLICT (event_id) DO NOTHING
|
|
933
|
+
""").format(
|
|
934
|
+
table=pg_sql.Identifier(self._memory_table), owner_id_col=pg_sql.Identifier(self._owner_id_column_name)
|
|
935
|
+
)
|
|
936
|
+
else:
|
|
937
|
+
query = pg_sql.SQL("""
|
|
938
|
+
INSERT INTO {table} (
|
|
939
|
+
id, session_id, app_name, user_id, event_id, author,
|
|
940
|
+
timestamp, content_json, content_text, metadata_json, inserted_at
|
|
941
|
+
) VALUES (
|
|
942
|
+
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
|
|
943
|
+
)
|
|
944
|
+
ON CONFLICT (event_id) DO NOTHING
|
|
945
|
+
""").format(table=pg_sql.Identifier(self._memory_table))
|
|
946
|
+
|
|
947
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
948
|
+
for entry in entries:
|
|
949
|
+
if self._owner_id_column_name:
|
|
950
|
+
cur.execute(query, _build_insert_params_with_owner(entry, owner_id))
|
|
951
|
+
else:
|
|
952
|
+
cur.execute(query, _build_insert_params(entry))
|
|
953
|
+
if cur.rowcount and cur.rowcount > 0:
|
|
954
|
+
inserted_count += cur.rowcount
|
|
955
|
+
return inserted_count
|
|
956
|
+
|
|
957
|
+
async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
|
|
958
|
+
"""Bulk insert memory entries with deduplication."""
|
|
959
|
+
return await async_(self._insert_memory_entries)(entries, owner_id)
|
|
960
|
+
|
|
961
|
+
def _search_entries(
|
|
962
|
+
self, query: str, app_name: str, user_id: str, limit: "int | None" = None
|
|
963
|
+
) -> "list[MemoryRecord]":
|
|
964
|
+
if not self._enabled:
|
|
965
|
+
msg = "Memory store is disabled"
|
|
966
|
+
raise RuntimeError(msg)
|
|
967
|
+
|
|
968
|
+
if not query:
|
|
969
|
+
return []
|
|
970
|
+
|
|
971
|
+
effective_limit = limit if limit is not None else self._max_results
|
|
972
|
+
|
|
973
|
+
if self._use_fts:
|
|
974
|
+
sql = f"""
|
|
975
|
+
SELECT * FROM {self._memory_table}
|
|
976
|
+
WHERE app_name = %s AND user_id = %s
|
|
977
|
+
AND to_tsvector('english', content_text) @@ plainto_tsquery('english', %s)
|
|
978
|
+
ORDER BY timestamp DESC
|
|
979
|
+
LIMIT %s
|
|
980
|
+
"""
|
|
981
|
+
else:
|
|
982
|
+
sql = f"""
|
|
983
|
+
SELECT * FROM {self._memory_table}
|
|
984
|
+
WHERE app_name = %s AND user_id = %s AND content_text ILIKE %s
|
|
985
|
+
ORDER BY timestamp DESC
|
|
986
|
+
LIMIT %s
|
|
987
|
+
"""
|
|
988
|
+
|
|
989
|
+
search_param = query if self._use_fts else f"%{query}%"
|
|
990
|
+
params = (app_name, user_id, search_param, effective_limit)
|
|
991
|
+
|
|
992
|
+
try:
|
|
993
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
994
|
+
cur.execute(sql.encode(), params)
|
|
995
|
+
rows = cur.fetchall()
|
|
996
|
+
columns = [col[0] for col in cur.description or []]
|
|
997
|
+
except errors.UndefinedTable:
|
|
998
|
+
return []
|
|
999
|
+
|
|
1000
|
+
return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows]
|
|
1001
|
+
|
|
1002
|
+
async def search_entries(
|
|
1003
|
+
self, query: str, app_name: str, user_id: str, limit: "int | None" = None
|
|
1004
|
+
) -> "list[MemoryRecord]":
|
|
1005
|
+
"""Search memory entries by text query."""
|
|
1006
|
+
return await async_(self._search_entries)(query, app_name, user_id, limit)
|
|
1007
|
+
|
|
1008
|
+
def _delete_entries_by_session(self, session_id: str) -> int:
|
|
1009
|
+
if not self._enabled:
|
|
1010
|
+
msg = "Memory store is disabled"
|
|
1011
|
+
raise RuntimeError(msg)
|
|
1012
|
+
|
|
1013
|
+
sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s"
|
|
1014
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
1015
|
+
cur.execute(sql.encode(), (session_id,))
|
|
1016
|
+
conn.commit()
|
|
1017
|
+
return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0
|
|
1018
|
+
|
|
1019
|
+
async def delete_entries_by_session(self, session_id: str) -> int:
|
|
1020
|
+
"""Delete all memory entries for a specific session."""
|
|
1021
|
+
return await async_(self._delete_entries_by_session)(session_id)
|
|
1022
|
+
|
|
1023
|
+
def _delete_entries_older_than(self, days: int) -> int:
|
|
1024
|
+
if not self._enabled:
|
|
1025
|
+
msg = "Memory store is disabled"
|
|
1026
|
+
raise RuntimeError(msg)
|
|
1027
|
+
|
|
1028
|
+
sql = f"""
|
|
1029
|
+
DELETE FROM {self._memory_table}
|
|
1030
|
+
WHERE inserted_at < CURRENT_TIMESTAMP - INTERVAL '{days} days'
|
|
1031
|
+
"""
|
|
1032
|
+
with self._config.provide_connection() as conn, conn.cursor() as cur:
|
|
1033
|
+
cur.execute(sql.encode())
|
|
1034
|
+
conn.commit()
|
|
1035
|
+
return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0
|
|
1036
|
+
|
|
1037
|
+
async def delete_entries_older_than(self, days: int) -> int:
|
|
1038
|
+
"""Delete memory entries older than specified days."""
|
|
1039
|
+
return await async_(self._delete_entries_older_than)(days)
|