<commit_before># --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------
# pylint: disable=line-too-long

from collections import OrderedDict
import sys
import os
import tempfile
import shutil
import zipfile
import traceback
import hashlib
from subprocess import check_output, STDOUT, CalledProcessError
from urllib.parse import urlparse

from packaging.version import parse

from azure.cli.core import CommandIndex
from azure.cli.core.util import CLIError, reload_module, rmtree_with_retry
from azure.cli.core.extension import (extension_exists, build_extension_path, get_extensions, get_extension_modname,
                                      get_extension, ext_compat_with_cli,
                                      EXT_METADATA_ISPREVIEW, EXT_METADATA_ISEXPERIMENTAL,
                                      WheelExtension, DevExtension, ExtensionNotInstalledException, WHEEL_INFO_RE)
from azure.cli.core.telemetry import set_extension_management_detail

from knack.log import get_logger

from ._homebrew_patch import HomebrewPipPatch
from ._index import get_index_extensions
from ._resolve import resolve_from_index, NoExtensionCandidatesError

logger = get_logger(__name__)

OUT_KEY_NAME = 'name'
OUT_KEY_VERSION = 'version'
OUT_KEY_TYPE = 'extensionType'
OUT_KEY_METADATA = 'metadata'
OUT_KEY_PREVIEW = 'preview'
OUT_KEY_EXPERIMENTAL = 'experimental'
OUT_KEY_PATH = 'path'

IS_WINDOWS = sys.platform.lower() in ['windows', 'win32']
LIST_FILE_PATH = os.path.join(os.sep, 'etc', 'apt', 'sources.list.d', 'azure-cli.list')
LSB_RELEASE_FILE = os.path.join(os.sep, 'etc', 'lsb-release')


def _run_pip(pip_exec_args, extension_path=None):
    cmd = [sys.executable, '-m', 'pip'] + pip_exec_args + ['-vv', '--disable-pip-version-check', '--no-cache-dir']
    logger.debug('Running: %s', cmd)
    try:
        log_output = check_output(cmd, stderr=STDOUT, universal_newlines=True)
        logger.debug(log_output)
        returncode = 0
    except CalledProcessError as e:
        logger.debug(e.output)
        logger.debug(e)
        if "PermissionError: [WinError 5]" in e.output:
            logger.warning("You do not have the permission to add extensions in the target directory%s. You may need to rerun on a shell as administrator.", ': ' + os.path.split(extension_path)[0] if extension_path else '')
        returncode = e.returncode
    return returncode


def _whl_download_from_url(url_parse_result, ext_file):
    import requests
    from azure.cli.core.util import should_disable_connection_verify
    url = url_parse_result.geturl()
    r = requests.get(url, stream=True, verify=(not should_disable_connection_verify()))
    if r.status_code != 200:
        raise CLIError("Request to {} failed with {}".format(url, r.status_code))
    with open(ext_file, 'wb') as f:
        for chunk in r.iter_content(chunk_size=1024):
            if chunk:  # ignore keep-alive new chunks
                f.write(chunk)


def _validate_whl_extension(ext_file):
    tmp_dir = tempfile.mkdtemp()
    zip_ref = zipfile.ZipFile(ext_file, 'r')
    zip_ref.extractall(tmp_dir)
    zip_ref.close()
    azext_metadata = WheelExtension.get_azext_metadata(tmp_dir)
    rmtree_with_retry(tmp_dir)
    check_version_compatibility(azext_metadata)


def _get_extension_info_from_source(source):
    url_parse_result = urlparse(source)
    is_url = (url_parse_result.scheme == 'http' or url_parse_result.scheme == 'https')
    whl_filename = os.path.basename(url_parse_result.path) if is_url else os.path.basename(source)
    parsed_filename = WHEEL_INFO_RE(whl_filename)
    # Extension names can have - but .whl format changes it to _ (PEP 0427). Undo this.
    extension_name = parsed_filename.groupdict().get('name').replace('_', '-') if parsed_filename else None
    extension_version = parsed_filename.groupdict().get('ver') if parsed_filename else None
    return extension_name, extension_version


def _add_whl_ext(cli_ctx, source, ext_sha256=None, pip_extra_index_urls=None, pip_proxy=None, system=None):  # pylint: disable=too-many-statements
    cli_ctx.get_progress_controller().add(message='Analyzing')
    if not source.endswith('.whl'):
        raise ValueError('Unknown extension type. Only Python wheels are supported.')
    url_parse_result = urlparse(source)
    is_url = (url_parse_result.scheme == 'http' or url_parse_result.scheme == 'https')
    logger.debug('Extension source is url? %s', is_url)
    whl_filename = os.path.basename(url_parse_result.path) if is_url else os.path.basename(source)
    parsed_filename = WHEEL_INFO_RE(whl_filename)
    # Extension names can have - but .whl format changes it to _ (PEP 0427). Undo this.
    extension_name = parsed_filename.groupdict().get('name').replace('_', '-') if parsed_filename else None
    if not extension_name:
        raise CLIError('Unable to determine extension name from {}. Is the file name correct?'.format(source))
    if extension_exists(extension_name, ext_type=WheelExtension):
        raise CLIError('The extension {} already exists.'.format(extension_name))
    if extension_name == 'rdbms-connect':
        _install_deps_for_psycopg2()
    ext_file = None
    if is_url:
        # Download from URL
        tmp_dir = tempfile.mkdtemp()
        ext_file = os.path.join(tmp_dir, whl_filename)
        logger.debug('Downloading %s to %s', source, ext_file)
        import requests
        try:
            cli_ctx.get_progress_controller().add(message='Downloading')
            _whl_download_from_url(url_parse_result, ext_file)
        except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError) as err:
            raise CLIError('Please ensure you have network connection. Error detail: {}'.format(str(err)))
        logger.debug('Downloaded to %s', ext_file)
    else:
        # Get file path
        ext_file = os.path.realpath(os.path.expanduser(source))
        if not os.path.isfile(ext_file):
            raise CLIError("File {} not found.".format(source))
    # Validate the extension
    logger.debug('Validating the extension %s', ext_file)
    if ext_sha256:
        valid_checksum, computed_checksum = is_valid_sha256sum(ext_file, ext_sha256)
        if valid_checksum:
            logger.debug("Checksum of %s is OK", ext_file)
        else:
            logger.debug("Invalid checksum for %s. Expected '%s', computed '%s'.",
                         ext_file, ext_sha256, computed_checksum)
            raise CLIError("The checksum of the extension does not match the expected value. "
                           "Use --debug for more information.")
    try:
        cli_ctx.get_progress_controller().add(message='Validating')
        _validate_whl_extension(ext_file)
    except AssertionError:
        logger.debug(traceback.format_exc())
        raise CLIError('The extension is invalid. Use --debug for more information.')
    except CLIError as e:
        raise e
    logger.debug('Validation successful on %s', ext_file)
    # Check for distro consistency
    check_distro_consistency()
    cli_ctx.get_progress_controller().add(message='Installing')
    # Install with pip
    extension_path = build_extension_path(extension_name, system)
    pip_args = ['install', '--target', extension_path, ext_file]

    if pip_proxy:
        pip_args = pip_args + ['--proxy', pip_proxy]
    if pip_extra_index_urls:
        for extra_index_url in pip_extra_index_urls:
            pip_args = pip_args + ['--extra-index-url', extra_index_url]

    logger.debug('Executing pip with args: %s', pip_args)
    with HomebrewPipPatch():
        pip_status_code = _run_pip(pip_args, extension_path)
    if pip_status_code > 0:
        logger.debug('Pip failed so deleting anything we might have installed at %s', extension_path)
        rmtree_with_retry(extension_path)
        raise CLIError('An error occurred. Pip failed with status code {}. '
                       'Use --debug for more information.'.format(pip_status_code))
    # Save the whl we used to install the extension in the extension dir.
    dst = os.path.join(extension_path, whl_filename)
    shutil.copyfile(ext_file, dst)
    logger.debug('Saved the whl to %s', dst)

    return extension_name


def _install_deps_for_psycopg2():  # pylint: disable=too-many-statements
    # Below system dependencies are required to install the psycopg2 dependency for Linux and macOS
    import platform
    import subprocess
    from azure.cli.core.util import get_linux_distro
    from azure.cli.core._environment import _ENV_AZ_INSTALLER
    installer = os.getenv(_ENV_AZ_INSTALLER)
    system = platform.system()
    if system == 'Darwin':
        subprocess.call(['xcode-select', '--install'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if installer != 'HOMEBREW':
            from shutil import which
            if which('brew') is None:
                logger.warning('You may need to install postgresql with homebrew first before you install this extension.')
                return
        exit_code = subprocess.call(['brew', 'list', 'postgresql'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if exit_code != 0:
            update_cmd = ['brew', 'install', 'postgresql']
            logger.warning('This extension depends on postgresql and it will be installed first.')
            logger.debug("Install dependencies with '%s'", " ".join(update_cmd))
            subprocess.call(update_cmd)
        # Fix the issue of -lssl not found during building psycopg2
        if os.environ.get('LIBRARY_PATH') is None:
            os.environ['LIBRARY_PATH'] = '/usr/local/opt/openssl/lib/'<commit_msg>[RDBMS] Fix operations.py file installing dependencies in CloudShell (#21553)<commit_after># --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------
# pylint: disable=line-too-long

from collections import OrderedDict
import sys
import os
import tempfile
import shutil
import zipfile
import traceback
import hashlib
from subprocess import check_output, STDOUT, CalledProcessError
from urllib.parse import urlparse

from packaging.version import parse

from azure.cli.core import CommandIndex
from azure.cli.core.util import CLIError, reload_module, rmtree_with_retry
from azure.cli.core.extension import (extension_exists, build_extension_path, get_extensions, get_extension_modname,
                                      get_extension, ext_compat_with_cli,
                                      EXT_METADATA_ISPREVIEW, EXT_METADATA_ISEXPERIMENTAL,
                                      WheelExtension, DevExtension, ExtensionNotInstalledException, WHEEL_INFO_RE)
from azure.cli.core.telemetry import set_extension_management_detail

from knack.log import get_logger

from ._homebrew_patch import HomebrewPipPatch
from ._index import get_index_extensions
from ._resolve import resolve_from_index, NoExtensionCandidatesError

logger = get_logger(__name__)

OUT_KEY_NAME = 'name'
OUT_KEY_VERSION = 'version'
OUT_KEY_TYPE = 'extensionType'
OUT_KEY_METADATA = 'metadata'
OUT_KEY_PREVIEW = 'preview'
OUT_KEY_EXPERIMENTAL = 'experimental'
OUT_KEY_PATH = 'path'

IS_WINDOWS = sys.platform.lower() in ['windows', 'win32']
LIST_FILE_PATH = os.path.join(os.sep, 'etc', 'apt', 'sources.list.d', 'azure-cli.list')
LSB_RELEASE_FILE = os.path.join(os.sep, 'etc', 'lsb-release')


def _run_pip(pip_exec_args, extension_path=None):
    cmd = [sys.executable, '-m', 'pip'] + pip_exec_args + ['-vv', '--disable-pip-version-check', '--no-cache-dir']
    logger.debug('Running: %s', cmd)
    try:
        log_output = check_output(cmd, stderr=STDOUT, universal_newlines=True)
        logger.debug(log_output)
        returncode = 0
    except CalledProcessError as e:
        logger.debug(e.output)
        logger.debug(e)
        if "PermissionError: [WinError 5]" in e.output:
            logger.warning("You do not have the permission to add extensions in the target directory%s. You may need to rerun on a shell as administrator.", ': ' + os.path.split(extension_path)[0] if extension_path else '')
        returncode = e.returncode
    return returncode


def _whl_download_from_url(url_parse_result, ext_file):
    import requests
    from azure.cli.core.util import should_disable_connection_verify
    url = url_parse_result.geturl()
    r = requests.get(url, stream=True, verify=(not should_disable_connection_verify()))
    if r.status_code != 200:
        raise CLIError("Request to {} failed with {}".format(url, r.status_code))
    with open(ext_file, 'wb') as f:
        for chunk in r.iter_content(chunk_size=1024):
            if chunk:  # ignore keep-alive new chunks
                f.write(chunk)


def _validate_whl_extension(ext_file):
    tmp_dir = tempfile.mkdtemp()
    zip_ref = zipfile.ZipFile(ext_file, 'r')
    zip_ref.extractall(tmp_dir)
    zip_ref.close()
    azext_metadata = WheelExtension.get_azext_metadata(tmp_dir)
    rmtree_with_retry(tmp_dir)
    check_version_compatibility(azext_metadata)


def _get_extension_info_from_source(source):
    url_parse_result = urlparse(source)
    is_url = (url_parse_result.scheme == 'http' or url_parse_result.scheme == 'https')
    whl_filename = os.path.basename(url_parse_result.path) if is_url else os.path.basename(source)
    parsed_filename = WHEEL_INFO_RE(whl_filename)
    # Extension names can have - but .whl format changes it to _ (PEP 0427). Undo this.
    extension_name = parsed_filename.groupdict().get('name').replace('_', '-') if parsed_filename else None
    extension_version = parsed_filename.groupdict().get('ver') if parsed_filename else None
    return extension_name, extension_version


def _add_whl_ext(cli_ctx, source, ext_sha256=None, pip_extra_index_urls=None, pip_proxy=None, system=None):  # pylint: disable=too-many-statements
    cli_ctx.get_progress_controller().add(message='Analyzing')
    if not source.endswith('.whl'):
        raise ValueError('Unknown extension type. Only Python wheels are supported.')
    url_parse_result = urlparse(source)
    is_url = (url_parse_result.scheme == 'http' or url_parse_result.scheme == 'https')
    logger.debug('Extension source is url? %s', is_url)
    whl_filename = os.path.basename(url_parse_result.path) if is_url else os.path.basename(source)
    parsed_filename = WHEEL_INFO_RE(whl_filename)
    # Extension names can have - but .whl format changes it to _ (PEP 0427). Undo this.
    extension_name = parsed_filename.groupdict().get('name').replace('_', '-') if parsed_filename else None
    if not extension_name:
        raise CLIError('Unable to determine extension name from {}. Is the file name correct?'.format(source))
    if extension_exists(extension_name, ext_type=WheelExtension):
        raise CLIError('The extension {} already exists.'.format(extension_name))
    if extension_name == 'rdbms-connect':
        _install_deps_for_psycopg2()
    ext_file = None
    if is_url:
        # Download from URL
        tmp_dir = tempfile.mkdtemp()
        ext_file = os.path.join(tmp_dir, whl_filename)
        logger.debug('Downloading %s to %s', source, ext_file)
        import requests
        try:
            cli_ctx.get_progress_controller().add(message='Downloading')
            _whl_download_from_url(url_parse_result, ext_file)
        except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError) as err:
            raise CLIError('Please ensure you have network connection. Error detail: {}'.format(str(err)))
        logger.debug('Downloaded to %s', ext_file)
    else:
        # Get file path
        ext_file = os.path.realpath(os.path.expanduser(source))
        if not os.path.isfile(ext_file):
            raise CLIError("File {} not found.".format(source))
    # Validate the extension
    logger.debug('Validating the extension %s', ext_file)
    if ext_sha256:
        valid_checksum, computed_checksum = is_valid_sha256sum(ext_file, ext_sha256)
        if valid_checksum:
            logger.debug("Checksum of %s is OK", ext_file)
        else:
            logger.debug("Invalid checksum for %s. Expected '%s', computed '%s'.",
                         ext_file, ext_sha256, computed_checksum)
            raise CLIError("The checksum of the extension does not match the expected value. "
                           "Use --debug for more information.")
    try:
        cli_ctx.get_progress_controller().add(message='Validating')
        _validate_whl_extension(ext_file)
    except AssertionError:
        logger.debug(traceback.format_exc())
        raise CLIError('The extension is invalid. Use --debug for more information.')
    except CLIError as e:
        raise e
    logger.debug('Validation successful on %s', ext_file)
    # Check for distro consistency
    check_distro_consistency()
    cli_ctx.get_progress_controller().add(message='Installing')
    # Install with pip
    extension_path = build_extension_path(extension_name, system)
    pip_args = ['install', '--target', extension_path, ext_file]

    if pip_proxy:
        pip_args = pip_args + ['--proxy', pip_proxy]
    if pip_extra_index_urls:
        for extra_index_url in pip_extra_index_urls:
            pip_args = pip_args + ['--extra-index-url', extra_index_url]

    logger.debug('Executing pip with args: %s', pip_args)
    with HomebrewPipPatch():
        pip_status_code = _run_pip(pip_args, extension_path)
    if pip_status_code > 0:
        logger.debug('Pip failed so deleting anything we might have installed at %s', extension_path)
        rmtree_with_retry(extension_path)
        raise CLIError('An error occurred. Pip failed with status code {}. '
                       'Use --debug for more information.'.format(pip_status_code))
    # Save the whl we used to install the extension in the extension dir.
    dst = os.path.join(extension_path, whl_filename)
    shutil.copyfile(ext_file, dst)
    logger.debug('Saved the whl to %s', dst)

    return extension_name


def _install_deps_for_psycopg2():  # pylint: disable=too-many-statements
    # If we are in Cloud Shell, dependencies should have already been installed.
    from azure.cli.core.util import in_cloud_console
    if in_cloud_console():
        return

    # Below system dependencies are required to install the psycopg2 dependency for Linux and macOS
    import platform
    import subprocess
    from azure.cli.core.util import get_linux_distro
    from azure.cli.core._environment import _ENV_AZ_INSTALLER
    installer = os.getenv(_ENV_AZ_INSTALLER)
    system = platform.system()
    if system == 'Darwin':
        subprocess.call(['xcode-select', '--install'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if installer != 'HOMEBREW':
            from shutil import which
            if which('brew') is None:
                logger.warning('You may need to install postgresql with homebrew first before you install this extension.')
                return
        exit_code = subprocess.call(['brew', 'list', 'postgresql'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if exit_code != 0:
            update_cmd = ['brew', 'install', 'postgresql']
            logger.warning('This extension depends on postgresql and it will be installed first.')
            logger.debug("Install dependencies with '%s'", " ".join(update_cmd))
            subprocess.call(update_cmd)
        # Fix the issue of -lssl not found during building psycopg2
        if os.environ.get('LIBRARY_PATH') is None:
            os.environ['LIBRARY_PATH'] = '/usr/local/opt/openssl/lib/'
{'fraction_non_alphanumeric': 0.11312171957010747, 'fraction_numerical': 0.004348912771807048, 'mean_word_length': 3.2126763529164033, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4667150', 'n_tokens_mistral': 6041, 'n_tokens_neox': 5858, 'n_words': 1502}