py10x-universe 0.1.3__py3-none-any.whl

infra_10x/mongodb_store.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from core_10x.global_cache import cache
+from core_10x.nucleus import Nucleus
+from core_10x.ts_store import TsCollection, TsDuplicateKeyError, TsStore, standard_key
+from py10x_infra import MongoCollectionHelper
+from pymongo import MongoClient, errors
+from pymongo.errors import DuplicateKeyError
+from pymongo.uri_parser import parse_uri as pymongo_parse_uri
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from core_10x.ts_store import f
+    from pymongo.collection import Collection
+    from pymongo.database import Database
+
+
+class MongoCollection(TsCollection):
+    s_id_tag = '_id'
+
+    assert Nucleus.ID_TAG() == s_id_tag, f"Nucleus.ID_TAG() must be '{s_id_tag}'"
+
+    def __init__(self, db, collection_name: str):
+        self.coll: Collection = db[collection_name]
+
+    def collection_name(self) -> str:
+        return self.coll.name
+
+    def id_exists(self, id_value: str) -> bool:
+        return self.coll.count_documents({self.s_id_tag: id_value}) > 0
+
+    def find(self, query: f = None, _at_most: int = 0, _order: dict = None) -> Iterable:
+        cursor = self.coll.find(query.prefix_notation()) if query else self.coll.find()
+        if _order:
+            cursor = cursor.sort(list(_order.items()))
+        if _at_most:
+            cursor = cursor.limit(_at_most)
+        return cursor
+
+    def count(self, query: f = None) -> int:
+        return self.coll.count_documents(query.prefix_notation()) if query else self.coll.count_documents({})
+
+    def save_new(self, serialized_traitable: dict, overwrite: bool = False) -> int:
+        needs_upsert = (set_values := serialized_traitable.get('$set')) or overwrite
+        id_tag = self.s_id_tag
+        id_value = (set_values or serialized_traitable)[id_tag]
+        e = None
+
+        # TODO: overwrite via save(), not save_new() so that revision is incremented rather than reset
+        (set_values or serialized_traitable)[Nucleus.REVISION_TAG()] = 1
+
+        if not needs_upsert:
+            try:
+                res = self.coll.insert_one(serialized_traitable)
+            except DuplicateKeyError:
+                ack, cnt = False, 1
+            else:
+                ack, cnt = res.acknowledged, 0
+        else:
+            res = self.coll.update_one({id_tag: id_value}, serialized_traitable if set_values else {'$set': serialized_traitable}, upsert=True)
+            ack, cnt = res.acknowledged, res.matched_count
+
+        if cnt and not overwrite:  # -- e.g. this id/revision already existed
+            raise TsDuplicateKeyError(self.collection_name(), {id_tag: id_value}) from e
+
+        return int(ack)
+
+    def save(self, serialized_traitable: dict) -> int:
+        rev_tag = Nucleus.REVISION_TAG()
+        id_tag = self.s_id_tag
+
+        revision = serialized_traitable.get(rev_tag, -1)
+        assert revision >= 0, 'revision must be >= 0'
+
+        if revision == 0:
+            return self.save_new(serialized_traitable)
+
+        id_value = serialized_traitable.get(id_tag)
+
+        filter = {}
+        pipeline = []
+        serialized_traitable = dict(serialized_traitable)  # -- copy to avoid modifying the input
+        MongoCollectionHelper.prepare_filter_and_pipeline(serialized_traitable, filter, pipeline)
+        # self.filter_and_pipeline(serialized_traitable, filter, pipeline)
+
+        res = self.coll.update_one(filter, pipeline)
+        if not res.acknowledged:
+            return revision
+
+        if not res.matched_count:  # -- e.g. restore from deleted
+            raise AssertionError(f'{self.coll} {id_value} has been most probably inappropriately restored from deleted')
+
+        if res.matched_count != 1:
+            return revision
+
+        return revision if res.modified_count != 1 else revision + 1
+
+    @classmethod
+    def filter_and_pipeline(cls, serialized_traitable, filter, pipeline):
+        rev_tag = Nucleus.REVISION_TAG()
+        id_tag = cls.s_id_tag
+        for key in (rev_tag, id_tag):
+            filter[key] = serialized_traitable.pop(key)
+
+        rev_condition = {'$and': [{'$eq': ['$' + name, {'$literal': value}]} for name, value in serialized_traitable.items()]}
+
+        # fmt: off
+        update_revision = {
+            '$cond': [
+                rev_condition,         #-- if each field is equal to its prev value
+                filter[rev_tag],       #       then, keep the revision as is
+                filter[rev_tag] + 1    #       else, increment it
+            ]
+        }
+
+        pipeline.append(
+            {
+                '$replaceRoot': {
+                    'newRoot': {
+                        id_tag:     filter[id_tag],
+                        rev_tag:    update_revision,
+                    }
+                }
+            }
+        )
+        # fmt: on
+
+        pipeline.extend(
+            {'$replaceWith': {'$setField': dict(field=field, input='$$ROOT', value={'$literal': value})}}
+            for field, value in serialized_traitable.items()
+        )
+
+    # #---- TODO: move to C++
+    # def save(self, serialized_traitable: dict) -> int:
+    #     """
+    #     Updates (and inc _rev) only if at least one traits has been changed
+    #     :returns new revision if successful, otherwise the old one
+    #     """
+    #
+    #     rev_tag = Nucleus.REVISION_TAG
+    #     id_tag = self.s_id_tag
+    #
+    #     revision = serialized_traitable.get(rev_tag, -1)
+    #     assert revision >= 0, 'revision must be >= 0'
+    #
+    #     if revision == 0:
+    #         return self.save_new(serialized_traitable)
+    #
+    #     id_value = serialized_traitable.get(id_tag)
+    #     del serialized_traitable[id_tag]
+    #     del serialized_traitable[rev_tag]
+    #
+    #     filter = {
+    #         id_tag:     id_value,
+    #         rev_tag:    revision,
+    #     }
+    #
+    #     rev_condition = {
+    #         '$and': [ {'$eq': ['$' + name, {'$literal': value}] } for name, value in serialized_traitable.items() ]
+    #     }
+    #
+    #     update_revision = {
+    #         '$cond': [
+    #             rev_condition,  #-- if each field is equal to its prev value
+    #             revision,       #       then, keep the revision as is
+    #             revision + 1    #       else, increment it
+    #         ]
+    #     }
+    #
+    #     cmds = [
+    #         {
+    #             '$replaceRoot': {
+    #                 'newRoot': {
+    #                     id_tag:     id_value,
+    #                     rev_tag:    update_revision,
+    #                 }
+    #             }
+    #         }
+    #     ]
+    #
+    #     cmds.extend(
+    #         {
+    #             '$replaceWith': {
+    #                 '$setField': dict(field = field, input = '$$ROOT', value = {'$literal': value})
+    #             }
+    #         }
+    #         for field, value in serialized_traitable.items()
+    #     )
+    #
+    #     res = self.coll.update_one(filter, cmds)
+    #     if not res.acknowledged:
+    #         return revision
+    #
+    #     if not res.matched_count:  # -- e.g. restore from deleted
+    #         raise AssertionError(f'{self.coll} {id_value} has been most probably inapropriately restored from deleted')
+    #
+    #     if res.matched_count != 1:
+    #         return revision
+    #
+    #     return revision if res.modified_count != 1 else revision + 1
+
+    def delete(self, id_value: str) -> bool:
+        q = {self.s_id_tag: id_value}
+        return self.coll.delete_one(q).acknowledged
+
+    def create_index(self, name: str, trait_name: str | list[tuple[str, int]], **index_args) -> str | None:
+        index_info = self.coll.index_information()
+        if name in index_info:
+            return None
+
+        return self.coll.create_index(trait_name, name=name, **index_args)
+
+    def max(self, trait_name: str, filter: f = None) -> dict | None:
+        if filter:
+            cur = self.coll.find(filter.prefix_notation()).sort({trait_name: -1}).limit(1)
+        else:
+            cur = self.coll.find().sort({trait_name: -1}).limit(1)
+        for data in cur:
+            return data
+
+        return None
+
+    def min(self, trait_name: str, filter: f = None) -> dict | None:
+        if filter:
+            cur = self.coll.find(filter.prefix_notation()).sort({trait_name: 1}).limit(1)
+        else:
+            cur = self.coll.find().sort({trait_name: 1}).limit(1)
+        for data in cur:
+            return data
+
+        return None
+
+    def load(self, id_value: str) -> dict | None:
+        for data in self.coll.find({self.s_id_tag: id_value}):
+            return data
+
+        return None
+
+
+class MongoStore(TsStore, resource_name='MONGO_DB'):
+    PROTOCOL = 'mongodb'
+
+    ADMIN = 'admin'
+
+    # fmt: off
+    s_instance_kwargs_map = dict(
+        port    = ('port',                      27017),
+        ssl     = ('ssl',                       False),
+        sst     = ('serverSelectionTimeoutMS',  10000),
+    )
+    # fmt: on
+
+    s_cached_connections: dict[tuple, MongoClient] = {}
+
+    @classmethod
+    def connect(cls, hostname: str, username: str, password: str, _cache: bool = True, _throw: bool = True, **kwargs) -> MongoClient:
+        connection_key = standard_key((hostname, username), kwargs) if _cache else None
+        client = cls.s_cached_connections.get(connection_key)
+        if not client:
+            client = MongoClient(hostname, username=username, password=password, **kwargs)
+            try:
+                client.server_info()
+            except Exception:
+                client.close()
+                if _throw:
+                    raise
+                client = None
+        if client and connection_key:
+            cls.s_cached_connections[connection_key] = client
+
+        return client
+
+    @classmethod
+    def uncache_connection(cls, hostname: str, username: str, password: str, **kwargs):
+        connection_key = standard_key((hostname, username), kwargs)
+        client = cls.s_cached_connections.pop(connection_key, None)
+        if client:
+            client.close()
+
+    # noinspection PyMethodOverriding
+    @classmethod
+    def new_instance(cls, hostname: str, dbname: str, username: str, password: str, **kwargs) -> TsStore:
+        client = cls.connect(hostname, username, password, **kwargs)
+        return cls(client, client[dbname], username)
+
+    @classmethod
+    def parse_uri(cls, uri: str) -> dict:
+        params = pymongo_parse_uri(uri)
+        try:
+            # fmt: off
+            hostname, port  = params['nodelist'][0]
+            kwargs          = params['options']
+            kwargs[cls.PORT_TAG]  = port
+            args = {
+                cls.HOSTNAME_TAG:   hostname,
+                cls.DBNAME_TAG:     params['database'],
+                cls.USERNAME_TAG:   params['username'],
+                cls.PASSWORD_TAG:   params['password'],
+            }
+            # fmt: on
+            args.update(kwargs)
+            return args
+
+        except Exception as e:
+            raise ValueError(f'Invalid URI = {uri}') from e
+
+    def __init__(self, client: MongoClient, db: Database, username: str):
+        self.client = client
+        self.db: Database = db
+        self.username = username
+
+    def collection_names(self, regexp: str = None) -> list:
+        filter = dict(name={'$regex': regexp}) if regexp else None
+        return self.db.list_collection_names(filter=filter)
+
+    def collection(self, collection_name: str) -> TsCollection:
+        return MongoCollection(self.db, collection_name)
+
+    def delete_collection(self, collection_name: str) -> bool:
+        self.db.drop_collection(collection_name)
+        return True
+
+    @classmethod
+    @cache
+    def is_running_with_auth(cls, host_name: str) -> tuple:  # -- (is_running, with_auth)
+        client = cls.connect(host_name, '', '', _cache=False, _throw=False)
+        if not client:
+            return False, False
+
+        admin_db = client[cls.ADMIN]
+        try:
+            res = admin_db.command('getCmdLineOpts')
+            auth = any(r == '--auth' for r in res['argv'][1:])
+            return True, auth
+
+        except errors.OperationFailure:  # -- auth is required
+            return True, True
+
+        finally:
+            client.close()
+
+    def auth_user(self) -> str | None:
+        return self.username

infra_10x/mongodb_utils.py

@@ -0,0 +1,129 @@
+import re
+import secrets
+import string
+import subprocess
+from getpass import getpass
+
+import keyring
+import requests
+from core_10x.backbone.namespace import AUTO_PASSWORD_LENGTH, USER_ROLE
+from core_10x.global_cache import cache
+from core_10x.rc import RC
+
+from infra_10x.mongodb_admin import MongodbAdmin
+from infra_10x.mongodb_store import MongoStore
+
+
+@cache
+def create_xx_role(admin: MongodbAdmin):
+    admin.db.command(
+        'updateRole' if admin.role_exists('xxUser') else 'createRole',
+        'xxUser',
+        privileges=[
+            {
+                'resource': {'anyResource': True},
+                'actions': USER_ROLE.WORKER.value,
+            }
+        ],
+        roles=[],
+    )
+
+
+def create_xx_user(admin: MongodbAdmin, xx_user: str, new_password: str) -> None:
+    password = keyring.get_password('xxuser', xx_user) or new_password
+    admin.update_user(xx_user, password, 'xxUser', keep_current_roles=False)
+    if password == new_password:
+        keyring.set_password('xxuser', xx_user, password)
+
+
+def create_xx_admin(admin: MongodbAdmin, xx_admin: str, new_password: str) -> None:
+    if xx_admin == admin.username:
+        raise ValueError('cannot create self!')
+    password = keyring.get_password('xxadmin', xx_admin) or new_password
+    admin.update_user(xx_admin, password, 'userAdminAnyDatabase', 'readWriteAnyDatabase')
+    if password == new_password:
+        keyring.set_password('xxadmin', xx_admin, password)
+
+
+def get_xx_admin(user: str):
+    return f'{user}-admin'
+
+
+@cache
+def git_user():
+    return requests.get('https://api.github.com/user', headers={'Authorization': f'token {git_token()}'}).json()['login'] or input('Username:')
+
+
+@cache
+def git_token():
+    return subprocess.getoutput("git credential fill <<< 'protocol=https\nhost=github.com\n' | grep password | cut -d= -f2") or input('Token:')
+
+
+@cache
+def all_git_users():
+    git_org = re.search(r'github\.com[/:]([\w\-\.]+)/', subprocess.check_output(['git', 'remote', 'get-url', 'origin'], text=True).strip()).group(1)
+    return [
+        user['login']
+        for user in requests.get(f'https://api.github.com/orgs/{git_org}/members', headers={'Authorization': f'token {git_token()}'}).json()
+    ]
+
+
+def generate_password():
+    return ''.join(secrets.choice(string.ascii_letters + string.digits + '!@#$%^&*') for _ in range(AUTO_PASSWORD_LENGTH))
+
+
+def create_users(hostname):
+    admin = MongodbAdmin(hostname=hostname, **get_credentials('xxadmin', hostname))
+    create_xx_role(admin)
+    for xx_user in all_git_users():
+        xx_admin = get_xx_admin(xx_user)
+        if xx_admin != admin.username:  # -- do not break current working account!
+            create_xx_admin(admin, xx_admin, generate_password())
+        create_xx_user(admin, xx_user, generate_password())
+
+
+def clear_vault():
+    for xx_user in all_git_users():
+        if keyring.get_password('xxuser', xx_user):
+            keyring.delete_password('xxuser', xx_user)
+        if xx_user != git_user():  # -- do not break own admin account!
+            xx_admin = get_xx_admin(xx_user)
+            if keyring.get_password('xxadmin', xx_admin):
+                keyring.delete_password('xxadmin', xx_admin)
+
+
+def drop_users(hostname):
+    admin = MongodbAdmin(hostname=hostname, **get_credentials('xxadmin', hostname))
+    for xx_user in all_git_users():
+        xx_admin = get_xx_admin(xx_user)
+        if xx_admin != admin.username:  # -- do not break current working account!
+            admin.delete_user(xx_admin)
+        admin.delete_user(xx_user)
+
+
+def get_credentials(service, hostname):
+    if hostname == 'localhost':
+        return dict(username='', password='')
+    username = git_user()
+    if service == 'xxadmin':
+        username = get_xx_admin(username)
+    password = keyring.get_password(service, username)
+    if not password:
+        password = getpass(f"{username}'s password for `{service}`:")
+        keyring.set_password(service, username, password)
+    return dict(username=username, password=password)
+
+
+def copy_db(from_host: str, to_host: str, dbname: str, overwrite=False) -> RC:
+    from_store = MongoStore.instance(hostname=from_host, dbname=dbname, **get_credentials('xxuser', from_host))
+    to_store = MongoStore.instance(hostname=to_host, dbname=dbname, **get_credentials('xxuser', to_host))
+    return from_store.copy_to(to_store, overwrite=overwrite)
+
+
+if __name__ == '__main__':
+    host = 'mongo10x.eastus2.cloudapp.azure.com'
+    # drop_users(host)
+    # clear_vault()
+    # create_users(host)
+
+    copy_db(from_host='localhost', to_host=host, dbname='mkt_data', overwrite=False).throw()

infra_10x/unit_tests/conftest.py

@@ -0,0 +1,13 @@
+import pytest
+from infra_10x.mongodb_store import MongoStore
+
+MONGO_URL = 'mongodb://localhost:27017/'
+# MONGO_URL="mongodb+srv://dev.qbultu3.ts_storedb.net/?authMechanism=MONGODB-X509&authSource=%24external&tls=true&tlsCertificateKeyFile=%2FUsers%2Fiap%2FDownloads%2FX509-cert-590074097809994161.pem"
+TEST_DB = 'test_db'
+
+
+@pytest.fixture(scope='module')
+def ts_instance():
+    instance = MongoStore.instance(hostname=MONGO_URL, dbname=TEST_DB)
+    instance.username = 'test_user'
+    return instance

infra_10x/unit_tests/test_mongo_db.py

@@ -0,0 +1,36 @@
+from core_10x.testlib.ts_tests import TestTSStore, ts_setup
+from infra_10x.mongodb_store import MongoCollection, MongoStore
+
+
+def test_mongo_parse_uri_round_trip():
+    uri = 'mongodb://user:pass@localhost:27017/testdb?ssl=false&serverSelectionTimeoutMS=5000'
+    args = MongoStore.parse_uri(uri)
+
+    assert args[MongoStore.HOSTNAME_TAG] == 'localhost'
+    assert args[MongoStore.DBNAME_TAG] == 'testdb'
+    assert args[MongoStore.USERNAME_TAG] == 'user'
+    assert args[MongoStore.PASSWORD_TAG] == 'pass'
+    assert args['port'] == 27017
+    # Options are driver-dependent; we only require that custom options are propagated.
+    assert 'serverSelectionTimeoutMS' in args
+
+
+def test_filter_and_pipeline_equivalence():
+    # This test mirrors infra_10x/manual_tests/test_prepare_filter_and_pipeline.py
+    from py10x_infra import MongoCollectionHelper
+
+    # The helper only constructs filter/pipeline; no need for a live Mongo instance here.
+    serialized_traitable = dict(_id='AAAA', _rev=10, name='test', age=60)
+
+    data1 = dict(serialized_traitable)
+    pipeline1: list = []
+    filter1: dict = {}
+    MongoCollection.filter_and_pipeline(data1, filter1, pipeline1)
+
+    data2 = dict(serialized_traitable)
+    pipeline2: list = []
+    filter2: dict = {}
+    MongoCollectionHelper.prepare_filter_and_pipeline(data2, filter2, pipeline2)
+
+    assert filter1 == filter2
+    assert pipeline1 == pipeline2

infra_10x/unit_tests/test_mongo_history.py

@@ -0,0 +1 @@
+from core_10x.testlib.traitable_history_tests import TestTraitableHistory, test_collection, test_store