text: # (C) Datadog, Inc. 2020-present
  # All rights reserved
  # Licensed under a 3-clause BSD style license (see LICENSE)
  
  from contextlib import contextmanager
  
  from six import raise_from
  
  from datadog_checks.base import AgentCheck
  
  try:
      import adodbapi
  except ImportError:
      adodbapi = None
  
  try:
      import pyodbc
  except ImportError:
      pyodbc = None
  
  DATABASE_EXISTS_QUERY = 'select name, collation_name from sys.databases;'
  
  
  class SQLConnectionError(Exception):
      """Exception raised for SQL instance connection issues"""
  
      pass
  
  
  class Connection(object):
      """Manages the connection to a SQL Server instance."""
  
      DEFAULT_COMMAND_TIMEOUT = 5
      DEFAULT_DATABASE = 'master'
      DEFAULT_DRIVER = 'SQL Server'
      DEFAULT_DB_KEY = 'database'
      PROC_GUARD_DB_KEY = 'proc_only_if_database'
  
      valid_adoproviders = ['SQLOLEDB', 'MSOLEDBSQL', 'SQLNCLI11']
      default_adoprovider = 'SQLOLEDB'
  
      def __init__(self, init_config, instance_config, service_check_handler, logger):
          self.instance = instance_config
          self.service_check_handler = service_check_handler
          self.log = logger
  
          # mapping of raw connections based on conn_key to different databases
          self._conns = {}
          self.timeout = int(self.instance.get('command_timeout', self.DEFAULT_COMMAND_TIMEOUT))
          self.existing_databases = None
  
          self.adoprovider = self.default_adoprovider
  
          self.valid_connectors = []
          if adodbapi is not None:
              self.valid_connectors.append('adodbapi')
          if pyodbc is not None:
              self.valid_connectors.append('odbc')
  
          self.connector = init_config.get('connector', 'adodbapi')
          if self.connector.lower() not in self.valid_connectors:
              self.log.error("Invalid database connector %s, defaulting to adodbapi", self.connector)
              self.connector = 'adodbapi'
  
          self.adoprovider = init_config.get('adoprovider', self.default_adoprovider)
          if self.adoprovider.upper() not in self.valid_adoproviders:
              self.log.error(
                  "Invalid ADODB provider string %s, defaulting to %s", self.adoprovider, self.default_adoprovider
              )
              self.adoprovider = self.default_adoprovider
  
          self.log.debug('Connection initialized.')
  
      @contextmanager
      def get_managed_cursor(self):
          cursor = self.get_cursor(self.DEFAULT_DB_KEY)
          yield cursor
          self.close_cursor(cursor)
  
      def get_cursor(self, db_key, db_name=None):
          """
          Return a cursor to execute query against the db
          Cursor are cached in the self.connections dict
          """
          conn_key = self._conn_key(db_key, db_name)
          try:
              conn = self._conns[conn_key]
          except KeyError:
              # We catch KeyError to avoid leaking the auth info used to compose the key
              # FIXME: we should find a better way to compute unique keys to map opened connections other than
              # using auth info in clear text!
              raise SQLConnectionError("Cannot find an opened connection for host: {}".format(self.instance.get('host')))
          return conn.cursor()
  
      def close_cursor(self, cursor):
          """
          We close the cursor explicitly b/c we had proven memory leaks
          We handle any exception from closing, although according to the doc:
          "in adodbapi, it is NOT an error to re-close a closed cursor"
          """
          try:
              cursor.close()
          except Exception as e:
              self.log.warning("Could not close adodbapi cursor\n%s", e)
  
      def check_database(self):
          with self.open_managed_default_database():
              db_exists, context = self._check_db_exists()
  
          return db_exists, context
  
      @contextmanager
      def open_managed_default_database(self):
          with self._open_managed_db_connections(None, db_name=self.DEFAULT_DATABASE):
              yield
  
      @contextmanager
      def open_managed_default_connection(self):
          with self._open_managed_db_connections(self.DEFAULT_DB_KEY):
              yield
  
      @contextmanager
      def _open_managed_db_connections(self, db_key, db_name=None):
          self.open_db_connections(db_key, db_name)
          yield
          self.close_db_connections(db_key, db_name)
  
      def open_db_connections(self, db_key, db_name=None):
          """
          We open the db connections explicitly, so we can ensure they are open
          before we use them, and are closable, once we are finished. Open db
          connections keep locks on the db, presenting issues such as the SQL
          Server Agent being unable to stop.
          """
  
          conn_key = self._conn_key(db_key, db_name)
  
          dsn, host, username, password, database, driver = self._get_access_info(db_key, db_name)
  
          cs = self.instance.get('connection_string', '')
          cs += ';' if cs != '' else ''
  
          try:
              if self.get_connector() == 'adodbapi':
                  cs += self._conn_string_adodbapi(db_key, db_name=db_name)
                  # autocommit: true disables implicit transaction
                  rawconn = adodbapi.connect(cs, {'timeout': self.timeout, 'autocommit': True})
              else:
                  cs += self._conn_string_odbc(db_key, db_name=db_name)
                  rawconn = pyodbc.connect(cs, timeout=self.timeout)
  
              self.service_check_handler(AgentCheck.OK, host, database)
  
              if conn_key not in self._conns:
                  self._conns[conn_key] = rawconn
              else:
                  try:
                      # explicitly trying to avoid leaks...
                      self._conns[conn_key].close()
                  except Exception as e:
                      self.log.info("Could not close adodbapi db connection\n%s", e)
  
                  self._conns[conn_key] = rawconn
          except Exception as e:
              cx = "{} - {}".format(host, database)
              message = "Unable to connect to SQL Server for instance {}: {}".format(cx, repr(e))
  
              password = self.instance.get('password')
              if password is not None:
                  message = message.replace(password, "*" * 6)
  
              self.service_check_handler(AgentCheck.CRITICAL, host, database, message)
  
              raise_from(SQLConnectionError(message), None)
  
      def close_db_connections(self, db_key, db_name=None):
          """
          We close the db connections explicitly b/c when we don't they keep
          locks on the db. This presents as issues such as the SQL Server Agent
          being unable to stop.
          """
          conn_key = self._conn_key(db_key, db_name)
          if conn_key not in self._conns:
              return
  
          try:
              self._conns[conn_key].close()
              del self._conns[conn_key]
          except Exception as e:
              self.log.warning("Could not close adodbapi db connection\n%s", e)
  
      def _check_db_exists(self):
          """
          Check for existence of a database, but take into consideration whether the db is case-sensitive or not.
  
          If not case-sensitive, then we normalize the database name to lowercase on both sides and check.
          If case-sensitive, then we only accept exact-name matches.
  
          If the check fails, then we won't do any checks if `ignore_missing_database` is enabled, or we will fail
          with a ConfigurationError otherwise.
          """
  
          dsn, host, username, password, database, driver = self._get_access_info(self.DEFAULT_DB_KEY)
          context = "{} - {}".format(host, database)
          if self.existing_databases is None:
              cursor = self.get_cursor(None, self.DEFAULT_DATABASE)
  
              try:
                  self.existing_databases = {}
                  cursor.execute(DATABASE_EXISTS_QUERY)
                  for row in cursor:
                      case_insensitive = 'CI' in row.collation_name
                      self.existing_databases[row.name.lower()] = case_insensitive, row.name
  
              except Exception as e:
                  self.log.error("Failed to check if database %s exists: %s", database, e)
                  return False, context
              finally:
                  self.close_cursor(cursor)
  
          exists = False
          if database.lower() in self.existing_databases:
              case_insensitive, cased_name = self.existing_databases[database.lower()]
              if case_insensitive or database == cased_name:
                  exists = True
  
          return exists, context
  
      def get_connector(self):
          connector = self.instance.get('connector', self.connector)
          if connector != self.connector:
              if connector.lower() not in self.valid_connectors:
                  self.log.warning("Invalid database connector %s using default %s", connector, self.connector)
                  connector = self.connector
              else:
                  self.log.debug("Overriding default connector for %s with %s", self.instance['host'], connector)
          return connector
  
      def _get_adoprovider(self):
          provider = self.instance.get('adoprovider', self.default_adoprovider)
          if provider != self.adoprovider:
              if provider.upper() not in self.valid_adoproviders:
                  self.log.warning("Invalid ADO provider %s using default %s", provider, self.adoprovider)
                  provider = self.adoprovider
              else:
                  self.log.debug("Overriding default ADO provider for %s with %s", self.instance['host'], provider)
          return provider
  
      def _get_access_info(self, db_key, db_name=None):
          """Convenience method to extract info from instance"""
          dsn = self.instance.get('dsn')
          host = self.instance.get('host')
          username = self.instance.get('username')
          password = self.instance.get('password')
          database = self.instance.get(db_key) if db_name is None else db_name
          driver = self.instance.get('driver')
          if not dsn:
              if not host:
                  host = '127.0.0.1,1433'
              if not database:
                  database = self.DEFAULT_DATABASE
              if not driver:
                  driver = self.DEFAULT_DRIVER
          return dsn, host, username, password, database, driver
  
      def _conn_key(self, db_key, db_name=None):
          """Return a key to use for the connection cache"""
          dsn, host, username, password, database, driver = self._get_access_info(db_key, db_name)
          return '{}:{}:{}:{}:{}:{}'.format(dsn, host, username, password, database, driver)
  
      def _conn_string_odbc(self, db_key, conn_key=None, db_name=None):
          """Return a connection string to use with odbc"""
          if conn_key:
              dsn, host, username, password, database, driver = conn_key.split(":")
          else:
              dsn, host, username, password, database, driver = self._get_access_info(db_key, db_name)
  
          conn_str = ''
          if dsn:
              conn_str = 'DSN={};'.format(dsn)
  
          if driver:
              conn_str += 'DRIVER={};'.format(driver)
          if host:
              conn_str += 'Server={};'.format(host)
          if database:
              conn_str += 'Database={};'.format(database)
  
          if username:
              conn_str += 'UID={};'.format(username)
          self.log.debug("Connection string (before password) %s", conn_str)
          if password:
              conn_str += 'PWD={};'.format(password)
          return conn_str
  
      def _conn_string_adodbapi(self, db_key, conn_key=None, db_name=None):
          """Return a connection string to use with adodbapi"""
          if conn_key:
              _, host, username, password, database, _ = conn_key.split(":")
          else:
              _, host, username, password, database, _ = self._get_access_info(db_key, db_name)
  
          provider = self._get_adoprovider()
          conn_str = 'Provider={};Data Source={};Initial Catalog={};'.format(provider, host, database)
  
          if username:
              conn_str += 'User ID={};'.format(username)
          self.log.debug("Connection string (before password) %s", conn_str)
          if password:
              conn_str += 'Password={};'.format(password)
          if not username and not password:
              conn_str += 'Integrated Security=SSPI;'
          return conn_str
  