Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
__init__.py544 linesDownload Raw Back to biganimal
1# ##########################################################################2# #3# # pgAdmin 4 - PostgreSQL Tools4# #5# # Copyright (C) 2013 - 2024, The pgAdmin Development Team6# # This software is released under the PostgreSQL Licence7# #8# ##########################################################################9 10# EDB BigAnimal Cloud Deployment Implementation11 12import requests13import json14import pickle15from flask_babel import gettext16from flask import session, current_app17from pgadmin.user_login_check import pga_login_required18from werkzeug.datastructures import Headers19from pgadmin.utils import PgAdminModule20from pgadmin.misc.cloud.utils import _create_server, CloudProcessDesc21from pgadmin.misc.bgprocess.processes import BatchProcess22from pgadmin.utils.ajax import make_json_response23from config import root24from pgadmin.utils.constants import MIMETYPE_APP_JSON25 26MODULE_NAME = 'biganimal'27 28SINGLE_CLUSTER_ARCH = 'single'29HA_CLUSTER_ARCH = 'ha'  # High Availability30EHA_CLUSTER_ARCH = 'eha'  # Extreme High Availability31 32 33class BigAnimalModule(PgAdminModule):34    """Cloud module to deploy on EDB BigAnimal"""35 36    def get_exposed_url_endpoints(self):37        return ['biganimal.verification',38                'biganimal.verification_ack',39                'biganimal.regions',40                'biganimal.db_types',41                'biganimal.db_versions',42                'biganimal.instance_types',43                'biganimal.volume_types',44                'biganimal.volume_properties',45                'biganimal.providers',46                'biganimal.projects']47 48 49blueprint = BigAnimalModule(MODULE_NAME, __name__,50                            static_url_path='/misc/cloud/biganimal')51 52 53@blueprint.route('/verification_ack/',54                 methods=['GET'], endpoint='verification_ack')55@pga_login_required56def biganimal_verification_ack():57    """Check the Verification is done or not."""58    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])59    status, error = biganimal_obj.polling_for_token()60    if status:61        session['biganimal']['provider_obj'] = pickle.dumps(biganimal_obj, -1)62    return make_json_response(success=status,63                              errormsg=error)64 65 66@blueprint.route('/verification/',67                 methods=['GET'], endpoint='verification')68@pga_login_required69def verification():70    """Verify Credentials."""71    biganimal = BigAnimalProvider()72    verification_uri = biganimal.get_device_code()73    session['biganimal'] = {}74    session['biganimal']['provider_obj'] = pickle.dumps(biganimal, -1)75 76    return make_json_response(data=verification_uri)77 78 79@blueprint.route('/projects/',80                 methods=['GET'], endpoint='projects')81@pga_login_required82def biganimal_projects():83    """Get Providers."""84    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])85    projects, error = biganimal_obj.get_projects()86    return make_json_response(data=projects, errormsg=error)87 88 89@blueprint.route('/providers/<project_id>',90                 methods=['GET'], endpoint='providers')91@pga_login_required92def biganimal_providers(project_id):93    """Get Providers."""94    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])95    providers, error = biganimal_obj.get_providers(project_id)96    session['biganimal']['provider_obj'] = pickle.dumps(biganimal_obj, -1)97    return make_json_response(data=providers, errormsg=error)98 99 100@blueprint.route('/regions/',101                 methods=['GET'], endpoint='regions')102@pga_login_required103def biganimal_regions():104    """Get Regions."""105    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])106    _, regions = biganimal_obj.get_regions()107    session['biganimal']['provider_obj'] = pickle.dumps(biganimal_obj, -1)108    return make_json_response(data=regions)109 110 111@blueprint.route('/db_types/',112                 methods=['GET'], endpoint='db_types')113@pga_login_required114def biganimal_db_types():115    """Get Database Types."""116    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])117    pg_types = biganimal_obj.get_postgres_types()118    return make_json_response(data=pg_types)119 120 121@blueprint.route('/db_versions/<cluster_type>/<pg_type>',122                 methods=['GET'], endpoint='db_versions')123@pga_login_required124def biganimal_db_versions(cluster_type, pg_type):125    """Get Database Version."""126    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])127    pg_versions = biganimal_obj.get_postgres_versions(cluster_type, pg_type)128    return make_json_response(data=pg_versions)129 130 131@blueprint.route('/instance_types/<region_id>/<provider_id>',132                 methods=['GET'], endpoint='instance_types')133@pga_login_required134def biganimal_instance_types(region_id, provider_id):135    """Get Instance Types."""136    if not region_id or not provider_id:137        return make_json_response(data=[])138    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])139    biganimal_instances = biganimal_obj.get_instance_types(region_id,140                                                           provider_id)141    return make_json_response(data=biganimal_instances)142 143 144@blueprint.route('/volume_types/<region_id>/<provider_id>',145                 methods=['GET'], endpoint='volume_types')146@pga_login_required147def biganimal_volume_types(region_id, provider_id):148    """Get Volume Types."""149    if not region_id or not provider_id:150        return make_json_response(data=[])151    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])152    biganimal_volumes = biganimal_obj.get_volume_types(region_id, provider_id)153    return make_json_response(data=biganimal_volumes)154 155 156@blueprint.route('/volume_properties/<region_id>/<provider_id>/<volume_type>',157                 methods=['GET'], endpoint='volume_properties')158@pga_login_required159def biganimal_volume_properties(region_id, provider_id, volume_type):160    """Get Volume Properties."""161    if not region_id or not provider_id:162        return make_json_response(data=[])163    biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])164    biganimal_volume_properties = biganimal_obj.get_volume_properties(165        region_id,166        provider_id,167        volume_type)168    return make_json_response(data=biganimal_volume_properties)169 170 171class BigAnimalProvider():172    """BigAnimal provider class"""173    BASE_URL = 'https://portal.biganimal.com/api/v3'174 175    def __init__(self):176        self.provider = {}177        self.device_code = {}178        self.token = {}179        self.raw_access_token = None180        self.access_token = None181        self.token_error = {}182        self.token_status = -1183        self.regions = []184        self.get_auth_provider()185        self.project_id = None186 187    def _get_headers(self):188        return {189            'content-type': MIMETYPE_APP_JSON,190            'Authorization': 'Bearer {0}'.format(self.access_token)191        }192 193    def get_auth_provider(self):194        """Get Authentication Provider Relevant Information."""195        provider_resp = requests.get("{0}/{1}".format(self.BASE_URL,196                                                      'auth/provider'))197        if provider_resp.status_code == 200 and provider_resp.content:198            self.provider = json.loads(provider_resp.content)199 200    def get_device_code(self):201        """Get device code"""202        _url = "{0}/{1}".format(self.provider['issuerUri'],203                                'oauth/device/code')204        _headers = {"content-type": "application/x-www-form-urlencoded"}205        _data = {206            'client_id': self.provider['clientId'],207            'audience': self.provider['audience'],208            'scope': self.provider['scope']209        }210        device_resp = requests.post(_url,211                                    headers=_headers,212                                    data=_data)213 214        if device_resp.status_code == 200 and device_resp.content:215            self.device_code = json.loads(device_resp.content)216            return self.device_code['verification_uri_complete']217 218    def polling_for_token(self):219        # Polling for the Token220        _url = "{0}/{1}".format(self.provider['issuerUri'], 'oauth/token')221        _headers = {"content-type": "application/x-www-form-urlencoded"}222        _data = {223            'grant_type': 'urn:ietf:params:oauth:grant-type:device_code',224            'device_code': self.device_code['device_code'],225            'client_id': self.provider['clientId']226        }227        token_resp = requests.post(_url,228                                   headers=_headers,229                                   data=_data)230        if token_resp.status_code == 200:231            self.token = json.loads(token_resp.content)232            self.raw_access_token = self.token['access_token']233            self.token_error['error'] = None234            self.token_status = 1235            status, msg = self.exchange_token()236            if status and not self._check_admin_permission():237                return False, gettext('forbidden')238            return status, msg239        elif token_resp.status_code == 403:240            self.token_error = json.loads(token_resp.content)241            if self.token_error['error'] == 'authorization_pending' or\242                    self.token_error['error'] == 'access_denied':243                self.token_status = 0244                return False, self.token_error['error']245        return False, None246 247    def exchange_token(self):248        _url = "{0}/{1}".format(self.BASE_URL, 'auth/token')249        _headers = {"content-type": "application/json",250                    "accept": "application/json"}251        _data = {'token': self.raw_access_token}252        token_resp = requests.post(_url,253                                   headers=_headers,254                                   data=json.dumps(_data))255 256        final_token = json.loads(token_resp.content)257        if token_resp.status_code == 200:258            self.access_token = final_token['token']259            return True, None260        else:261            return False, self.token_error['error']262 263    def _check_admin_permission(self):264        """265        Check wehether the user has valid role or not.266        There is no direct way to do this, so just checking the create cluster267        permission.268        """269        _url = "{0}/{1}".format(270            self.BASE_URL,271            'user-info')272        resp = requests.get(_url, headers=self._get_headers())273        if resp.status_code != 200:274            return False275        if resp.status_code == 200 and resp.content:276            content = json.loads(resp.content)277            if 'data' in content:278                # BigAnimal introduced Project feature in v3,279                # so all the existing clusters moved to the default Project.280                # For now, we can get the Proj Id by replacing 'org' to 'prj'281                # in organization ID: org_1234  -> prj_1234282                proj_id = content['data']['organizationId'].replace('org',283                                                                    'prj')284                for permission in content['data']['scopedPermissions']:285                    if proj_id == permission['scope'] and\286                            'create:clusters' in permission['permissions']:287                        return True288        return False289 290    def get_providers(self, project_id):291        """Get cloud providers"""292        if not project_id:293            return False, gettext('Project not provided.')294        _url = '{0}/projects/{1}/cloud-providers'.format(295            self.BASE_URL, project_id)296        providers = []297        resp = requests.get(_url, headers=self._get_headers())298        if resp.status_code == 200 and resp.content:299            self.project_id = project_id300            provider_resp = json.loads(resp.content)301            for value in provider_resp['data']:302                providers.append({303                    'label': value['cloudProviderName'],304                    'value': value['cloudProviderId'],305                    'connected': value['connected']})306            return providers, None307        elif resp.content:308            provider_resp = json.loads(resp.content)309            return [], provider_resp['error']['message']310        else:311            return [], gettext('Error retrieving providers.')312 313    def get_regions(self):314        """Get regions"""315        _url = '{0}/projects/{1}/regions'.format(316            self.BASE_URL, self.project_id)317        regions = []318        resp = requests.get(_url, headers=self._get_headers())319        if resp.status_code == 200 and resp.content:320            regions_resp = json.loads(resp.content)321            for value in regions_resp['data']:322                regions.append({323                    'label': value['regionName'],324                    'value': value['regionId']325                })326                self.regions.append(value['regionId'])327            return True, regions328        elif resp.content:329            regions_resp = json.loads(resp.content)330            return False, regions_resp['error']['message']331        else:332            return False, gettext('Error retrieving regions.')333 334    def get_postgres_types(self):335        """Get Postgres Types."""336        _url = "{0}/projects/{1}/pg-types".format(337            self.BASE_URL, self.project_id)338        pg_types = []339        resp = requests.get(_url, headers=self._get_headers())340        if resp.status_code == 200 and resp.content:341            pg_types_resp = json.loads(resp.content)342            for value in pg_types_resp['data']:343                # Extreme HA is in Beta, so avoid it344                if len(value['supportedClusterArchitectureIds']) != 1:345                    pg_types.append({346                        'label': value['pgTypeName'],347                        'value': value['pgTypeId']348                    })349        return pg_types350 351    def get_postgres_versions(self, cluster_type, pg_type):352        """Get Postgres Versions."""353        if not cluster_type or not pg_type:354            return []355 356        _url = "{0}/projects/{1}/pg-versions?clusterArchitectureIds={2}" \357               "&pgTypeIds={3}".format(self.BASE_URL, self.project_id,358                                       cluster_type, pg_type)359        pg_versions = []360        resp = requests.get(_url, headers=self._get_headers())361        if resp.status_code == 200 and resp.content:362            pg_versions_resp = json.loads(resp.content)363            for value in pg_versions_resp['data']:364                pg_versions.append({365                    'label': value['pgVersionName'],366                    'value': value['pgVersionId']367                })368        return pg_versions369 370    def get_instance_types(self, region_id, provider_id):371        """GEt Instance Types."""372        if region_id not in self.regions or not provider_id:373            return []374        _url = '{0}/projects/{1}/cloud-providers/{2}/regions/{3}/' \375               'instance-types?sort=instanceTypeName'.\376            format(self.BASE_URL, self.project_id, provider_id, region_id)377        resp = requests.get(_url, headers=self._get_headers())378        if resp.status_code == 200 and resp.content:379            pg_types = json.loads(resp.content)380            _sorted_data = sorted(pg_types['data'],381                                  key=lambda x: int(x['cpu']))382            return _sorted_data383        return []384 385    def get_volume_types(self, region_id, provider_id):386        """Get Volume Types."""387        if region_id not in self.regions:388            return []389 390        _url = '{0}/projects/{1}/cloud-providers/{2}/regions/{3}/volume-types'\391            .format(self.BASE_URL, self.project_id, provider_id, region_id)392        volume_types = []393        resp = requests.get(_url, headers=self._get_headers())394        if resp.status_code == 200 and resp.content:395            volume_resp = json.loads(resp.content)396            for value in volume_resp['data']:397                if value['enabledInRegion']:398                    volume_types.append({399                        'label': value['volumeTypeName'],400                        'value': value['volumeTypeId'],401                        'supportedInstanceFamilyNames': value[402                            'supportedInstanceFamilyNames']})403        return volume_types404 405    def get_volume_properties(self, region_id, provider_id, volume_type):406        """Get Volume Properties."""407        if region_id not in self.regions:408            return []409 410        _url = '{0}/projects/{1}/cloud-providers/{2}/regions/{3}/' \411               'volume-types/{4}/volume-properties'\412            .format(self.BASE_URL, self.project_id, provider_id, region_id,413                    volume_type)414        volume_properties = []415        resp = requests.get(_url, headers=self._get_headers())416        if resp.status_code == 200 and resp.content:417            volume_prop = json.loads(resp.content)418            for value in volume_prop['data']:419                volume_properties.append({420                    'label': value['volumePropertiesName'],421                    'value': value['volumePropertiesId']422                })423        return volume_properties424 425    def get_projects(self):426        projects = []427        _url = '{0}/projects'.format(self.BASE_URL)428        resp = requests.get(_url, headers=self._get_headers())429        if resp.status_code == 200 and resp.content:430            project_resp = json.loads(resp.content)431            for value in project_resp['data']:432                projects.append({433                    'label': value['projectName'],434                    'value': value['projectId']435                })436            return projects, None437        elif resp.content:438            project_resp = json.loads(resp.content)439            return [], project_resp['error']['message']440        else:441            return [], gettext('Error retrieving projects.')442 443 444def clear_biganimal_session():445    """Clear session data."""446    if 'biganimal' in session:447        session.pop('biganimal')448 449 450def deploy_on_biganimal(data):451    """Deploy Postgres instance on BigAnimal"""452    _cmd = 'python'453    _cmd_script = '{0}/pgacloud/pgacloud.py'.format(root)454    _label = data['instance_details']['name']455    _private_network = '1' if str(data['instance_details']['cloud_type']456                                  ) == 'private' else '0'457    _instance_size = data['instance_details']['instance_size'].split('||')[1]458    nodes = 1459 460    if data['cluster_details']['cluster_type'] == HA_CLUSTER_ARCH:461        nodes = int(data['cluster_details']['replicas']) + nodes462    elif data['cluster_details']['cluster_type'] == EHA_CLUSTER_ARCH:463        nodes = 5464 465    args = [_cmd_script,466            data['cloud'],467            'create-instance',468            '--name',469            data['instance_details']['name'],470            '--project',471            str(data['cluster_details']['project']),472            '--cloud-provider',473            str(data['cluster_details']['provider']),474            '--region',475            str(data['instance_details']['region']),476            '--db-type',477            str(data['db_details']['database_type']),478            '--db-version',479            str(data['db_details']['postgres_version']),480            '--volume-type',481            str(data['instance_details']['volume_type']),482            '--volume-properties',483            str(data['instance_details'].get('volume_properties',484                                             data['instance_details'][485                                                 'volume_type'])),486            '--volume-size',487            str(data['instance_details'].get('volume_size', None)),488            '--volume-IOPS',489            str(data['instance_details'].get('volume_IOPS', None)),490            '--throughput',491            str(data['instance_details'].get('disk_throughput', None)),492            '--instance-type',493            str(_instance_size),494            '--private-network',495            _private_network,496            '--cluster-arch',497            data['cluster_details']['cluster_type'],498            '--nodes',499            str(nodes),500            '--replicas',501            str(data['cluster_details']['replicas'])]502 503    if 'biganimal_public_ip' in data['instance_details']:504        args.append('--public-ip')505        args.append(str(data['instance_details']['biganimal_public_ip']))506 507    _cmd_msg = '{0} {1} {2}'.format(_cmd, _cmd_script, ' '.join(args))508    try:509        sid = _create_server({510            'gid': data['db_details']['gid'],511            'name': data['instance_details']['name'],512            'db': 'edb_admin',513            'username': 'edb_admin',514            'port': 5432,515            'cloud_status': -1516        })517 518        p = BatchProcess(519            desc=CloudProcessDesc(sid, _cmd_msg,520                                  data['cloud'],521                                  data['instance_details']['name']522                                  ),523            cmd=_cmd,524            args=args525        )526 527        env = dict()528        biganimal_obj = pickle.loads(session['biganimal']['provider_obj'])529        env['BIGANIMAL_ACCESS_KEY'] = biganimal_obj.access_token530 531        if 'password' in data['db_details']:532            env['BIGANIMAL_DATABASE_PASSWORD'] = data[533                'db_details']['password']534 535        p.set_env_variables(None, env=env)536        p.update_server_id(p.id, sid)537        p.start()538 539        return True, p, {'label': _label, 'sid': sid}540 541    except Exception as e:542        current_app.logger.exception(e)543        return False, None, str(e)544 
codekingpro/portable-devtools · Team Ai