nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__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.
Files changed (418) hide show
  1. migrations/0003_allfields_key.py +1 -35
  2. migrations/0009_upgrade_relations_and_texts_to_v2.py +4 -2
  3. migrations/0010_fix_corrupt_indexes.py +10 -10
  4. migrations/0011_materialize_labelset_ids.py +1 -16
  5. migrations/0012_rollover_shards.py +5 -10
  6. migrations/0014_rollover_shards.py +4 -5
  7. migrations/0015_targeted_rollover.py +5 -10
  8. migrations/0016_upgrade_to_paragraphs_v2.py +25 -28
  9. migrations/0017_multiple_writable_shards.py +2 -4
  10. migrations/0018_purge_orphan_kbslugs.py +5 -7
  11. migrations/0019_upgrade_to_paragraphs_v3.py +25 -28
  12. migrations/0020_drain_nodes_from_cluster.py +3 -3
  13. nucliadb/standalone/tests/unit/test_run.py → migrations/0021_overwrite_vectorsets_key.py +16 -19
  14. nucliadb/tests/unit/test_openapi.py → migrations/0022_fix_paragraph_deletion_bug.py +16 -11
  15. migrations/0023_backfill_pg_catalog.py +80 -0
  16. migrations/0025_assign_models_to_kbs_v2.py +113 -0
  17. migrations/0026_fix_high_cardinality_content_types.py +61 -0
  18. migrations/0027_rollover_texts3.py +73 -0
  19. nucliadb/ingest/fields/date.py → migrations/pg/0001_bootstrap.py +10 -12
  20. migrations/pg/0002_catalog.py +42 -0
  21. nucliadb/ingest/tests/unit/test_settings.py → migrations/pg/0003_catalog_kbid_index.py +5 -3
  22. nucliadb/common/cluster/base.py +30 -16
  23. nucliadb/common/cluster/discovery/base.py +6 -14
  24. nucliadb/common/cluster/discovery/k8s.py +9 -19
  25. nucliadb/common/cluster/discovery/manual.py +1 -3
  26. nucliadb/common/cluster/discovery/utils.py +1 -3
  27. nucliadb/common/cluster/grpc_node_dummy.py +3 -11
  28. nucliadb/common/cluster/index_node.py +10 -19
  29. nucliadb/common/cluster/manager.py +174 -59
  30. nucliadb/common/cluster/rebalance.py +27 -29
  31. nucliadb/common/cluster/rollover.py +353 -194
  32. nucliadb/common/cluster/settings.py +6 -0
  33. nucliadb/common/cluster/standalone/grpc_node_binding.py +13 -64
  34. nucliadb/common/cluster/standalone/index_node.py +4 -11
  35. nucliadb/common/cluster/standalone/service.py +2 -6
  36. nucliadb/common/cluster/standalone/utils.py +2 -6
  37. nucliadb/common/cluster/utils.py +29 -22
  38. nucliadb/common/constants.py +20 -0
  39. nucliadb/common/context/__init__.py +3 -0
  40. nucliadb/common/context/fastapi.py +8 -5
  41. nucliadb/{tests/knowledgeboxes/__init__.py → common/counters.py} +8 -2
  42. nucliadb/common/datamanagers/__init__.py +7 -1
  43. nucliadb/common/datamanagers/atomic.py +22 -4
  44. nucliadb/common/datamanagers/cluster.py +5 -5
  45. nucliadb/common/datamanagers/entities.py +6 -16
  46. nucliadb/common/datamanagers/fields.py +84 -0
  47. nucliadb/common/datamanagers/kb.py +83 -37
  48. nucliadb/common/datamanagers/labels.py +26 -56
  49. nucliadb/common/datamanagers/processing.py +2 -6
  50. nucliadb/common/datamanagers/resources.py +41 -103
  51. nucliadb/common/datamanagers/rollover.py +76 -15
  52. nucliadb/common/datamanagers/synonyms.py +1 -1
  53. nucliadb/common/datamanagers/utils.py +15 -6
  54. nucliadb/common/datamanagers/vectorsets.py +110 -0
  55. nucliadb/common/external_index_providers/base.py +257 -0
  56. nucliadb/{ingest/tests/unit/orm/test_orm_utils.py → common/external_index_providers/exceptions.py} +9 -8
  57. nucliadb/common/external_index_providers/manager.py +101 -0
  58. nucliadb/common/external_index_providers/pinecone.py +933 -0
  59. nucliadb/common/external_index_providers/settings.py +52 -0
  60. nucliadb/common/http_clients/auth.py +3 -6
  61. nucliadb/common/http_clients/processing.py +6 -11
  62. nucliadb/common/http_clients/utils.py +1 -3
  63. nucliadb/common/ids.py +240 -0
  64. nucliadb/common/locking.py +29 -7
  65. nucliadb/common/maindb/driver.py +11 -35
  66. nucliadb/common/maindb/exceptions.py +3 -0
  67. nucliadb/common/maindb/local.py +22 -9
  68. nucliadb/common/maindb/pg.py +206 -111
  69. nucliadb/common/maindb/utils.py +11 -42
  70. nucliadb/common/models_utils/from_proto.py +479 -0
  71. nucliadb/common/models_utils/to_proto.py +60 -0
  72. nucliadb/common/nidx.py +260 -0
  73. nucliadb/export_import/datamanager.py +25 -19
  74. nucliadb/export_import/exporter.py +5 -11
  75. nucliadb/export_import/importer.py +5 -7
  76. nucliadb/export_import/models.py +3 -3
  77. nucliadb/export_import/tasks.py +4 -4
  78. nucliadb/export_import/utils.py +25 -37
  79. nucliadb/health.py +1 -3
  80. nucliadb/ingest/app.py +15 -11
  81. nucliadb/ingest/consumer/auditing.py +21 -19
  82. nucliadb/ingest/consumer/consumer.py +82 -47
  83. nucliadb/ingest/consumer/materializer.py +5 -12
  84. nucliadb/ingest/consumer/pull.py +12 -27
  85. nucliadb/ingest/consumer/service.py +19 -17
  86. nucliadb/ingest/consumer/shard_creator.py +2 -4
  87. nucliadb/ingest/consumer/utils.py +1 -3
  88. nucliadb/ingest/fields/base.py +137 -105
  89. nucliadb/ingest/fields/conversation.py +18 -5
  90. nucliadb/ingest/fields/exceptions.py +1 -4
  91. nucliadb/ingest/fields/file.py +7 -16
  92. nucliadb/ingest/fields/link.py +5 -10
  93. nucliadb/ingest/fields/text.py +9 -4
  94. nucliadb/ingest/orm/brain.py +200 -213
  95. nucliadb/ingest/orm/broker_message.py +181 -0
  96. nucliadb/ingest/orm/entities.py +36 -51
  97. nucliadb/ingest/orm/exceptions.py +12 -0
  98. nucliadb/ingest/orm/knowledgebox.py +322 -197
  99. nucliadb/ingest/orm/processor/__init__.py +2 -700
  100. nucliadb/ingest/orm/processor/auditing.py +4 -23
  101. nucliadb/ingest/orm/processor/data_augmentation.py +164 -0
  102. nucliadb/ingest/orm/processor/pgcatalog.py +84 -0
  103. nucliadb/ingest/orm/processor/processor.py +752 -0
  104. nucliadb/ingest/orm/processor/sequence_manager.py +1 -1
  105. nucliadb/ingest/orm/resource.py +249 -402
  106. nucliadb/ingest/orm/utils.py +4 -4
  107. nucliadb/ingest/partitions.py +3 -9
  108. nucliadb/ingest/processing.py +64 -73
  109. nucliadb/ingest/py.typed +0 -0
  110. nucliadb/ingest/serialize.py +37 -167
  111. nucliadb/ingest/service/__init__.py +1 -3
  112. nucliadb/ingest/service/writer.py +185 -412
  113. nucliadb/ingest/settings.py +10 -20
  114. nucliadb/ingest/utils.py +3 -6
  115. nucliadb/learning_proxy.py +242 -55
  116. nucliadb/metrics_exporter.py +30 -19
  117. nucliadb/middleware/__init__.py +1 -3
  118. nucliadb/migrator/command.py +1 -3
  119. nucliadb/migrator/datamanager.py +13 -13
  120. nucliadb/migrator/migrator.py +47 -30
  121. nucliadb/migrator/utils.py +18 -10
  122. nucliadb/purge/__init__.py +139 -33
  123. nucliadb/purge/orphan_shards.py +7 -13
  124. nucliadb/reader/__init__.py +1 -3
  125. nucliadb/reader/api/models.py +1 -12
  126. nucliadb/reader/api/v1/__init__.py +0 -1
  127. nucliadb/reader/api/v1/download.py +21 -88
  128. nucliadb/reader/api/v1/export_import.py +1 -1
  129. nucliadb/reader/api/v1/knowledgebox.py +10 -10
  130. nucliadb/reader/api/v1/learning_config.py +2 -6
  131. nucliadb/reader/api/v1/resource.py +62 -88
  132. nucliadb/reader/api/v1/services.py +64 -83
  133. nucliadb/reader/app.py +12 -29
  134. nucliadb/reader/lifecycle.py +18 -4
  135. nucliadb/reader/py.typed +0 -0
  136. nucliadb/reader/reader/notifications.py +10 -28
  137. nucliadb/search/__init__.py +1 -3
  138. nucliadb/search/api/v1/__init__.py +1 -2
  139. nucliadb/search/api/v1/ask.py +17 -10
  140. nucliadb/search/api/v1/catalog.py +184 -0
  141. nucliadb/search/api/v1/feedback.py +16 -24
  142. nucliadb/search/api/v1/find.py +36 -36
  143. nucliadb/search/api/v1/knowledgebox.py +89 -60
  144. nucliadb/search/api/v1/resource/ask.py +2 -8
  145. nucliadb/search/api/v1/resource/search.py +49 -70
  146. nucliadb/search/api/v1/search.py +44 -210
  147. nucliadb/search/api/v1/suggest.py +39 -54
  148. nucliadb/search/app.py +12 -32
  149. nucliadb/search/lifecycle.py +10 -3
  150. nucliadb/search/predict.py +136 -187
  151. nucliadb/search/py.typed +0 -0
  152. nucliadb/search/requesters/utils.py +25 -58
  153. nucliadb/search/search/cache.py +149 -20
  154. nucliadb/search/search/chat/ask.py +571 -123
  155. nucliadb/search/{tests/unit/test_run.py → search/chat/exceptions.py} +14 -14
  156. nucliadb/search/search/chat/images.py +41 -17
  157. nucliadb/search/search/chat/prompt.py +817 -266
  158. nucliadb/search/search/chat/query.py +213 -309
  159. nucliadb/{tests/migrations/__init__.py → search/search/cut.py} +8 -8
  160. nucliadb/search/search/fetch.py +43 -36
  161. nucliadb/search/search/filters.py +9 -15
  162. nucliadb/search/search/find.py +214 -53
  163. nucliadb/search/search/find_merge.py +408 -391
  164. nucliadb/search/search/hydrator.py +191 -0
  165. nucliadb/search/search/merge.py +187 -223
  166. nucliadb/search/search/metrics.py +73 -2
  167. nucliadb/search/search/paragraphs.py +64 -106
  168. nucliadb/search/search/pgcatalog.py +233 -0
  169. nucliadb/search/search/predict_proxy.py +1 -1
  170. nucliadb/search/search/query.py +305 -150
  171. nucliadb/search/search/query_parser/exceptions.py +22 -0
  172. nucliadb/search/search/query_parser/models.py +101 -0
  173. nucliadb/search/search/query_parser/parser.py +183 -0
  174. nucliadb/search/search/rank_fusion.py +204 -0
  175. nucliadb/search/search/rerankers.py +270 -0
  176. nucliadb/search/search/shards.py +3 -32
  177. nucliadb/search/search/summarize.py +7 -18
  178. nucliadb/search/search/utils.py +27 -4
  179. nucliadb/search/settings.py +15 -1
  180. nucliadb/standalone/api_router.py +4 -10
  181. nucliadb/standalone/app.py +8 -14
  182. nucliadb/standalone/auth.py +7 -21
  183. nucliadb/standalone/config.py +7 -10
  184. nucliadb/standalone/lifecycle.py +26 -25
  185. nucliadb/standalone/migrations.py +1 -3
  186. nucliadb/standalone/purge.py +1 -1
  187. nucliadb/standalone/py.typed +0 -0
  188. nucliadb/standalone/run.py +3 -6
  189. nucliadb/standalone/settings.py +9 -16
  190. nucliadb/standalone/versions.py +15 -5
  191. nucliadb/tasks/consumer.py +8 -12
  192. nucliadb/tasks/producer.py +7 -6
  193. nucliadb/tests/config.py +53 -0
  194. nucliadb/train/__init__.py +1 -3
  195. nucliadb/train/api/utils.py +1 -2
  196. nucliadb/train/api/v1/shards.py +1 -1
  197. nucliadb/train/api/v1/trainset.py +2 -4
  198. nucliadb/train/app.py +10 -31
  199. nucliadb/train/generator.py +10 -19
  200. nucliadb/train/generators/field_classifier.py +7 -19
  201. nucliadb/train/generators/field_streaming.py +156 -0
  202. nucliadb/train/generators/image_classifier.py +12 -18
  203. nucliadb/train/generators/paragraph_classifier.py +5 -9
  204. nucliadb/train/generators/paragraph_streaming.py +6 -9
  205. nucliadb/train/generators/question_answer_streaming.py +19 -20
  206. nucliadb/train/generators/sentence_classifier.py +9 -15
  207. nucliadb/train/generators/token_classifier.py +48 -39
  208. nucliadb/train/generators/utils.py +14 -18
  209. nucliadb/train/lifecycle.py +7 -3
  210. nucliadb/train/nodes.py +23 -32
  211. nucliadb/train/py.typed +0 -0
  212. nucliadb/train/servicer.py +13 -21
  213. nucliadb/train/settings.py +2 -6
  214. nucliadb/train/types.py +13 -10
  215. nucliadb/train/upload.py +3 -6
  216. nucliadb/train/uploader.py +19 -23
  217. nucliadb/train/utils.py +1 -1
  218. nucliadb/writer/__init__.py +1 -3
  219. nucliadb/{ingest/fields/keywordset.py → writer/api/utils.py} +13 -10
  220. nucliadb/writer/api/v1/export_import.py +67 -14
  221. nucliadb/writer/api/v1/field.py +16 -269
  222. nucliadb/writer/api/v1/knowledgebox.py +218 -68
  223. nucliadb/writer/api/v1/resource.py +68 -88
  224. nucliadb/writer/api/v1/services.py +51 -70
  225. nucliadb/writer/api/v1/slug.py +61 -0
  226. nucliadb/writer/api/v1/transaction.py +67 -0
  227. nucliadb/writer/api/v1/upload.py +114 -113
  228. nucliadb/writer/app.py +6 -43
  229. nucliadb/writer/back_pressure.py +16 -38
  230. nucliadb/writer/exceptions.py +0 -4
  231. nucliadb/writer/lifecycle.py +21 -15
  232. nucliadb/writer/py.typed +0 -0
  233. nucliadb/writer/resource/audit.py +2 -1
  234. nucliadb/writer/resource/basic.py +48 -46
  235. nucliadb/writer/resource/field.py +25 -127
  236. nucliadb/writer/resource/origin.py +1 -2
  237. nucliadb/writer/settings.py +6 -2
  238. nucliadb/writer/tus/__init__.py +17 -15
  239. nucliadb/writer/tus/azure.py +111 -0
  240. nucliadb/writer/tus/dm.py +17 -5
  241. nucliadb/writer/tus/exceptions.py +1 -3
  242. nucliadb/writer/tus/gcs.py +49 -84
  243. nucliadb/writer/tus/local.py +21 -37
  244. nucliadb/writer/tus/s3.py +28 -68
  245. nucliadb/writer/tus/storage.py +5 -56
  246. nucliadb/writer/vectorsets.py +125 -0
  247. nucliadb-6.2.1.post2777.dist-info/METADATA +148 -0
  248. nucliadb-6.2.1.post2777.dist-info/RECORD +343 -0
  249. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2777.dist-info}/WHEEL +1 -1
  250. nucliadb/common/maindb/redis.py +0 -194
  251. nucliadb/common/maindb/tikv.py +0 -433
  252. nucliadb/ingest/fields/layout.py +0 -58
  253. nucliadb/ingest/tests/conftest.py +0 -30
  254. nucliadb/ingest/tests/fixtures.py +0 -764
  255. nucliadb/ingest/tests/integration/consumer/__init__.py +0 -18
  256. nucliadb/ingest/tests/integration/consumer/test_auditing.py +0 -78
  257. nucliadb/ingest/tests/integration/consumer/test_materializer.py +0 -126
  258. nucliadb/ingest/tests/integration/consumer/test_pull.py +0 -144
  259. nucliadb/ingest/tests/integration/consumer/test_service.py +0 -81
  260. nucliadb/ingest/tests/integration/consumer/test_shard_creator.py +0 -68
  261. nucliadb/ingest/tests/integration/ingest/test_ingest.py +0 -684
  262. nucliadb/ingest/tests/integration/ingest/test_processing_engine.py +0 -95
  263. nucliadb/ingest/tests/integration/ingest/test_relations.py +0 -272
  264. nucliadb/ingest/tests/unit/consumer/__init__.py +0 -18
  265. nucliadb/ingest/tests/unit/consumer/test_auditing.py +0 -139
  266. nucliadb/ingest/tests/unit/consumer/test_consumer.py +0 -69
  267. nucliadb/ingest/tests/unit/consumer/test_pull.py +0 -60
  268. nucliadb/ingest/tests/unit/consumer/test_shard_creator.py +0 -140
  269. nucliadb/ingest/tests/unit/consumer/test_utils.py +0 -67
  270. nucliadb/ingest/tests/unit/orm/__init__.py +0 -19
  271. nucliadb/ingest/tests/unit/orm/test_brain.py +0 -247
  272. nucliadb/ingest/tests/unit/orm/test_brain_vectors.py +0 -74
  273. nucliadb/ingest/tests/unit/orm/test_processor.py +0 -131
  274. nucliadb/ingest/tests/unit/orm/test_resource.py +0 -331
  275. nucliadb/ingest/tests/unit/test_cache.py +0 -31
  276. nucliadb/ingest/tests/unit/test_partitions.py +0 -40
  277. nucliadb/ingest/tests/unit/test_processing.py +0 -171
  278. nucliadb/middleware/transaction.py +0 -117
  279. nucliadb/reader/api/v1/learning_collector.py +0 -63
  280. nucliadb/reader/tests/__init__.py +0 -19
  281. nucliadb/reader/tests/conftest.py +0 -31
  282. nucliadb/reader/tests/fixtures.py +0 -136
  283. nucliadb/reader/tests/test_list_resources.py +0 -75
  284. nucliadb/reader/tests/test_reader_file_download.py +0 -273
  285. nucliadb/reader/tests/test_reader_resource.py +0 -353
  286. nucliadb/reader/tests/test_reader_resource_field.py +0 -219
  287. nucliadb/search/api/v1/chat.py +0 -263
  288. nucliadb/search/api/v1/resource/chat.py +0 -174
  289. nucliadb/search/tests/__init__.py +0 -19
  290. nucliadb/search/tests/conftest.py +0 -33
  291. nucliadb/search/tests/fixtures.py +0 -199
  292. nucliadb/search/tests/node.py +0 -466
  293. nucliadb/search/tests/unit/__init__.py +0 -18
  294. nucliadb/search/tests/unit/api/__init__.py +0 -19
  295. nucliadb/search/tests/unit/api/v1/__init__.py +0 -19
  296. nucliadb/search/tests/unit/api/v1/resource/__init__.py +0 -19
  297. nucliadb/search/tests/unit/api/v1/resource/test_chat.py +0 -98
  298. nucliadb/search/tests/unit/api/v1/test_ask.py +0 -120
  299. nucliadb/search/tests/unit/api/v1/test_chat.py +0 -96
  300. nucliadb/search/tests/unit/api/v1/test_predict_proxy.py +0 -98
  301. nucliadb/search/tests/unit/api/v1/test_summarize.py +0 -99
  302. nucliadb/search/tests/unit/search/__init__.py +0 -18
  303. nucliadb/search/tests/unit/search/requesters/__init__.py +0 -18
  304. nucliadb/search/tests/unit/search/requesters/test_utils.py +0 -211
  305. nucliadb/search/tests/unit/search/search/__init__.py +0 -19
  306. nucliadb/search/tests/unit/search/search/test_shards.py +0 -45
  307. nucliadb/search/tests/unit/search/search/test_utils.py +0 -82
  308. nucliadb/search/tests/unit/search/test_chat_prompt.py +0 -270
  309. nucliadb/search/tests/unit/search/test_fetch.py +0 -108
  310. nucliadb/search/tests/unit/search/test_filters.py +0 -125
  311. nucliadb/search/tests/unit/search/test_paragraphs.py +0 -157
  312. nucliadb/search/tests/unit/search/test_predict_proxy.py +0 -106
  313. nucliadb/search/tests/unit/search/test_query.py +0 -153
  314. nucliadb/search/tests/unit/test_app.py +0 -79
  315. nucliadb/search/tests/unit/test_find_merge.py +0 -112
  316. nucliadb/search/tests/unit/test_merge.py +0 -34
  317. nucliadb/search/tests/unit/test_predict.py +0 -525
  318. nucliadb/standalone/tests/__init__.py +0 -19
  319. nucliadb/standalone/tests/conftest.py +0 -33
  320. nucliadb/standalone/tests/fixtures.py +0 -38
  321. nucliadb/standalone/tests/unit/__init__.py +0 -18
  322. nucliadb/standalone/tests/unit/test_api_router.py +0 -61
  323. nucliadb/standalone/tests/unit/test_auth.py +0 -169
  324. nucliadb/standalone/tests/unit/test_introspect.py +0 -35
  325. nucliadb/standalone/tests/unit/test_migrations.py +0 -63
  326. nucliadb/standalone/tests/unit/test_versions.py +0 -68
  327. nucliadb/tests/benchmarks/__init__.py +0 -19
  328. nucliadb/tests/benchmarks/test_search.py +0 -99
  329. nucliadb/tests/conftest.py +0 -32
  330. nucliadb/tests/fixtures.py +0 -735
  331. nucliadb/tests/knowledgeboxes/philosophy_books.py +0 -202
  332. nucliadb/tests/knowledgeboxes/ten_dummy_resources.py +0 -107
  333. nucliadb/tests/migrations/test_migration_0017.py +0 -76
  334. nucliadb/tests/migrations/test_migration_0018.py +0 -95
  335. nucliadb/tests/tikv.py +0 -240
  336. nucliadb/tests/unit/__init__.py +0 -19
  337. nucliadb/tests/unit/common/__init__.py +0 -19
  338. nucliadb/tests/unit/common/cluster/__init__.py +0 -19
  339. nucliadb/tests/unit/common/cluster/discovery/__init__.py +0 -19
  340. nucliadb/tests/unit/common/cluster/discovery/test_k8s.py +0 -172
  341. nucliadb/tests/unit/common/cluster/standalone/__init__.py +0 -18
  342. nucliadb/tests/unit/common/cluster/standalone/test_service.py +0 -114
  343. nucliadb/tests/unit/common/cluster/standalone/test_utils.py +0 -61
  344. nucliadb/tests/unit/common/cluster/test_cluster.py +0 -408
  345. nucliadb/tests/unit/common/cluster/test_kb_shard_manager.py +0 -173
  346. nucliadb/tests/unit/common/cluster/test_rebalance.py +0 -38
  347. nucliadb/tests/unit/common/cluster/test_rollover.py +0 -282
  348. nucliadb/tests/unit/common/maindb/__init__.py +0 -18
  349. nucliadb/tests/unit/common/maindb/test_driver.py +0 -127
  350. nucliadb/tests/unit/common/maindb/test_tikv.py +0 -53
  351. nucliadb/tests/unit/common/maindb/test_utils.py +0 -92
  352. nucliadb/tests/unit/common/test_context.py +0 -36
  353. nucliadb/tests/unit/export_import/__init__.py +0 -19
  354. nucliadb/tests/unit/export_import/test_datamanager.py +0 -37
  355. nucliadb/tests/unit/export_import/test_utils.py +0 -301
  356. nucliadb/tests/unit/migrator/__init__.py +0 -19
  357. nucliadb/tests/unit/migrator/test_migrator.py +0 -87
  358. nucliadb/tests/unit/tasks/__init__.py +0 -19
  359. nucliadb/tests/unit/tasks/conftest.py +0 -42
  360. nucliadb/tests/unit/tasks/test_consumer.py +0 -92
  361. nucliadb/tests/unit/tasks/test_producer.py +0 -95
  362. nucliadb/tests/unit/tasks/test_tasks.py +0 -58
  363. nucliadb/tests/unit/test_field_ids.py +0 -49
  364. nucliadb/tests/unit/test_health.py +0 -86
  365. nucliadb/tests/unit/test_kb_slugs.py +0 -54
  366. nucliadb/tests/unit/test_learning_proxy.py +0 -252
  367. nucliadb/tests/unit/test_metrics_exporter.py +0 -77
  368. nucliadb/tests/unit/test_purge.py +0 -136
  369. nucliadb/tests/utils/__init__.py +0 -74
  370. nucliadb/tests/utils/aiohttp_session.py +0 -44
  371. nucliadb/tests/utils/broker_messages/__init__.py +0 -171
  372. nucliadb/tests/utils/broker_messages/fields.py +0 -197
  373. nucliadb/tests/utils/broker_messages/helpers.py +0 -33
  374. nucliadb/tests/utils/entities.py +0 -78
  375. nucliadb/train/api/v1/check.py +0 -60
  376. nucliadb/train/tests/__init__.py +0 -19
  377. nucliadb/train/tests/conftest.py +0 -29
  378. nucliadb/train/tests/fixtures.py +0 -342
  379. nucliadb/train/tests/test_field_classification.py +0 -122
  380. nucliadb/train/tests/test_get_entities.py +0 -80
  381. nucliadb/train/tests/test_get_info.py +0 -51
  382. nucliadb/train/tests/test_get_ontology.py +0 -34
  383. nucliadb/train/tests/test_get_ontology_count.py +0 -63
  384. nucliadb/train/tests/test_image_classification.py +0 -221
  385. nucliadb/train/tests/test_list_fields.py +0 -39
  386. nucliadb/train/tests/test_list_paragraphs.py +0 -73
  387. nucliadb/train/tests/test_list_resources.py +0 -39
  388. nucliadb/train/tests/test_list_sentences.py +0 -71
  389. nucliadb/train/tests/test_paragraph_classification.py +0 -123
  390. nucliadb/train/tests/test_paragraph_streaming.py +0 -118
  391. nucliadb/train/tests/test_question_answer_streaming.py +0 -239
  392. nucliadb/train/tests/test_sentence_classification.py +0 -143
  393. nucliadb/train/tests/test_token_classification.py +0 -136
  394. nucliadb/train/tests/utils.py +0 -101
  395. nucliadb/writer/layouts/__init__.py +0 -51
  396. nucliadb/writer/layouts/v1.py +0 -59
  397. nucliadb/writer/tests/__init__.py +0 -19
  398. nucliadb/writer/tests/conftest.py +0 -31
  399. nucliadb/writer/tests/fixtures.py +0 -191
  400. nucliadb/writer/tests/test_fields.py +0 -475
  401. nucliadb/writer/tests/test_files.py +0 -740
  402. nucliadb/writer/tests/test_knowledgebox.py +0 -49
  403. nucliadb/writer/tests/test_reprocess_file_field.py +0 -133
  404. nucliadb/writer/tests/test_resources.py +0 -476
  405. nucliadb/writer/tests/test_service.py +0 -137
  406. nucliadb/writer/tests/test_tus.py +0 -203
  407. nucliadb/writer/tests/utils.py +0 -35
  408. nucliadb/writer/tus/pg.py +0 -125
  409. nucliadb-4.0.0.post542.dist-info/METADATA +0 -135
  410. nucliadb-4.0.0.post542.dist-info/RECORD +0 -462
  411. {nucliadb/ingest/tests → migrations/pg}/__init__.py +0 -0
  412. /nucliadb/{ingest/tests/integration → common/external_index_providers}/__init__.py +0 -0
  413. /nucliadb/{ingest/tests/integration/ingest → common/models_utils}/__init__.py +0 -0
  414. /nucliadb/{ingest/tests/unit → search/search/query_parser}/__init__.py +0 -0
  415. /nucliadb/{ingest/tests → tests}/vectors.py +0 -0
  416. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2777.dist-info}/entry_points.txt +0 -0
  417. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2777.dist-info}/top_level.txt +0 -0
  418. {nucliadb-4.0.0.post542.dist-info → nucliadb-6.2.1.post2777.dist-info}/zip-safe +0 -0
nucliadb/tests/tikv.py DELETED
@@ -1,240 +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 os
21
- import platform
22
- import signal
23
- import subprocess
24
- import tarfile
25
- import tempfile
26
- import time
27
- from io import BytesIO
28
-
29
- import pytest
30
- import requests
31
- from tikv_client import TransactionClient # type: ignore
32
-
33
-
34
- class TiKVd(object):
35
- def __init__(
36
- self,
37
- port=20160,
38
- pd_port=2379,
39
- peer_port=2380,
40
- tikv_bin_name="tikv-server",
41
- pd_bin_name="pd-server",
42
- host="127.0.0.1",
43
- path="",
44
- debug=False,
45
- ):
46
- self.port = port
47
- self.pd_port = pd_port
48
- self.peer_port = peer_port
49
- self.tikv_bin_name = tikv_bin_name
50
- self.pd_bin_name = pd_bin_name
51
- self.path = path
52
- self.host = host
53
- self.tmpfolder = None
54
- self.debug = debug
55
- self.proc = None
56
- self.proc2 = None
57
-
58
- def start(self):
59
- self.tmpfolder = tempfile.TemporaryDirectory()
60
-
61
- cmd = [
62
- f"{self.path}/{self.tikv_bin_name}",
63
- f"--pd-endpoints={self.host}:{self.pd_port}",
64
- f"--addr={self.host}:{self.port}",
65
- f"--data-dir={self.tmpfolder.name}/tikv1",
66
- f"--log-file={self.tmpfolder.name}/tikv1.log",
67
- ]
68
-
69
- if self.debug:
70
- self.proc2 = subprocess.Popen(cmd)
71
- else:
72
- self.proc2 = subprocess.Popen(
73
- cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
74
- )
75
-
76
- if self.debug:
77
- if self.proc2 is None:
78
- print(
79
- "[\031[0;33mDEBUG\033[0;0m] Failed to start server listening on port %d started."
80
- % self.port
81
- )
82
- else:
83
- print(
84
- "[\033[0;33mDEBUG\033[0;0m] Server listening on port %d started."
85
- % self.port
86
- )
87
- cmd = [
88
- f"{self.path}/{self.pd_bin_name}",
89
- "--name=pd",
90
- f"--data-dir={self.tmpfolder.name}",
91
- f"--client-urls=http://{self.host}:{self.pd_port}",
92
- f"--peer-urls=http://{self.host}:{self.peer_port}",
93
- f"--initial-cluster=pd=http://{self.host}:{self.peer_port}",
94
- ]
95
-
96
- if self.debug:
97
- self.proc = subprocess.Popen(cmd)
98
- else:
99
- self.proc = subprocess.Popen(
100
- cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
101
- )
102
-
103
- if self.debug:
104
- if self.proc is None:
105
- print(
106
- "[\031[0;33mDEBUG\033[0;0m] Failed to start server listening on port %d started."
107
- % self.pd_port
108
- )
109
- else:
110
- print(
111
- "[\033[0;33mDEBUG\033[0;0m] Server listening on port %d started."
112
- % self.pd_port
113
- )
114
- return self.proc
115
-
116
- def stop(self):
117
- if self.debug:
118
- print(
119
- "[\033[0;33mDEBUG\033[0;0m] Server listening on %d will stop."
120
- % self.port
121
- )
122
-
123
- if self.debug:
124
- if self.proc is None:
125
- print(
126
- "[\033[0;31mDEBUG\033[0;0m] Failed terminating server listening on port %d"
127
- % self.port
128
- )
129
- if self.proc2 is None:
130
- print(
131
- "[\033[0;31mDEBUG\033[0;0m] Failed terminating server listening on port %d"
132
- % self.port
133
- )
134
-
135
- if self.proc.returncode is not None:
136
- if self.debug:
137
- print(
138
- "[\033[0;31mDEBUG\033[0;0m] Server listening on port {port} finished running already with exit {ret}".format( # noqa
139
- port=self.pd_port, ret=self.proc.returncode
140
- )
141
- )
142
- elif self.proc2.returncode is not None:
143
- if self.debug:
144
- print(
145
- "[\033[0;31mDEBUG\033[0;0m] Server listening on port {port} finished running already with exit {ret}".format( # noqa
146
- port=self.port, ret=self.proc2.returncode
147
- )
148
- )
149
- else:
150
- os.kill(self.proc.pid, signal.SIGKILL)
151
- self.proc.wait()
152
- os.kill(self.proc2.pid, signal.SIGKILL)
153
- self.proc2.wait()
154
- if self.debug:
155
- print(
156
- "[\033[0;33mDEBUG\033[0;0m] Server listening on %d was stopped."
157
- % self.port
158
- )
159
- if self.tmpfolder is not None:
160
- self.tmpfolder.cleanup()
161
- self.tmpfolder = None
162
-
163
-
164
- def start_tikvd(tikvd: TiKVd):
165
- tikvd.start()
166
-
167
- endpoint = "127.0.0.1:{port}".format(port=tikvd.port)
168
- retries = 0
169
- while True:
170
- if retries > 100:
171
- break
172
-
173
- try:
174
- connection = TransactionClient.connect(endpoint)
175
- txn = connection.begin(pessimistic=True)
176
- txn.rollback()
177
- except: # noqa
178
- retries += 1
179
- time.sleep(0.1)
180
-
181
-
182
- @pytest.fixture(scope="session")
183
- def tikvd():
184
- if os.environ.get("TESTING_TIKV_LOCAL", None):
185
- yield "localhost", "XX", "2379"
186
- return
187
-
188
- if not os.path.isfile("tikv-server"):
189
- version = "v5.3.1"
190
- arch = platform.machine()
191
- if arch == "x86_64":
192
- arch = "amd64"
193
- system = platform.system().lower()
194
-
195
- resp = requests.get(
196
- f"https://tiup-mirrors.pingcap.com/tikv-{version}-{system}-{arch}.tar.gz"
197
- )
198
-
199
- zipfile = tarfile.open(fileobj=BytesIO(resp.content), mode="r:gz")
200
-
201
- zipfile.extract(f"tikv-server")
202
- os.chmod("tikv-server", 755)
203
-
204
- if not os.path.isfile("pd-server"):
205
- version = "v5.3.1"
206
- arch = platform.machine()
207
- if arch == "x86_64":
208
- arch = "amd64"
209
- system = platform.system().lower()
210
-
211
- resp = requests.get(
212
- f"https://tiup-mirrors.pingcap.com/pd-{version}-{system}-{arch}.tar.gz"
213
- )
214
-
215
- zipfile = tarfile.open(fileobj=BytesIO(resp.content), mode="r:gz")
216
-
217
- zipfile.extract(f"pd-server")
218
- os.chmod("pd-server", 755)
219
-
220
- server = TiKVd(debug=True)
221
- server.tikv_bin_name = "tikv-server"
222
- server.pd_bin_name = "pd-server"
223
- server.path = os.getcwd()
224
-
225
- start_tikvd(server)
226
- print("Started TiKVd")
227
-
228
- for i in range(100):
229
- resp = requests.get(f"http://{server.host}:{server.pd_port}/pd/api/v1/stores")
230
- if (
231
- resp.status_code == 200
232
- and resp.json()["stores"][0]["store"]["state_name"] == "Up"
233
- ):
234
- break
235
- print(resp.status_code)
236
- print(resp.json())
237
- time.sleep(1)
238
-
239
- yield server.host, server.port, server.pd_port
240
- server.stop()
@@ -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,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,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,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,172 +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
- import time
22
- from unittest.mock import ANY, AsyncMock, MagicMock, patch
23
-
24
- import kubernetes_asyncio.client.models.v1_container_status # type: ignore
25
- import kubernetes_asyncio.client.models.v1_object_meta # type: ignore
26
- import kubernetes_asyncio.client.models.v1_pod # type: ignore
27
- import kubernetes_asyncio.client.models.v1_pod_status # type: ignore
28
- import pytest
29
-
30
- from nucliadb.common.cluster import manager
31
- from nucliadb.common.cluster.discovery.k8s import EventType, KubernetesDiscovery
32
- from nucliadb.common.cluster.discovery.types import IndexNodeMetadata
33
- from nucliadb.common.cluster.settings import Settings
34
- from nucliadb_protos import nodewriter_pb2
35
-
36
- pytestmark = pytest.mark.asyncio
37
-
38
-
39
- @pytest.fixture()
40
- def writer_stub():
41
- writer_stub = MagicMock()
42
- writer_stub.GetMetadata = AsyncMock(
43
- return_value=nodewriter_pb2.NodeMetadata(
44
- node_id="node_id", shard_count=1, available_disk=10, total_disk=10
45
- )
46
- )
47
- with (
48
- patch(
49
- "nucliadb.common.cluster.discovery.base.nodewriter_pb2_grpc.NodeWriterStub",
50
- return_value=writer_stub,
51
- ),
52
- patch(
53
- "nucliadb.common.cluster.discovery.base.replication_pb2_grpc.ReplicationServiceStub",
54
- return_value=writer_stub,
55
- ),
56
- patch("nucliadb.common.cluster.discovery.base.get_traced_grpc_channel"),
57
- ):
58
- yield writer_stub
59
-
60
-
61
- @pytest.fixture()
62
- def k8s_discovery(writer_stub):
63
- disc = KubernetesDiscovery(Settings())
64
- disc.node_heartbeat_interval = 0.1
65
- manager.INDEX_NODES.clear()
66
- yield disc
67
- manager.INDEX_NODES.clear()
68
-
69
-
70
- def create_k8s_event(
71
- type_: str = "ADDED",
72
- name: str = "node-0",
73
- container_name="reader",
74
- ready: bool = True,
75
- ) -> EventType:
76
- return {
77
- "type": type_,
78
- "object": kubernetes_asyncio.client.models.v1_pod.V1Pod(
79
- status=kubernetes_asyncio.client.models.v1_pod_status.V1PodStatus(
80
- container_statuses=[
81
- kubernetes_asyncio.client.models.v1_container_status.V1ContainerStatus(
82
- name=container_name,
83
- ready=ready,
84
- image="image",
85
- image_id="image_id",
86
- restart_count=0,
87
- )
88
- ],
89
- pod_ip="1.2.3.4",
90
- ),
91
- metadata=kubernetes_asyncio.client.models.v1_object_meta.V1ObjectMeta(
92
- name=name, labels={"readReplica": "false"}
93
- ),
94
- ),
95
- }
96
-
97
-
98
- async def test_get_node_metadata(k8s_discovery: KubernetesDiscovery, writer_stub):
99
- assert await k8s_discovery.get_node_metadata(
100
- "node-0", "1.1.1.1", read_replica=False
101
- ) == IndexNodeMetadata(
102
- node_id="node_id",
103
- shard_count=1,
104
- name="node_id",
105
- address="1.1.1.1",
106
- updated_at=ANY,
107
- available_disk=10,
108
- )
109
-
110
- writer_stub.GetMetadata.assert_awaited_once()
111
-
112
- # should be cached now
113
- await k8s_discovery.get_node_metadata("node-0", "1.1.1.1", read_replica=False)
114
-
115
- assert len(writer_stub.GetMetadata.mock_calls) == 1
116
-
117
-
118
- async def test_update_node(k8s_discovery: KubernetesDiscovery):
119
- await k8s_discovery.update_node(create_k8s_event())
120
- assert len(manager.INDEX_NODES) == 1
121
-
122
- # update again
123
- await k8s_discovery.update_node(create_k8s_event())
124
- assert len(manager.INDEX_NODES) == 1
125
-
126
-
127
- async def test_remove_node(k8s_discovery: KubernetesDiscovery):
128
- await k8s_discovery.update_node(create_k8s_event())
129
-
130
- assert len(manager.INDEX_NODES) == 1
131
-
132
- await k8s_discovery.update_node(create_k8s_event(ready=False))
133
-
134
- assert len(manager.INDEX_NODES) == 0
135
-
136
-
137
- async def test_update_node_data_cache(k8s_discovery: KubernetesDiscovery, writer_stub):
138
- await k8s_discovery.update_node(create_k8s_event())
139
-
140
- task = asyncio.create_task(k8s_discovery.update_node_data_cache())
141
-
142
- await asyncio.sleep(k8s_discovery.node_heartbeat_interval + 0.1)
143
-
144
- task.cancel()
145
-
146
- assert len(writer_stub.GetMetadata.mock_calls) == 2
147
-
148
-
149
- async def test_remove_stale_nodes(k8s_discovery: KubernetesDiscovery):
150
- manager.add_index_node(
151
- id="node_id",
152
- address="1.1.1.1",
153
- shard_count=1,
154
- available_disk=100,
155
- )
156
- nmd = IndexNodeMetadata(
157
- node_id="node_id",
158
- shard_count=1,
159
- available_disk=100,
160
- name="node_id",
161
- address="1.1.1.1",
162
- updated_at=time.time() - k8s_discovery.node_heartbeat_interval * 3,
163
- )
164
- k8s_discovery.node_id_cache["node_id"] = nmd
165
-
166
- assert len(k8s_discovery.node_id_cache) == 1
167
- assert len(manager.get_index_nodes()) == 1
168
-
169
- k8s_discovery._maybe_remove_stale_node("node_id")
170
-
171
- assert len(k8s_discovery.node_id_cache) == 0
172
- assert len(manager.get_index_nodes()) == 0
@@ -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,114 +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 os
21
- import tempfile
22
- from unittest.mock import AsyncMock, MagicMock, patch
23
-
24
- import pytest
25
- from nucliadb_protos.noderesources_pb2 import ShardId
26
- from nucliadb_protos.utils_pb2 import RelationNode
27
-
28
- from nucliadb.common.cluster.settings import Settings
29
- from nucliadb.common.cluster.standalone import service
30
- from nucliadb_protos import nodereader_pb2, standalone_pb2
31
-
32
- pytestmark = pytest.mark.asyncio
33
-
34
-
35
- @pytest.fixture
36
- def cluster_settings():
37
- settings = Settings()
38
- with (
39
- patch("nucliadb.common.cluster.standalone.service.cluster_settings", settings),
40
- tempfile.TemporaryDirectory() as tmpdir,
41
- ):
42
- settings.data_path = tmpdir
43
- os.makedirs(os.path.join(tmpdir, "shards"))
44
- yield settings
45
-
46
-
47
- @pytest.fixture
48
- def self_node(cluster_settings):
49
- self_node = MagicMock(id="id", address="address", shard_count=0, available_disk=10)
50
- self_node.reader = AsyncMock()
51
- self_node.writer = AsyncMock()
52
- self_node.reader.Search.return_value = nodereader_pb2.SearchResponse()
53
-
54
- relation_edges = nodereader_pb2.EdgeList()
55
- edge = nodereader_pb2.RelationEdge(
56
- edge_type=RelationNode.NodeType.ENTITY, property="foo"
57
- )
58
- relation_edges.list.append(edge)
59
- self_node.reader.RelationEdges.return_value = relation_edges
60
-
61
- with patch("nucliadb.common.cluster.standalone.service.get_self") as mock_get_self:
62
- mock_get_self.return_value = self_node
63
- yield self_node
64
-
65
-
66
- @pytest.fixture
67
- def servicer(self_node):
68
- yield service.StandaloneClusterServiceServicer()
69
-
70
-
71
- async def test_node_action(
72
- servicer: service.StandaloneClusterServiceServicer,
73
- self_node,
74
- cluster_settings,
75
- ):
76
- resp = await servicer.NodeAction(
77
- standalone_pb2.NodeActionRequest(
78
- service="reader",
79
- action="Search",
80
- payload=nodereader_pb2.SearchRequest(body="test").SerializeToString(),
81
- ),
82
- None,
83
- )
84
- assert resp == standalone_pb2.NodeActionResponse(
85
- payload=self_node.reader.Search.return_value.SerializeToString()
86
- )
87
-
88
-
89
- async def test_reader_node_action_relations(
90
- servicer: service.StandaloneClusterServiceServicer,
91
- self_node,
92
- cluster_settings,
93
- ):
94
- resp = await servicer.NodeAction(
95
- standalone_pb2.NodeActionRequest(
96
- service="reader",
97
- action="RelationEdges",
98
- payload=ShardId(id="test").SerializeToString(),
99
- ),
100
- None,
101
- )
102
- assert resp == standalone_pb2.NodeActionResponse(
103
- payload=self_node.reader.RelationEdges.return_value.SerializeToString()
104
- )
105
-
106
-
107
- async def test_node_info(
108
- servicer: service.StandaloneClusterServiceServicer, self_node, cluster_settings
109
- ):
110
- resp = await servicer.NodeInfo(standalone_pb2.NodeInfoRequest(), None)
111
- assert resp.id == self_node.id
112
- assert resp.address == self_node.address
113
- assert resp.shard_count == self_node.shard_count
114
- assert resp.available_disk > 0
@@ -1,61 +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 os
21
- import tempfile
22
- from unittest.mock import patch
23
-
24
- import pytest
25
-
26
- from nucliadb.common.cluster.settings import Settings
27
- from nucliadb.common.cluster.standalone import utils
28
-
29
-
30
- @pytest.fixture
31
- def cluster_settings():
32
- settings = Settings()
33
- with (
34
- patch("nucliadb.common.cluster.standalone.utils.cluster_settings", settings),
35
- tempfile.TemporaryDirectory() as tmpdir,
36
- ):
37
- settings.data_path = tmpdir
38
- yield settings
39
-
40
-
41
- def test_get_standalone_node_id(cluster_settings: Settings):
42
- assert utils.get_standalone_node_id()
43
- assert os.path.exists(os.path.join(cluster_settings.data_path, "node.key"))
44
-
45
-
46
- def test_get_self(cluster_settings: Settings):
47
- os.makedirs(os.path.join(cluster_settings.data_path, "shards", "1"))
48
-
49
- assert utils.get_self().shard_count == 1
50
-
51
-
52
- def test_get_self_k8s_host(cluster_settings: Settings, monkeypatch):
53
- monkeypatch.setenv("NUCLIADB_SERVICE_HOST", "host")
54
- monkeypatch.setenv("HOSTNAME", "nucliadb-0")
55
-
56
- with (
57
- patch("nucliadb.common.cluster.standalone.grpc_node_binding.NodeWriter"),
58
- patch("nucliadb.common.cluster.standalone.grpc_node_binding.NodeReader"),
59
- ):
60
- # patch because loading settings validates address now
61
- assert utils.get_self().address == "nucliadb-0.nucliadb"