<s> import json <EOL> import logging <EOL> import os <EOL> import shutil <EOL> from django . utils . translation import ugettext as _ <EOL> from desktop . lib . exceptions_renderable import PopupException <EOL> from libsolr . api import SolrApi <EOL> from libzookeeper . conf import ENSEMBLE <EOL> from libzookeeper . models import ZookeeperClient <EOL> from search . conf import SOLR_URL , SECURITY_ENABLED <EOL> from indexer . conf import CORE_INSTANCE_DIR <EOL> from indexer . utils import copy_configs , field_values_from_log , field_values_from_separated_file <EOL> LOG = logging . getLogger ( __name__ ) <EOL> MAX_UPLOAD_SIZE = <NUM_LIT:100> * <NUM_LIT> * <NUM_LIT> <EOL> ALLOWED_FIELD_ATTRIBUTES = set ( [ '<STR_LIT:name>' , '<STR_LIT:type>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> FLAGS = [ ( '<STR_LIT:I>' , '<STR_LIT>' ) , ( '<STR_LIT:T>' , '<STR_LIT>' ) , ( '<STR_LIT:S>' , '<STR_LIT>' ) ] <EOL> ZK_SOLR_CONFIG_NAMESPACE = '<STR_LIT>' <EOL> def get_solr_ensemble ( ) : <EOL> return '<STR_LIT>' % ENSEMBLE . get ( ) <EOL> class CollectionManagerController ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , user ) : <EOL> self . user = user <EOL> def _format_flags ( self , fields ) : <EOL> for name , properties in fields . items ( ) : <EOL> for ( code , value ) in FLAGS : <EOL> if code in properties [ '<STR_LIT>' ] : <EOL> properties [ value ] = True <EOL> return fields <EOL> def is_solr_cloud_mode ( self ) : <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> try : <EOL> api . collections ( ) <EOL> setattr ( self , '<STR_LIT>' , True ) <EOL> except Exception , e : <EOL> LOG . info ( '<STR_LIT>' % e ) <EOL> setattr ( self , '<STR_LIT>' , False ) <EOL> return getattr ( self , '<STR_LIT>' ) <EOL> def collection_exists ( self , collection ) : <EOL> return collection in self . get_collections ( ) <EOL> def get_collections ( self ) : <EOL> try : <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> if self . is_solr_cloud_mode ( ) : <EOL> solr_collections = api . collections ( ) <EOL> for name in solr_collections : <EOL> solr_collections [ name ] [ '<STR_LIT>' ] = False <EOL> else : <EOL> solr_collections = { } <EOL> solr_cores = api . cores ( ) <EOL> for name in solr_cores : <EOL> solr_cores [ name ] [ '<STR_LIT>' ] = True <EOL> solr_aliases = api . aliases ( ) <EOL> for name in solr_aliases : <EOL> solr_aliases [ name ] = { <EOL> '<STR_LIT>' : False , <EOL> '<STR_LIT>' : True , <EOL> '<STR_LIT>' : solr_aliases [ name ] <EOL> } <EOL> except Exception , e : <EOL> LOG . warn ( '<STR_LIT>' % e ) <EOL> solr_collections = { } <EOL> solr_cores = { } <EOL> solr_aliases = { } <EOL> solr_cores . update ( solr_collections ) <EOL> solr_cores . update ( solr_aliases ) <EOL> return solr_cores <EOL> def get_fields ( self , collection_or_core_name ) : <EOL> try : <EOL> field_data = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) . fields ( collection_or_core_name ) <EOL> fields = self . _format_flags ( field_data [ '<STR_LIT>' ] [ '<STR_LIT>' ] ) <EOL> except : <EOL> LOG . exception ( _ ( '<STR_LIT>' ) % collection_or_core_name ) <EOL> raise PopupException ( _ ( '<STR_LIT>' ) % collection_or_core_name ) <EOL> try : <EOL> uniquekey = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) . uniquekey ( collection_or_core_name ) <EOL> except : <EOL> LOG . exception ( _ ( '<STR_LIT>' ) % collection_or_core_name ) <EOL> raise PopupException ( _ ( '<STR_LIT>' ) % collection_or_core_name ) <EOL> return uniquekey , fields <EOL> def create_collection ( self , name , fields , unique_key_field = '<STR_LIT:id>' , df = '<STR_LIT:text>' ) : <EOL> """<STR_LIT>""" <EOL> if self . is_solr_cloud_mode ( ) : <EOL> self . _create_solr_cloud_collection ( name , fields , unique_key_field , df ) <EOL> else : <EOL> self . _create_non_solr_cloud_collection ( name , fields , unique_key_field , df ) <EOL> def _create_solr_cloud_collection ( self , name , fields , unique_key_field , df ) : <EOL> with ZookeeperClient ( hosts = get_solr_ensemble ( ) , read_only = False ) as zc : <EOL> root_node = '<STR_LIT>' % ( ZK_SOLR_CONFIG_NAMESPACE , name ) <EOL> tmp_path , solr_config_path = copy_configs ( fields , unique_key_field , df , True ) <EOL> try : <EOL> config_root_path = '<STR_LIT>' % ( solr_config_path , '<STR_LIT>' ) <EOL> try : <EOL> zc . copy_path ( root_node , config_root_path ) <EOL> except Exception , e : <EOL> zc . delete_path ( root_node ) <EOL> raise PopupException ( _ ( '<STR_LIT>' ) , detail = e ) <EOL> finally : <EOL> shutil . rmtree ( tmp_path ) <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> if not api . create_collection ( name ) : <EOL> try : <EOL> zc . delete_path ( root_node ) <EOL> except Exception , e : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) , detail = e ) <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> def _create_non_solr_cloud_collection ( self , name , fields , unique_key_field , df ) : <EOL> instancedir = os . path . join ( CORE_INSTANCE_DIR . get ( ) , name ) <EOL> if os . path . exists ( instancedir ) : <EOL> raise PopupException ( _ ( "<STR_LIT>" ) % instancedir ) <EOL> tmp_path , solr_config_path = copy_configs ( fields , unique_key_field , df , False ) <EOL> try : <EOL> shutil . move ( solr_config_path , instancedir ) <EOL> finally : <EOL> shutil . rmtree ( tmp_path ) <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> if not api . create_core ( name , instancedir ) : <EOL> shutil . rmtree ( instancedir ) <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> def delete_collection ( self , name , core ) : <EOL> """<STR_LIT>""" <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> if core : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> if api . remove_collection ( name ) : <EOL> try : <EOL> root_node = '<STR_LIT>' % ( ZK_SOLR_CONFIG_NAMESPACE , name ) <EOL> with ZookeeperClient ( hosts = get_solr_ensemble ( ) , read_only = False ) as zc : <EOL> zc . delete_path ( root_node ) <EOL> except Exception , e : <EOL> api . add_collection ( name ) <EOL> raise PopupException ( _ ( '<STR_LIT>' ) , detail = e ) <EOL> else : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> def update_collection ( self , name , fields ) : <EOL> """<STR_LIT>""" <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> old_field_names = api . fields ( name ) [ '<STR_LIT>' ] [ '<STR_LIT>' ] . keys ( ) <EOL> new_fields = filter ( lambda field : field [ '<STR_LIT:name>' ] not in old_field_names , fields ) <EOL> new_fields_filtered = [ ] <EOL> for field in new_fields : <EOL> new_field = { } <EOL> for attribute in filter ( lambda attribute : attribute in field , ALLOWED_FIELD_ATTRIBUTES ) : <EOL> new_field [ attribute ] = field [ attribute ] <EOL> new_fields_filtered . append ( new_field ) <EOL> api . add_fields ( name , new_fields_filtered ) <EOL> def update_data_from_hdfs ( self , fs , collection_or_core_name , fields , path , data_type = '<STR_LIT>' , indexing_strategy = '<STR_LIT>' , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> if indexing_strategy == '<STR_LIT>' : <EOL> stats = fs . stats ( path ) <EOL> if stats . size > MAX_UPLOAD_SIZE : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> else : <EOL> unique_key , fields = self . get_fields ( collection_or_core_name ) <EOL> fields = [ { '<STR_LIT:name>' : field , '<STR_LIT:type>' : fields [ field ] [ '<STR_LIT:type>' ] } for field in fields ] <EOL> fh = fs . open ( path ) <EOL> if data_type == '<STR_LIT>' : <EOL> data = json . dumps ( [ value for value in field_values_from_log ( fh , fields ) ] ) <EOL> content_type = '<STR_LIT>' <EOL> elif data_type == '<STR_LIT>' : <EOL> data = json . dumps ( [ value for value in field_values_from_separated_file ( fh , kwargs . get ( '<STR_LIT>' , '<STR_LIT:U+002C>' ) , kwargs . get ( '<STR_LIT>' , '<STR_LIT:">' ) , fields ) ] , indent = <NUM_LIT:2> ) <EOL> content_type = '<STR_LIT>' <EOL> else : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) % data_type ) <EOL> fh . close ( ) <EOL> if not api . update ( collection_or_core_name , data , content_type = content_type ) : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> else : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) % indexing_strategy ) <EOL> def update_data_from_hive ( self , db , collection_or_core_name , database , table , columns , indexing_strategy = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> from beeswax . server import dbms <EOL> import tablib <EOL> api = SolrApi ( SOLR_URL . get ( ) , self . user , SECURITY_ENABLED . get ( ) ) <EOL> if indexing_strategy == '<STR_LIT>' : <EOL> table = db . get_table ( database , table ) <EOL> hql = "<STR_LIT>" % ( '<STR_LIT:U+002C>' . join ( columns ) , database , table . name , db . _get_browse_limit_clause ( table ) ) <EOL> query = dbms . hql_query ( hql ) <EOL> try : <EOL> handle = db . execute_and_wait ( query ) <EOL> if handle : <EOL> result = db . fetch ( handle , rows = <NUM_LIT:100> ) <EOL> db . close ( handle ) <EOL> dataset = tablib . Dataset ( ) <EOL> dataset . append ( columns ) <EOL> for row in result . rows ( ) : <EOL> dataset . append ( row ) <EOL> if not api . update ( collection_or_core_name , dataset . csv , content_type = '<STR_LIT>' ) : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> else : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) ) <EOL> except Exception , e : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) , detail = e ) <EOL> else : <EOL> raise PopupException ( _ ( '<STR_LIT>' ) % indexing_strategy ) </s>
