nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.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.
- migrations/0003_allfields_key.py +1 -35
- migrations/0009_upgrade_relations_and_texts_to_v2.py +4 -2
- migrations/0010_fix_corrupt_indexes.py +10 -10
- migrations/0011_materialize_labelset_ids.py +1 -16
- migrations/0012_rollover_shards.py +5 -10
- migrations/0014_rollover_shards.py +4 -5
- migrations/0015_targeted_rollover.py +5 -10
- migrations/0016_upgrade_to_paragraphs_v2.py +25 -28
- migrations/0017_multiple_writable_shards.py +2 -4
- migrations/0018_purge_orphan_kbslugs.py +5 -7
- migrations/0019_upgrade_to_paragraphs_v3.py +25 -28
- migrations/0020_drain_nodes_from_cluster.py +3 -3
- nucliadb/standalone/tests/unit/test_run.py → migrations/0021_overwrite_vectorsets_key.py +16 -19
- nucliadb/tests/unit/test_openapi.py → migrations/0022_fix_paragraph_deletion_bug.py +16 -11
- migrations/0023_backfill_pg_catalog.py +80 -0
- migrations/0025_assign_models_to_kbs_v2.py +113 -0
- migrations/0026_fix_high_cardinality_content_types.py +61 -0
- migrations/0027_rollover_texts3.py +73 -0
- nucliadb/ingest/fields/date.py → migrations/pg/0001_bootstrap.py +10 -12
- migrations/pg/0002_catalog.py +42 -0
- nucliadb/ingest/tests/unit/test_settings.py → migrations/pg/0003_catalog_kbid_index.py +5 -3
- nucliadb/common/cluster/base.py +30 -16
- nucliadb/common/cluster/discovery/base.py +6 -14
- nucliadb/common/cluster/discovery/k8s.py +9 -19
- nucliadb/common/cluster/discovery/manual.py +1 -3
- nucliadb/common/cluster/discovery/utils.py +1 -3
- nucliadb/common/cluster/grpc_node_dummy.py +3 -11
- nucliadb/common/cluster/index_node.py +10 -19
- nucliadb/common/cluster/manager.py +174 -59
- nucliadb/common/cluster/rebalance.py +27 -29
- nucliadb/common/cluster/rollover.py +353 -194
- nucliadb/common/cluster/settings.py +6 -0
- nucliadb/common/cluster/standalone/grpc_node_binding.py +13 -64
- nucliadb/common/cluster/standalone/index_node.py +4 -11
- nucliadb/common/cluster/standalone/service.py +2 -6
- nucliadb/common/cluster/standalone/utils.py +2 -6
- nucliadb/common/cluster/utils.py +29 -22
- nucliadb/common/constants.py +20 -0
- nucliadb/common/context/__init__.py +3 -0
- nucliadb/common/context/fastapi.py +8 -5
- nucliadb/{tests/knowledgeboxes/__init__.py → common/counters.py} +8 -2
- nucliadb/common/datamanagers/__init__.py +7 -1
- nucliadb/common/datamanagers/atomic.py +22 -4
- nucliadb/common/datamanagers/cluster.py +5 -5
- nucliadb/common/datamanagers/entities.py +6 -16
- nucliadb/common/datamanagers/fields.py +84 -0
- nucliadb/common/datamanagers/kb.py +83 -37
- nucliadb/common/datamanagers/labels.py +26 -56
- nucliadb/common/datamanagers/processing.py +2 -6
- nucliadb/common/datamanagers/resources.py +41 -103
- nucliadb/common/datamanagers/rollover.py +76 -15
- nucliadb/common/datamanagers/synonyms.py +1 -1
- nucliadb/common/datamanagers/utils.py +15 -6
- nucliadb/common/datamanagers/vectorsets.py +110 -0
- nucliadb/common/external_index_providers/base.py +257 -0
- nucliadb/{ingest/tests/unit/orm/test_orm_utils.py → common/external_index_providers/exceptions.py} +9 -8
- nucliadb/common/external_index_providers/manager.py +101 -0
- nucliadb/common/external_index_providers/pinecone.py +933 -0
- nucliadb/common/external_index_providers/settings.py +52 -0
- nucliadb/common/http_clients/auth.py +3 -6
- nucliadb/common/http_clients/processing.py +6 -11
- nucliadb/common/http_clients/utils.py +1 -3
- nucliadb/common/ids.py +240 -0
- nucliadb/common/locking.py +29 -7
- nucliadb/common/maindb/driver.py +11 -35
- nucliadb/common/maindb/exceptions.py +3 -0
- nucliadb/common/maindb/local.py +22 -9
- nucliadb/common/maindb/pg.py +206 -111
- nucliadb/common/maindb/utils.py +11 -42
- nucliadb/common/models_utils/from_proto.py +479 -0
- nucliadb/common/models_utils/to_proto.py +60 -0
- nucliadb/common/nidx.py +260 -0
- nucliadb/export_import/datamanager.py +25 -19
- nucliadb/export_import/exporter.py +5 -11
- nucliadb/export_import/importer.py +5 -7
- nucliadb/export_import/models.py +3 -3
- nucliadb/export_import/tasks.py +4 -4
- nucliadb/export_import/utils.py +25 -37
- nucliadb/health.py +1 -3
- nucliadb/ingest/app.py +15 -11
- nucliadb/ingest/consumer/auditing.py +21 -19
- nucliadb/ingest/consumer/consumer.py +82 -47
- nucliadb/ingest/consumer/materializer.py +5 -12
- nucliadb/ingest/consumer/pull.py +12 -27
- nucliadb/ingest/consumer/service.py +19 -17
- nucliadb/ingest/consumer/shard_creator.py +2 -4
- nucliadb/ingest/consumer/utils.py +1 -3
- nucliadb/ingest/fields/base.py +137 -105
- nucliadb/ingest/fields/conversation.py +18 -5
- nucliadb/ingest/fields/exceptions.py +1 -4
- nucliadb/ingest/fields/file.py +7 -16
- nucliadb/ingest/fields/link.py +5 -10
- nucliadb/ingest/fields/text.py +9 -4
- nucliadb/ingest/orm/brain.py +200 -213
- nucliadb/ingest/orm/broker_message.py +181 -0
- nucliadb/ingest/orm/entities.py +36 -51
- nucliadb/ingest/orm/exceptions.py +12 -0
- nucliadb/ingest/orm/knowledgebox.py +322 -197
- nucliadb/ingest/orm/processor/__init__.py +2 -700
- nucliadb/ingest/orm/processor/auditing.py +4 -23
- nucliadb/ingest/orm/processor/data_augmentation.py +164 -0
- nucliadb/ingest/orm/processor/pgcatalog.py +84 -0
- nucliadb/ingest/orm/processor/processor.py +752 -0
- nucliadb/ingest/orm/processor/sequence_manager.py +1 -1
- nucliadb/ingest/orm/resource.py +249 -403
- nucliadb/ingest/orm/utils.py +4 -4
- nucliadb/ingest/partitions.py +3 -9
- nucliadb/ingest/processing.py +70 -73
- nucliadb/ingest/py.typed +0 -0
- nucliadb/ingest/serialize.py +37 -167
- nucliadb/ingest/service/__init__.py +1 -3
- nucliadb/ingest/service/writer.py +185 -412
- nucliadb/ingest/settings.py +10 -20
- nucliadb/ingest/utils.py +3 -6
- nucliadb/learning_proxy.py +242 -55
- nucliadb/metrics_exporter.py +30 -19
- nucliadb/middleware/__init__.py +1 -3
- nucliadb/migrator/command.py +1 -3
- nucliadb/migrator/datamanager.py +13 -13
- nucliadb/migrator/migrator.py +47 -30
- nucliadb/migrator/utils.py +18 -10
- nucliadb/purge/__init__.py +139 -33
- nucliadb/purge/orphan_shards.py +7 -13
- nucliadb/reader/__init__.py +1 -3
- nucliadb/reader/api/models.py +1 -12
- nucliadb/reader/api/v1/__init__.py +0 -1
- nucliadb/reader/api/v1/download.py +21 -88
- nucliadb/reader/api/v1/export_import.py +1 -1
- nucliadb/reader/api/v1/knowledgebox.py +10 -10
- nucliadb/reader/api/v1/learning_config.py +2 -6
- nucliadb/reader/api/v1/resource.py +62 -88
- nucliadb/reader/api/v1/services.py +64 -83
- nucliadb/reader/app.py +12 -29
- nucliadb/reader/lifecycle.py +18 -4
- nucliadb/reader/py.typed +0 -0
- nucliadb/reader/reader/notifications.py +10 -28
- nucliadb/search/__init__.py +1 -3
- nucliadb/search/api/v1/__init__.py +1 -2
- nucliadb/search/api/v1/ask.py +17 -10
- nucliadb/search/api/v1/catalog.py +184 -0
- nucliadb/search/api/v1/feedback.py +16 -24
- nucliadb/search/api/v1/find.py +36 -36
- nucliadb/search/api/v1/knowledgebox.py +89 -60
- nucliadb/search/api/v1/resource/ask.py +2 -8
- nucliadb/search/api/v1/resource/search.py +49 -70
- nucliadb/search/api/v1/search.py +44 -210
- nucliadb/search/api/v1/suggest.py +39 -54
- nucliadb/search/app.py +12 -32
- nucliadb/search/lifecycle.py +10 -3
- nucliadb/search/predict.py +136 -187
- nucliadb/search/py.typed +0 -0
- nucliadb/search/requesters/utils.py +25 -58
- nucliadb/search/search/cache.py +149 -20
- nucliadb/search/search/chat/ask.py +571 -123
- nucliadb/search/{tests/unit/test_run.py → search/chat/exceptions.py} +14 -14
- nucliadb/search/search/chat/images.py +41 -17
- nucliadb/search/search/chat/prompt.py +817 -266
- nucliadb/search/search/chat/query.py +213 -309
- nucliadb/{tests/migrations/__init__.py → search/search/cut.py} +8 -8
- nucliadb/search/search/fetch.py +43 -36
- nucliadb/search/search/filters.py +9 -15
- nucliadb/search/search/find.py +214 -53
- nucliadb/search/search/find_merge.py +408 -391
- nucliadb/search/search/hydrator.py +191 -0
- nucliadb/search/search/merge.py +187 -223
- nucliadb/search/search/metrics.py +73 -2
- nucliadb/search/search/paragraphs.py +64 -106
- nucliadb/search/search/pgcatalog.py +233 -0
- nucliadb/search/search/predict_proxy.py +1 -1
- nucliadb/search/search/query.py +305 -150
- nucliadb/search/search/query_parser/exceptions.py +22 -0
- nucliadb/search/search/query_parser/models.py +101 -0
- nucliadb/search/search/query_parser/parser.py +183 -0
- nucliadb/search/search/rank_fusion.py +204 -0
- nucliadb/search/search/rerankers.py +270 -0
- nucliadb/search/search/shards.py +3 -32
- nucliadb/search/search/summarize.py +7 -18
- nucliadb/search/search/utils.py +27 -4
- nucliadb/search/settings.py +15 -1
- nucliadb/standalone/api_router.py +4 -10
- nucliadb/standalone/app.py +8 -14
- nucliadb/standalone/auth.py +7 -21
- nucliadb/standalone/config.py +7 -10
- nucliadb/standalone/lifecycle.py +26 -25
- nucliadb/standalone/migrations.py +1 -3
- nucliadb/standalone/purge.py +1 -1
- nucliadb/standalone/py.typed +0 -0
- nucliadb/standalone/run.py +3 -6
- nucliadb/standalone/settings.py +9 -16
- nucliadb/standalone/versions.py +15 -5
- nucliadb/tasks/consumer.py +8 -12
- nucliadb/tasks/producer.py +7 -6
- nucliadb/tests/config.py +53 -0
- nucliadb/train/__init__.py +1 -3
- nucliadb/train/api/utils.py +1 -2
- nucliadb/train/api/v1/shards.py +1 -1
- nucliadb/train/api/v1/trainset.py +2 -4
- nucliadb/train/app.py +10 -31
- nucliadb/train/generator.py +10 -19
- nucliadb/train/generators/field_classifier.py +7 -19
- nucliadb/train/generators/field_streaming.py +156 -0
- nucliadb/train/generators/image_classifier.py +12 -18
- nucliadb/train/generators/paragraph_classifier.py +5 -9
- nucliadb/train/generators/paragraph_streaming.py +6 -9
- nucliadb/train/generators/question_answer_streaming.py +19 -20
- nucliadb/train/generators/sentence_classifier.py +9 -15
- nucliadb/train/generators/token_classifier.py +48 -39
- nucliadb/train/generators/utils.py +14 -18
- nucliadb/train/lifecycle.py +7 -3
- nucliadb/train/nodes.py +23 -32
- nucliadb/train/py.typed +0 -0
- nucliadb/train/servicer.py +13 -21
- nucliadb/train/settings.py +2 -6
- nucliadb/train/types.py +13 -10
- nucliadb/train/upload.py +3 -6
- nucliadb/train/uploader.py +19 -23
- nucliadb/train/utils.py +1 -1
- nucliadb/writer/__init__.py +1 -3
- nucliadb/{ingest/fields/keywordset.py → writer/api/utils.py} +13 -10
- nucliadb/writer/api/v1/export_import.py +67 -14
- nucliadb/writer/api/v1/field.py +16 -269
- nucliadb/writer/api/v1/knowledgebox.py +218 -68
- nucliadb/writer/api/v1/resource.py +68 -88
- nucliadb/writer/api/v1/services.py +51 -70
- nucliadb/writer/api/v1/slug.py +61 -0
- nucliadb/writer/api/v1/transaction.py +67 -0
- nucliadb/writer/api/v1/upload.py +143 -117
- nucliadb/writer/app.py +6 -43
- nucliadb/writer/back_pressure.py +16 -38
- nucliadb/writer/exceptions.py +0 -4
- nucliadb/writer/lifecycle.py +21 -15
- nucliadb/writer/py.typed +0 -0
- nucliadb/writer/resource/audit.py +2 -1
- nucliadb/writer/resource/basic.py +48 -46
- nucliadb/writer/resource/field.py +37 -128
- nucliadb/writer/resource/origin.py +1 -2
- nucliadb/writer/settings.py +6 -2
- nucliadb/writer/tus/__init__.py +17 -15
- nucliadb/writer/tus/azure.py +111 -0
- nucliadb/writer/tus/dm.py +17 -5
- nucliadb/writer/tus/exceptions.py +1 -3
- nucliadb/writer/tus/gcs.py +49 -84
- nucliadb/writer/tus/local.py +21 -37
- nucliadb/writer/tus/s3.py +28 -68
- nucliadb/writer/tus/storage.py +5 -56
- nucliadb/writer/vectorsets.py +125 -0
- nucliadb-6.2.1.post2798.dist-info/METADATA +148 -0
- nucliadb-6.2.1.post2798.dist-info/RECORD +343 -0
- {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/WHEEL +1 -1
- nucliadb/common/maindb/redis.py +0 -194
- nucliadb/common/maindb/tikv.py +0 -433
- nucliadb/ingest/fields/layout.py +0 -58
- nucliadb/ingest/tests/conftest.py +0 -30
- nucliadb/ingest/tests/fixtures.py +0 -764
- nucliadb/ingest/tests/integration/consumer/__init__.py +0 -18
- nucliadb/ingest/tests/integration/consumer/test_auditing.py +0 -78
- nucliadb/ingest/tests/integration/consumer/test_materializer.py +0 -126
- nucliadb/ingest/tests/integration/consumer/test_pull.py +0 -144
- nucliadb/ingest/tests/integration/consumer/test_service.py +0 -81
- nucliadb/ingest/tests/integration/consumer/test_shard_creator.py +0 -68
- nucliadb/ingest/tests/integration/ingest/test_ingest.py +0 -684
- nucliadb/ingest/tests/integration/ingest/test_processing_engine.py +0 -95
- nucliadb/ingest/tests/integration/ingest/test_relations.py +0 -272
- nucliadb/ingest/tests/unit/consumer/__init__.py +0 -18
- nucliadb/ingest/tests/unit/consumer/test_auditing.py +0 -139
- nucliadb/ingest/tests/unit/consumer/test_consumer.py +0 -69
- nucliadb/ingest/tests/unit/consumer/test_pull.py +0 -60
- nucliadb/ingest/tests/unit/consumer/test_shard_creator.py +0 -140
- nucliadb/ingest/tests/unit/consumer/test_utils.py +0 -67
- nucliadb/ingest/tests/unit/orm/__init__.py +0 -19
- nucliadb/ingest/tests/unit/orm/test_brain.py +0 -247
- nucliadb/ingest/tests/unit/orm/test_brain_vectors.py +0 -74
- nucliadb/ingest/tests/unit/orm/test_processor.py +0 -131
- nucliadb/ingest/tests/unit/orm/test_resource.py +0 -331
- nucliadb/ingest/tests/unit/test_cache.py +0 -31
- nucliadb/ingest/tests/unit/test_partitions.py +0 -40
- nucliadb/ingest/tests/unit/test_processing.py +0 -171
- nucliadb/middleware/transaction.py +0 -117
- nucliadb/reader/api/v1/learning_collector.py +0 -63
- nucliadb/reader/tests/__init__.py +0 -19
- nucliadb/reader/tests/conftest.py +0 -31
- nucliadb/reader/tests/fixtures.py +0 -136
- nucliadb/reader/tests/test_list_resources.py +0 -75
- nucliadb/reader/tests/test_reader_file_download.py +0 -273
- nucliadb/reader/tests/test_reader_resource.py +0 -353
- nucliadb/reader/tests/test_reader_resource_field.py +0 -219
- nucliadb/search/api/v1/chat.py +0 -263
- nucliadb/search/api/v1/resource/chat.py +0 -174
- nucliadb/search/tests/__init__.py +0 -19
- nucliadb/search/tests/conftest.py +0 -33
- nucliadb/search/tests/fixtures.py +0 -199
- nucliadb/search/tests/node.py +0 -466
- nucliadb/search/tests/unit/__init__.py +0 -18
- nucliadb/search/tests/unit/api/__init__.py +0 -19
- nucliadb/search/tests/unit/api/v1/__init__.py +0 -19
- nucliadb/search/tests/unit/api/v1/resource/__init__.py +0 -19
- nucliadb/search/tests/unit/api/v1/resource/test_chat.py +0 -98
- nucliadb/search/tests/unit/api/v1/test_ask.py +0 -120
- nucliadb/search/tests/unit/api/v1/test_chat.py +0 -96
- nucliadb/search/tests/unit/api/v1/test_predict_proxy.py +0 -98
- nucliadb/search/tests/unit/api/v1/test_summarize.py +0 -99
- nucliadb/search/tests/unit/search/__init__.py +0 -18
- nucliadb/search/tests/unit/search/requesters/__init__.py +0 -18
- nucliadb/search/tests/unit/search/requesters/test_utils.py +0 -211
- nucliadb/search/tests/unit/search/search/__init__.py +0 -19
- nucliadb/search/tests/unit/search/search/test_shards.py +0 -45
- nucliadb/search/tests/unit/search/search/test_utils.py +0 -82
- nucliadb/search/tests/unit/search/test_chat_prompt.py +0 -270
- nucliadb/search/tests/unit/search/test_fetch.py +0 -108
- nucliadb/search/tests/unit/search/test_filters.py +0 -125
- nucliadb/search/tests/unit/search/test_paragraphs.py +0 -157
- nucliadb/search/tests/unit/search/test_predict_proxy.py +0 -106
- nucliadb/search/tests/unit/search/test_query.py +0 -153
- nucliadb/search/tests/unit/test_app.py +0 -79
- nucliadb/search/tests/unit/test_find_merge.py +0 -112
- nucliadb/search/tests/unit/test_merge.py +0 -34
- nucliadb/search/tests/unit/test_predict.py +0 -525
- nucliadb/standalone/tests/__init__.py +0 -19
- nucliadb/standalone/tests/conftest.py +0 -33
- nucliadb/standalone/tests/fixtures.py +0 -38
- nucliadb/standalone/tests/unit/__init__.py +0 -18
- nucliadb/standalone/tests/unit/test_api_router.py +0 -61
- nucliadb/standalone/tests/unit/test_auth.py +0 -169
- nucliadb/standalone/tests/unit/test_introspect.py +0 -35
- nucliadb/standalone/tests/unit/test_migrations.py +0 -63
- nucliadb/standalone/tests/unit/test_versions.py +0 -68
- nucliadb/tests/benchmarks/__init__.py +0 -19
- nucliadb/tests/benchmarks/test_search.py +0 -99
- nucliadb/tests/conftest.py +0 -32
- nucliadb/tests/fixtures.py +0 -735
- nucliadb/tests/knowledgeboxes/philosophy_books.py +0 -202
- nucliadb/tests/knowledgeboxes/ten_dummy_resources.py +0 -107
- nucliadb/tests/migrations/test_migration_0017.py +0 -76
- nucliadb/tests/migrations/test_migration_0018.py +0 -95
- nucliadb/tests/tikv.py +0 -240
- nucliadb/tests/unit/__init__.py +0 -19
- nucliadb/tests/unit/common/__init__.py +0 -19
- nucliadb/tests/unit/common/cluster/__init__.py +0 -19
- nucliadb/tests/unit/common/cluster/discovery/__init__.py +0 -19
- nucliadb/tests/unit/common/cluster/discovery/test_k8s.py +0 -172
- nucliadb/tests/unit/common/cluster/standalone/__init__.py +0 -18
- nucliadb/tests/unit/common/cluster/standalone/test_service.py +0 -114
- nucliadb/tests/unit/common/cluster/standalone/test_utils.py +0 -61
- nucliadb/tests/unit/common/cluster/test_cluster.py +0 -408
- nucliadb/tests/unit/common/cluster/test_kb_shard_manager.py +0 -173
- nucliadb/tests/unit/common/cluster/test_rebalance.py +0 -38
- nucliadb/tests/unit/common/cluster/test_rollover.py +0 -282
- nucliadb/tests/unit/common/maindb/__init__.py +0 -18
- nucliadb/tests/unit/common/maindb/test_driver.py +0 -127
- nucliadb/tests/unit/common/maindb/test_tikv.py +0 -53
- nucliadb/tests/unit/common/maindb/test_utils.py +0 -92
- nucliadb/tests/unit/common/test_context.py +0 -36
- nucliadb/tests/unit/export_import/__init__.py +0 -19
- nucliadb/tests/unit/export_import/test_datamanager.py +0 -37
- nucliadb/tests/unit/export_import/test_utils.py +0 -301
- nucliadb/tests/unit/migrator/__init__.py +0 -19
- nucliadb/tests/unit/migrator/test_migrator.py +0 -87
- nucliadb/tests/unit/tasks/__init__.py +0 -19
- nucliadb/tests/unit/tasks/conftest.py +0 -42
- nucliadb/tests/unit/tasks/test_consumer.py +0 -92
- nucliadb/tests/unit/tasks/test_producer.py +0 -95
- nucliadb/tests/unit/tasks/test_tasks.py +0 -58
- nucliadb/tests/unit/test_field_ids.py +0 -49
- nucliadb/tests/unit/test_health.py +0 -86
- nucliadb/tests/unit/test_kb_slugs.py +0 -54
- nucliadb/tests/unit/test_learning_proxy.py +0 -252
- nucliadb/tests/unit/test_metrics_exporter.py +0 -77
- nucliadb/tests/unit/test_purge.py +0 -136
- nucliadb/tests/utils/__init__.py +0 -74
- nucliadb/tests/utils/aiohttp_session.py +0 -44
- nucliadb/tests/utils/broker_messages/__init__.py +0 -171
- nucliadb/tests/utils/broker_messages/fields.py +0 -197
- nucliadb/tests/utils/broker_messages/helpers.py +0 -33
- nucliadb/tests/utils/entities.py +0 -78
- nucliadb/train/api/v1/check.py +0 -60
- nucliadb/train/tests/__init__.py +0 -19
- nucliadb/train/tests/conftest.py +0 -29
- nucliadb/train/tests/fixtures.py +0 -342
- nucliadb/train/tests/test_field_classification.py +0 -122
- nucliadb/train/tests/test_get_entities.py +0 -80
- nucliadb/train/tests/test_get_info.py +0 -51
- nucliadb/train/tests/test_get_ontology.py +0 -34
- nucliadb/train/tests/test_get_ontology_count.py +0 -63
- nucliadb/train/tests/test_image_classification.py +0 -221
- nucliadb/train/tests/test_list_fields.py +0 -39
- nucliadb/train/tests/test_list_paragraphs.py +0 -73
- nucliadb/train/tests/test_list_resources.py +0 -39
- nucliadb/train/tests/test_list_sentences.py +0 -71
- nucliadb/train/tests/test_paragraph_classification.py +0 -123
- nucliadb/train/tests/test_paragraph_streaming.py +0 -118
- nucliadb/train/tests/test_question_answer_streaming.py +0 -239
- nucliadb/train/tests/test_sentence_classification.py +0 -143
- nucliadb/train/tests/test_token_classification.py +0 -136
- nucliadb/train/tests/utils.py +0 -101
- nucliadb/writer/layouts/__init__.py +0 -51
- nucliadb/writer/layouts/v1.py +0 -59
- nucliadb/writer/tests/__init__.py +0 -19
- nucliadb/writer/tests/conftest.py +0 -31
- nucliadb/writer/tests/fixtures.py +0 -191
- nucliadb/writer/tests/test_fields.py +0 -475
- nucliadb/writer/tests/test_files.py +0 -740
- nucliadb/writer/tests/test_knowledgebox.py +0 -49
- nucliadb/writer/tests/test_reprocess_file_field.py +0 -133
- nucliadb/writer/tests/test_resources.py +0 -476
- nucliadb/writer/tests/test_service.py +0 -137
- nucliadb/writer/tests/test_tus.py +0 -203
- nucliadb/writer/tests/utils.py +0 -35
- nucliadb/writer/tus/pg.py +0 -125
- nucliadb-4.0.0.post542.dist-info/METADATA +0 -135
- nucliadb-4.0.0.post542.dist-info/RECORD +0 -462
- {nucliadb/ingest/tests → migrations/pg}/__init__.py +0 -0
- /nucliadb/{ingest/tests/integration → common/external_index_providers}/__init__.py +0 -0
- /nucliadb/{ingest/tests/integration/ingest → common/models_utils}/__init__.py +0 -0
- /nucliadb/{ingest/tests/unit → search/search/query_parser}/__init__.py +0 -0
- /nucliadb/{ingest/tests → tests}/vectors.py +0 -0
- {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/entry_points.txt +0 -0
- {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/top_level.txt +0 -0
- {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2798.dist-info}/zip-safe +0 -0
@@ -1,282 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
19
|
-
#
|
20
|
-
from datetime import datetime
|
21
|
-
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
22
|
-
|
23
|
-
import pytest
|
24
|
-
|
25
|
-
from nucliadb.common.cluster import manager, rollover
|
26
|
-
from nucliadb.common.cluster.index_node import IndexNode
|
27
|
-
from nucliadb_protos import writer_pb2
|
28
|
-
|
29
|
-
pytestmark = pytest.mark.asyncio
|
30
|
-
|
31
|
-
|
32
|
-
@pytest.fixture()
|
33
|
-
def available_nodes():
|
34
|
-
nodes = {
|
35
|
-
"0": IndexNode(
|
36
|
-
id="0", address="node-0", available_disk=100, shard_count=0, dummy=True
|
37
|
-
),
|
38
|
-
"1": IndexNode(
|
39
|
-
id="1", address="node-1", available_disk=100, shard_count=0, dummy=True
|
40
|
-
),
|
41
|
-
"2": IndexNode(
|
42
|
-
id="2", address="node-2", available_disk=100, shard_count=0, dummy=True
|
43
|
-
),
|
44
|
-
}
|
45
|
-
with patch.object(manager, "INDEX_NODES", new=nodes):
|
46
|
-
yield nodes
|
47
|
-
|
48
|
-
|
49
|
-
@pytest.fixture()
|
50
|
-
def shards():
|
51
|
-
yield writer_pb2.Shards(
|
52
|
-
shards=[
|
53
|
-
writer_pb2.ShardObject(
|
54
|
-
shard="1",
|
55
|
-
replicas=[
|
56
|
-
writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="1")),
|
57
|
-
writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="2")),
|
58
|
-
],
|
59
|
-
read_only=True,
|
60
|
-
),
|
61
|
-
writer_pb2.ShardObject(
|
62
|
-
shard="2",
|
63
|
-
replicas=[
|
64
|
-
writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="3")),
|
65
|
-
writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="4")),
|
66
|
-
],
|
67
|
-
read_only=False,
|
68
|
-
),
|
69
|
-
],
|
70
|
-
kbid="kbid",
|
71
|
-
actual=1,
|
72
|
-
)
|
73
|
-
|
74
|
-
|
75
|
-
@pytest.fixture()
|
76
|
-
def resource_ids():
|
77
|
-
yield ["1", "2", "3"]
|
78
|
-
|
79
|
-
|
80
|
-
@pytest.fixture()
|
81
|
-
def resources_datamanager(resource_ids):
|
82
|
-
mock = MagicMock()
|
83
|
-
|
84
|
-
async def iterate_resource_ids(kbid):
|
85
|
-
for id in resource_ids:
|
86
|
-
yield id
|
87
|
-
|
88
|
-
mock.iterate_resource_ids = iterate_resource_ids
|
89
|
-
mock.get_resource_shard_id = AsyncMock()
|
90
|
-
mock.get_resource_shard_id.return_value = "1"
|
91
|
-
|
92
|
-
res = MagicMock()
|
93
|
-
|
94
|
-
res.basic.modified.ToDatetime.return_value = datetime.now()
|
95
|
-
|
96
|
-
mock.get_resource = AsyncMock()
|
97
|
-
mock.get_resource.return_value = res
|
98
|
-
|
99
|
-
mock.get_resource_index_message = AsyncMock()
|
100
|
-
metadata = MagicMock()
|
101
|
-
metadata.modified.ToDatetime.return_value = datetime.now()
|
102
|
-
mock.get_resource_index_message.return_value = metadata
|
103
|
-
|
104
|
-
with patch("nucliadb.common.cluster.rollover.datamanagers.resources", mock):
|
105
|
-
yield mock
|
106
|
-
|
107
|
-
|
108
|
-
@pytest.fixture()
|
109
|
-
def cluster_datamanager(resource_ids, shards):
|
110
|
-
mock = MagicMock()
|
111
|
-
mock.get_kb_shards = AsyncMock()
|
112
|
-
mock.get_kb_shards.return_value = shards
|
113
|
-
mock.update_kb_shards = AsyncMock()
|
114
|
-
|
115
|
-
with patch("nucliadb.common.cluster.rollover.datamanagers.cluster", mock):
|
116
|
-
yield mock
|
117
|
-
|
118
|
-
|
119
|
-
@pytest.fixture()
|
120
|
-
def rollover_datamanager(resource_ids, cluster_datamanager):
|
121
|
-
mock = MagicMock()
|
122
|
-
mock.get_kb_rollover_shards = AsyncMock()
|
123
|
-
mock.get_kb_rollover_shards.return_value = None
|
124
|
-
mock.update_kb_rollover_shards = AsyncMock()
|
125
|
-
mock.delete_kb_rollover_shard = AsyncMock()
|
126
|
-
mock.delete_kb_rollover_shards = AsyncMock()
|
127
|
-
mock.get_to_index = AsyncMock(side_effect=["1", None])
|
128
|
-
mock.add_indexed = AsyncMock()
|
129
|
-
mock.remove_to_index = AsyncMock()
|
130
|
-
mock.get_indexed_data = AsyncMock(return_value=("1", 1))
|
131
|
-
mock.remove_indexed = AsyncMock()
|
132
|
-
|
133
|
-
async def _mock_indexed_keys(kbid):
|
134
|
-
yield "1"
|
135
|
-
|
136
|
-
mock.iter_indexed_keys = _mock_indexed_keys
|
137
|
-
|
138
|
-
with (
|
139
|
-
patch("nucliadb.common.cluster.rollover.datamanagers.rollover", mock),
|
140
|
-
patch(
|
141
|
-
"nucliadb.common.cluster.rollover.datamanagers.with_transaction",
|
142
|
-
return_value=AsyncMock(),
|
143
|
-
),
|
144
|
-
patch(
|
145
|
-
"nucliadb.ingest.consumer.shard_creator.locking.distributed_lock",
|
146
|
-
return_value=AsyncMock(),
|
147
|
-
),
|
148
|
-
):
|
149
|
-
yield mock
|
150
|
-
|
151
|
-
|
152
|
-
@pytest.fixture()
|
153
|
-
def app_context(rollover_datamanager, resources_datamanager, available_nodes):
|
154
|
-
mock = MagicMock()
|
155
|
-
mock.shard_manager = MagicMock()
|
156
|
-
mock.shard_manager.rollback_shard = AsyncMock()
|
157
|
-
mock.shard_manager.add_resource = AsyncMock()
|
158
|
-
mock.shard_manager.delete_resource = AsyncMock()
|
159
|
-
mock.kv_driver = MagicMock()
|
160
|
-
|
161
|
-
consumer_info = MagicMock()
|
162
|
-
consumer_info.delivered.stream_seq = 0
|
163
|
-
mock.nats_manager.js.consumer_info.return_value = consumer_info
|
164
|
-
yield mock
|
165
|
-
|
166
|
-
|
167
|
-
async def test_create_rollover_shards(
|
168
|
-
app_context, available_nodes, shards: writer_pb2.Shards, rollover_datamanager
|
169
|
-
):
|
170
|
-
new_shards = await rollover.create_rollover_shards(app_context, "kbid")
|
171
|
-
|
172
|
-
assert new_shards.kbid == "kbid"
|
173
|
-
assert sum(
|
174
|
-
[len(node.writer.calls["NewShard"]) for node in available_nodes.values()]
|
175
|
-
) == sum([len(s.replicas) for s in shards.shards])
|
176
|
-
rollover_datamanager.update_kb_rollover_shards.assert_called_with(
|
177
|
-
ANY, kbid="kbid", kb_shards=new_shards
|
178
|
-
)
|
179
|
-
|
180
|
-
|
181
|
-
async def test_create_rollover_shards_does_not_recreate(
|
182
|
-
app_context, shards: writer_pb2.Shards, rollover_datamanager
|
183
|
-
):
|
184
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
185
|
-
|
186
|
-
await rollover.create_rollover_shards(app_context, "kbid")
|
187
|
-
|
188
|
-
app_context.shard_manager.rollback_shard.assert_not_called()
|
189
|
-
rollover_datamanager.update_kb_rollover_shards.assert_not_called()
|
190
|
-
|
191
|
-
|
192
|
-
async def test_index_rollover_shards(
|
193
|
-
app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
|
194
|
-
):
|
195
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
196
|
-
|
197
|
-
await rollover.index_rollover_shards(app_context, "kbid")
|
198
|
-
rollover_datamanager.add_indexed.assert_called_with(
|
199
|
-
ANY, kbid="kbid", resource_id="1", shard_id="1", modification_time=1
|
200
|
-
)
|
201
|
-
|
202
|
-
|
203
|
-
async def test_index_rollover_shards_handles_missing_shards(
|
204
|
-
app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
|
205
|
-
):
|
206
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = None
|
207
|
-
with pytest.raises(rollover.UnexpectedRolloverError):
|
208
|
-
await rollover.index_rollover_shards(app_context, "kbid")
|
209
|
-
|
210
|
-
|
211
|
-
async def test_index_rollover_shards_handles_missing_shard_id(
|
212
|
-
app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
|
213
|
-
):
|
214
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
215
|
-
resources_datamanager.get_resource_shard_id.return_value = None
|
216
|
-
await rollover.index_rollover_shards(app_context, "kbid")
|
217
|
-
|
218
|
-
|
219
|
-
async def test_index_rollover_shards_handles_missing_res(
|
220
|
-
app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
|
221
|
-
):
|
222
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
223
|
-
resources_datamanager.get_resource_index_message.return_value = None
|
224
|
-
|
225
|
-
await rollover.index_rollover_shards(app_context, "kbid")
|
226
|
-
|
227
|
-
rollover_datamanager.remove_to_index.assert_called_with(
|
228
|
-
ANY, kbid="kbid", resource="1"
|
229
|
-
)
|
230
|
-
|
231
|
-
|
232
|
-
async def test_cutover_shards(
|
233
|
-
app_context, rollover_datamanager, cluster_datamanager, shards
|
234
|
-
):
|
235
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
236
|
-
|
237
|
-
await rollover.cutover_shards(app_context, "kbid")
|
238
|
-
|
239
|
-
cluster_datamanager.update_kb_shards.assert_called_with(
|
240
|
-
ANY, kbid="kbid", shards=ANY
|
241
|
-
)
|
242
|
-
[
|
243
|
-
app_context.shard_manager.rollback_shard.assert_any_call(shard)
|
244
|
-
for shard in shards.shards
|
245
|
-
]
|
246
|
-
|
247
|
-
|
248
|
-
async def test_cutover_shards_missing(app_context, rollover_datamanager):
|
249
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = None
|
250
|
-
|
251
|
-
with pytest.raises(rollover.UnexpectedRolloverError):
|
252
|
-
await rollover.cutover_shards(app_context, "kbid")
|
253
|
-
|
254
|
-
|
255
|
-
async def test_validate_indexed_data(
|
256
|
-
app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
|
257
|
-
):
|
258
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
259
|
-
|
260
|
-
indexed_res = await rollover.validate_indexed_data(app_context, "kbid")
|
261
|
-
assert len(indexed_res) == len(resource_ids)
|
262
|
-
[
|
263
|
-
resources_datamanager.get_resource_index_message.assert_any_call(
|
264
|
-
ANY, kbid="kbid", rid=res_id
|
265
|
-
)
|
266
|
-
for res_id in resource_ids
|
267
|
-
]
|
268
|
-
|
269
|
-
|
270
|
-
async def test_validate_indexed_data_handles_missing_res(
|
271
|
-
app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
|
272
|
-
):
|
273
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
274
|
-
resources_datamanager.get_resource_index_message.return_value = None
|
275
|
-
assert len(await rollover.validate_indexed_data(app_context, "kbid")) == 0
|
276
|
-
|
277
|
-
|
278
|
-
async def test_rollover_shards(app_context, rollover_datamanager, shards, resource_ids):
|
279
|
-
rollover_datamanager.get_kb_rollover_shards.return_value = shards
|
280
|
-
resource_ids.clear()
|
281
|
-
|
282
|
-
await rollover.rollover_kb_shards(app_context, "kbid")
|
@@ -1,18 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
@@ -1,127 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
19
|
-
#
|
20
|
-
import asyncio
|
21
|
-
from unittest import mock
|
22
|
-
|
23
|
-
import pytest
|
24
|
-
|
25
|
-
from nucliadb.common.maindb.driver import Driver
|
26
|
-
|
27
|
-
|
28
|
-
class TransactionTest:
|
29
|
-
def __init__(self):
|
30
|
-
self.abort = mock.AsyncMock()
|
31
|
-
self.open = True
|
32
|
-
|
33
|
-
async def commit(self, **kw):
|
34
|
-
self.open = False
|
35
|
-
|
36
|
-
|
37
|
-
@pytest.fixture(scope="function")
|
38
|
-
def txn():
|
39
|
-
return TransactionTest()
|
40
|
-
|
41
|
-
|
42
|
-
@pytest.fixture(scope="function")
|
43
|
-
def driver(txn) -> Driver: # type: ignore
|
44
|
-
driver = Driver()
|
45
|
-
with mock.patch.object(driver, "begin", new=mock.AsyncMock(return_value=txn)):
|
46
|
-
yield driver
|
47
|
-
|
48
|
-
|
49
|
-
@pytest.mark.asyncio
|
50
|
-
async def test_driver_async_abort(driver, txn):
|
51
|
-
async with driver.transaction(wait_for_abort=False):
|
52
|
-
pass
|
53
|
-
|
54
|
-
assert len(driver._abort_tasks) == 1
|
55
|
-
await asyncio.sleep(0.1)
|
56
|
-
|
57
|
-
txn.abort.assert_called_once()
|
58
|
-
assert len(driver._abort_tasks) == 0
|
59
|
-
|
60
|
-
|
61
|
-
@pytest.mark.asyncio
|
62
|
-
async def test_driver_finalize_aborts_transactions(driver, txn):
|
63
|
-
async with driver.transaction(wait_for_abort=False):
|
64
|
-
pass
|
65
|
-
|
66
|
-
assert len(driver._abort_tasks) == 1
|
67
|
-
|
68
|
-
await driver.finalize()
|
69
|
-
|
70
|
-
txn.abort.assert_called_once()
|
71
|
-
|
72
|
-
assert len(driver._abort_tasks) == 0
|
73
|
-
|
74
|
-
|
75
|
-
@pytest.mark.asyncio
|
76
|
-
async def test_transaction_handles_txn_begin_errors(driver):
|
77
|
-
driver.begin.side_effect = ValueError()
|
78
|
-
testmock = mock.AsyncMock()
|
79
|
-
with pytest.raises(ValueError):
|
80
|
-
async with driver.transaction():
|
81
|
-
testmock()
|
82
|
-
testmock.assert_not_called()
|
83
|
-
|
84
|
-
|
85
|
-
@pytest.mark.asyncio
|
86
|
-
async def test_transaction_does_not_abotr_if_commited(driver):
|
87
|
-
async with driver.transaction() as txn:
|
88
|
-
await txn.commit()
|
89
|
-
txn.abort.assert_not_awaited()
|
90
|
-
|
91
|
-
|
92
|
-
@pytest.mark.asyncio
|
93
|
-
async def test_transaction_aborts_if_txn_open(driver):
|
94
|
-
async with driver.transaction() as txn:
|
95
|
-
pass
|
96
|
-
txn.abort.assert_awaited_once()
|
97
|
-
|
98
|
-
|
99
|
-
@pytest.mark.asyncio
|
100
|
-
async def test_transaction_aborts_on_errors(driver):
|
101
|
-
with pytest.raises(ValueError):
|
102
|
-
async with driver.transaction() as txn:
|
103
|
-
raise ValueError()
|
104
|
-
txn.abort.assert_awaited_once()
|
105
|
-
|
106
|
-
|
107
|
-
@pytest.mark.asyncio
|
108
|
-
async def test_transaction_wait_for_abort(driver):
|
109
|
-
async with driver.transaction(wait_for_abort=False) as txn:
|
110
|
-
pass
|
111
|
-
txn.abort.assert_not_awaited()
|
112
|
-
txn.abort.assert_called_once()
|
113
|
-
txn.abort.reset_mock()
|
114
|
-
|
115
|
-
# Should wait for abort if there are errors
|
116
|
-
with pytest.raises(ValueError):
|
117
|
-
async with driver.transaction(wait_for_abort=False) as txn:
|
118
|
-
raise ValueError()
|
119
|
-
txn.abort.assert_awaited_once()
|
120
|
-
txn.abort.assert_called_once()
|
121
|
-
txn.abort.reset_mock()
|
122
|
-
|
123
|
-
async with driver.transaction(wait_for_abort=True) as txn:
|
124
|
-
pass
|
125
|
-
txn.abort.assert_awaited_once()
|
126
|
-
txn.abort.assert_called_once()
|
127
|
-
txn.abort.reset_mock()
|
@@ -1,53 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
19
|
-
|
20
|
-
from unittest.mock import AsyncMock, MagicMock
|
21
|
-
|
22
|
-
import pytest
|
23
|
-
|
24
|
-
from nucliadb.common.maindb.exceptions import ConflictError
|
25
|
-
from nucliadb.common.maindb.tikv import LeaderNotFoundError, TiKVDataLayer
|
26
|
-
|
27
|
-
|
28
|
-
@pytest.mark.parametrize(
|
29
|
-
"tikv_exception,handled_exception",
|
30
|
-
[
|
31
|
-
(
|
32
|
-
Exception("gRPC error: RpcFailure: 4-DEADLINE_EXCEEDED Deadline Exceeded"),
|
33
|
-
TimeoutError,
|
34
|
-
),
|
35
|
-
(Exception("Leader of region 34234 is not found"), LeaderNotFoundError),
|
36
|
-
],
|
37
|
-
)
|
38
|
-
async def test_get_retrials(tikv_exception, handled_exception):
|
39
|
-
inner_txn = MagicMock(get=AsyncMock(side_effect=tikv_exception))
|
40
|
-
tikv_txn = TiKVDataLayer(inner_txn)
|
41
|
-
|
42
|
-
with pytest.raises(handled_exception):
|
43
|
-
await tikv_txn.get("key")
|
44
|
-
|
45
|
-
assert inner_txn.get.call_count == 2
|
46
|
-
|
47
|
-
|
48
|
-
async def test_commit_raises_conflict_error():
|
49
|
-
inner_txn = MagicMock(commit=AsyncMock(side_effect=Exception("WriteConflict")))
|
50
|
-
tikv_txn = TiKVDataLayer(inner_txn)
|
51
|
-
|
52
|
-
with pytest.raises(ConflictError):
|
53
|
-
await tikv_txn.commit()
|
@@ -1,92 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
19
|
-
|
20
|
-
from unittest.mock import AsyncMock, patch
|
21
|
-
|
22
|
-
import pytest
|
23
|
-
|
24
|
-
from nucliadb.common.maindb.utils import settings, setup_driver
|
25
|
-
from nucliadb.ingest.settings import DriverConfig
|
26
|
-
from nucliadb_utils.exceptions import ConfigurationError
|
27
|
-
from nucliadb_utils.utilities import Utility, clean_utility
|
28
|
-
|
29
|
-
|
30
|
-
@pytest.fixture(autouse=True)
|
31
|
-
def reset_driver_utils():
|
32
|
-
clean_utility(Utility.MAINDB_DRIVER)
|
33
|
-
yield
|
34
|
-
clean_utility(Utility.MAINDB_DRIVER)
|
35
|
-
|
36
|
-
|
37
|
-
@pytest.mark.asyncio
|
38
|
-
async def test_setup_driver_redis():
|
39
|
-
mock = AsyncMock(initialized=False)
|
40
|
-
with (
|
41
|
-
patch.object(settings, "driver", DriverConfig("redis")),
|
42
|
-
patch.object(settings, "driver_redis_url", "driver_redis_url"),
|
43
|
-
patch("nucliadb.common.maindb.utils.RedisDriver", return_value=mock),
|
44
|
-
):
|
45
|
-
assert await setup_driver() == mock
|
46
|
-
mock.initialize.assert_awaited_once()
|
47
|
-
|
48
|
-
|
49
|
-
@pytest.mark.asyncio
|
50
|
-
async def test_setup_driver_tikv():
|
51
|
-
mock = AsyncMock(initialized=False)
|
52
|
-
with (
|
53
|
-
patch.object(settings, "driver", DriverConfig("tikv")),
|
54
|
-
patch.object(settings, "driver_tikv_url", "driver_tikv_url"),
|
55
|
-
patch("nucliadb.common.maindb.utils.TiKVDriver", return_value=mock),
|
56
|
-
):
|
57
|
-
assert await setup_driver() == mock
|
58
|
-
mock.initialize.assert_awaited_once()
|
59
|
-
|
60
|
-
|
61
|
-
@pytest.mark.asyncio
|
62
|
-
async def test_setup_driver_pg():
|
63
|
-
mock = AsyncMock(initialized=False)
|
64
|
-
with (
|
65
|
-
patch.object(settings, "driver", DriverConfig("pg")),
|
66
|
-
patch.object(settings, "driver_pg_url", "driver_pg_url"),
|
67
|
-
patch("nucliadb.common.maindb.utils.PGDriver", return_value=mock),
|
68
|
-
):
|
69
|
-
assert await setup_driver() == mock
|
70
|
-
mock.initialize.assert_awaited_once()
|
71
|
-
|
72
|
-
|
73
|
-
@pytest.mark.asyncio
|
74
|
-
async def test_setup_driver_local():
|
75
|
-
mock = AsyncMock(initialized=False)
|
76
|
-
with (
|
77
|
-
patch.object(settings, "driver", DriverConfig("local")),
|
78
|
-
patch.object(settings, "driver_local_url", "driver_local_url"),
|
79
|
-
patch("nucliadb.common.maindb.utils.LocalDriver", return_value=mock),
|
80
|
-
):
|
81
|
-
assert await setup_driver() == mock
|
82
|
-
mock.initialize.assert_awaited_once()
|
83
|
-
|
84
|
-
|
85
|
-
@pytest.mark.asyncio
|
86
|
-
async def test_setup_driver_error():
|
87
|
-
with (
|
88
|
-
patch.object(settings, "driver", DriverConfig("pg")),
|
89
|
-
patch.object(settings, "driver_pg_url", None),
|
90
|
-
pytest.raises(ConfigurationError),
|
91
|
-
):
|
92
|
-
await setup_driver()
|
@@ -1,36 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
19
|
-
#
|
20
|
-
import asyncio
|
21
|
-
from unittest.mock import AsyncMock
|
22
|
-
|
23
|
-
from nucliadb.common.context import ApplicationContext
|
24
|
-
|
25
|
-
|
26
|
-
async def test_initialize_happens_only_once():
|
27
|
-
context = ApplicationContext()
|
28
|
-
context._initialize = AsyncMock() # noqa
|
29
|
-
|
30
|
-
tasks = []
|
31
|
-
for _ in range(10):
|
32
|
-
tasks.append(context.initialize())
|
33
|
-
await asyncio.gather(*tasks)
|
34
|
-
|
35
|
-
context._initialize.assert_awaited_once()
|
36
|
-
context._initialized is True
|
@@ -1,19 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
19
|
-
#
|
@@ -1,37 +0,0 @@
|
|
1
|
-
# Copyright (C) 2021 Bosutech XXI S.L.
|
2
|
-
#
|
3
|
-
# nucliadb is offered under the AGPL v3.0 and as commercial software.
|
4
|
-
# For commercial licensing, contact us at info@nuclia.com.
|
5
|
-
#
|
6
|
-
# AGPL:
|
7
|
-
# This program is free software: you can redistribute it and/or modify
|
8
|
-
# it under the terms of the GNU Affero General Public License as
|
9
|
-
# published by the Free Software Foundation, either version 3 of the
|
10
|
-
# License, or (at your option) any later version.
|
11
|
-
#
|
12
|
-
# This program is distributed in the hope that it will be useful,
|
13
|
-
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
14
|
-
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
15
|
-
# GNU Affero General Public License for more details.
|
16
|
-
#
|
17
|
-
# You should have received a copy of the GNU Affero General Public License
|
18
|
-
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
19
|
-
#
|
20
|
-
from unittest.mock import AsyncMock, Mock
|
21
|
-
|
22
|
-
from nucliadb.export_import.datamanager import ExportImportDataManager
|
23
|
-
|
24
|
-
|
25
|
-
async def test_try_delete_from_storage():
|
26
|
-
driver = Mock()
|
27
|
-
storage = Mock()
|
28
|
-
|
29
|
-
dm = ExportImportDataManager(driver, storage)
|
30
|
-
dm.delete_export = AsyncMock(side_effect=ValueError()) # type: ignore
|
31
|
-
dm.delete_import = AsyncMock(side_effect=ValueError()) # type: ignore
|
32
|
-
|
33
|
-
await dm.try_delete_from_storage("export", "kbid", "foo")
|
34
|
-
await dm.try_delete_from_storage("import", "kbid", "foo")
|
35
|
-
|
36
|
-
dm.delete_export.assert_called()
|
37
|
-
dm.delete_import.assert_called()
|