mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-01-07 19:31:51 +08:00
87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
import logging
|
|
from sqlalchemy import exc as sqlalchemy_exc
|
|
from sqlalchemy import and_, or_
|
|
from mishards.models import Tables
|
|
from mishards.router import RouterMixin
|
|
from mishards import exceptions, db
|
|
from mishards.hash_ring import HashRing
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Factory(RouterMixin):
|
|
name = 'FileBasedHashRingRouter'
|
|
|
|
def __init__(self, writable_topo, readonly_topo, **kwargs):
|
|
super(Factory, self).__init__(writable_topo=writable_topo,
|
|
readonly_topo=readonly_topo)
|
|
|
|
def routing(self, table_name, partition_tags=None, metadata=None, **kwargs):
|
|
range_array = kwargs.pop('range_array', None)
|
|
return self._route(table_name, range_array, partition_tags, metadata, **kwargs)
|
|
|
|
def _route(self, table_name, range_array, partition_tags=None, metadata=None, **kwargs):
|
|
# PXU TODO: Implement Thread-local Context
|
|
# PXU TODO: Session life mgt
|
|
|
|
if not partition_tags:
|
|
cond = and_(
|
|
or_(Tables.table_id == table_name, Tables.owner_table == table_name),
|
|
Tables.state != Tables.TO_DELETE)
|
|
else:
|
|
cond = and_(Tables.state != Tables.TO_DELETE,
|
|
Tables.owner_table == table_name,
|
|
Tables.partition_tag.in_(partition_tags))
|
|
try:
|
|
tables = db.Session.query(Tables).filter(cond).all()
|
|
except sqlalchemy_exc.SQLAlchemyError as e:
|
|
raise exceptions.DBError(message=str(e), metadata=metadata)
|
|
|
|
if not tables:
|
|
raise exceptions.TableNotFoundError('{}:{}'.format(table_name, partition_tags), metadata=metadata)
|
|
|
|
total_files = []
|
|
for table in tables:
|
|
files = table.files_to_search(range_array)
|
|
total_files.append(files)
|
|
|
|
db.remove_session()
|
|
|
|
servers = self.readonly_topo.group_names
|
|
logger.info('Available servers: {}'.format(list(servers)))
|
|
|
|
ring = HashRing(servers)
|
|
|
|
routing = {}
|
|
|
|
for files in total_files:
|
|
for f in files:
|
|
target_host = ring.get_node(str(f.id))
|
|
sub = routing.get(target_host, None)
|
|
if not sub:
|
|
sub = {}
|
|
routing[target_host] = sub
|
|
kv = sub.get(f.table_id, None)
|
|
if not kv:
|
|
kv = []
|
|
sub[f.table_id] = kv
|
|
sub[f.table_id].append(str(f.id))
|
|
|
|
return routing
|
|
|
|
@classmethod
|
|
def Create(cls, **kwargs):
|
|
writable_topo = kwargs.pop('writable_topo', None)
|
|
if not writable_topo:
|
|
raise RuntimeError('Cannot find \'writable_topo\' to initialize \'{}\''.format(self.name))
|
|
readonly_topo = kwargs.pop('readonly_topo', None)
|
|
if not readonly_topo:
|
|
raise RuntimeError('Cannot find \'readonly_topo\' to initialize \'{}\''.format(self.name))
|
|
router = cls(writable_topo=writable_topo, readonly_topo=readonly_topo, **kwargs)
|
|
return router
|
|
|
|
|
|
def setup(app):
|
|
logger.info('Plugin \'{}\' Installed In Package: {}'.format(__file__, app.plugin_package_name))
|
|
app.on_plugin_setup(Factory)
|