<commit_before>import collections
import re
import threading
from sqlalchemy.dialects.postgresql.base import PGDialect
from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
from sqlalchemy.dialects.postgresql import INET
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.util import warn
import sqlalchemy.sql as sql

import sqlalchemy.types as sqltypes

from .stmt_compiler import CockroachCompiler, CockroachIdentifierPreparer

# Map type names (as returned by information_schema) to sqlalchemy type
# objects.
#
# TODO(bdarnell): test more of these. The stock test suite only covers
# a few basic ones.
_type_map = {
    'bool': sqltypes.BOOLEAN,  # introspection returns "BOOL" not boolean
    'boolean': sqltypes.BOOLEAN,

    'bigint': sqltypes.INT,
    'int': sqltypes.INT,
    'int2': sqltypes.INT,
    'int4': sqltypes.INT,
    'int64': sqltypes.INT,
    'int8': sqltypes.INT,
    'integer': sqltypes.INT,
    'smallint': sqltypes.INT,

    'double precision': sqltypes.FLOAT,
    'float': sqltypes.FLOAT,
    'float4': sqltypes.FLOAT,
    'float8': sqltypes.FLOAT,
    'real': sqltypes.FLOAT,

    'dec': sqltypes.DECIMAL,
    'decimal': sqltypes.DECIMAL,
    'numeric': sqltypes.DECIMAL,

    'date': sqltypes.DATE,

    'time': sqltypes.Time,
    'time without time zone': sqltypes.Time,

    'timestamp': sqltypes.TIMESTAMP,
    'timestamptz': sqltypes.TIMESTAMP,
    'timestamp with time zone': sqltypes.TIMESTAMP,
    'timestamp without time zone': sqltypes.TIMESTAMP,

    'interval': sqltypes.Interval,

    'char': sqltypes.VARCHAR,
    'char varying': sqltypes.VARCHAR,
    'character': sqltypes.VARCHAR,
    'character varying': sqltypes.VARCHAR,
    'string': sqltypes.VARCHAR,
    'text': sqltypes.Text,
    'varchar': sqltypes.VARCHAR,

    'blob': sqltypes.BLOB,
    'bytea': sqltypes.BLOB,
    'bytes': sqltypes.BLOB,

    'json': sqltypes.JSON,
    'jsonb': sqltypes.JSON,

    'uuid': UUID,

    'inet': INET,
}


class _SavepointState(threading.local):
    """Hack to override names used in savepoint statements.

    To get the Session to do the right thing with transaction retries,
    we use the begin_nested() method, which executes a savepoint. We
    need to transform the savepoint statements that are a part of this
    retry loop, while leaving other savepoints alone. Unfortunately
    the interface leaves us with no way to pass this information along
    except via a thread-local variable.
    """
    def __init__(self):
        self.cockroach_restart = False


savepoint_state = _SavepointState()


class CockroachDBDialect(PGDialect_psycopg2):
    name = 'cockroachdb'
    supports_comments = False
    supports_sequences = False
    statement_compiler = CockroachCompiler
    preparer = CockroachIdentifierPreparer

    def __init__(self, *args, **kwargs):
        super(CockroachDBDialect, self).__init__(*args,
                                                 use_native_hstore=False,
                                                 server_side_cursors=False,
                                                 **kwargs)

    def initialize(self, connection):
        # Bypass PGDialect's initialize implementation, which looks at
        # server_version_info and performs postgres-specific queries
        # to detect certain features on the server. Set the attributes
        # by hand and hope things don't change out from under us too
        # often.
        super(PGDialect, self).initialize(connection)
        self.implicit_returning = True
        self.supports_native_enum = False
        self.supports_smallserial = False
        self._backslash_escapes = False
        self._has_native_jsonb = True
        sversion = connection.scalar("select version()")
        self._is_v2plus = " v2." in sversion
        self._is_v21plus = self._is_v2plus and (" v2.0." not in sversion)
        self._has_native_json = self._is_v2plus
        self._has_native_jsonb = self._is_v2plus

    def _get_server_version_info(self, conn):
        # PGDialect expects a postgres server version number here,
        # although we've overridden most of the places where it's
        # used.
        return (9, 5, 0)

    def get_table_names(self, conn, schema=None, **kw):
        # Upstream implementation needs correlated subqueries.

        if not self._is_v2plus:
            # v1.1 or earlier.
            return [row.Table for row in conn.execute("SHOW TABLES")]

        # v2.0+ have a good information schema. Use it.
        return [row.table_name for row in conn.execute(
            "SELECT table_name FROM information_schema.tables WHERE table_schema=%s",
            (schema or self.default_schema_name,))]

    def has_table(self, conn, table, schema=None):
        # Upstream implementation needs pg_table_is_visible().
        return any(t == table for t in self.get_table_names(conn, schema=schema))

    # The upstream implementations of the reflection functions below depend on
    # correlated subqueries which are not yet supported.
    def get_columns(self, conn, table_name, schema=None, **kw):
        if not self._is_v2plus:
            # v1.1.
            # Bad: the table name is not properly escaped.
            # Oh well. Hoping 1.1 won't be around for long.
            rows = conn.execute('SHOW COLUMNS FROM "%s"."%s"' %
                                (schema or self.default_schema_name, table_name))
        else:
            # v2.0 or later. Information schema is usable.
            rows = conn.execute(
                'SELECT column_name, data_type, is_nullable::bool, column_default, '
                'numeric_precision, numeric_scale '
                'FROM information_schema.columns '
                'WHERE table_schema = %s AND table_name = %s AND NOT is_hidden::bool',
                (schema or self.default_schema_name, table_name),
            )

        res = []
        for row in rows:
            name, type_str, nullable, default = row[:4]
            # When there are type parameters, attach them to the
            # returned type object.
            m = re.match(r'^(\w+(?: \w+)*)(?:\(([0-9, ]*)\))?$', type_str)
            if m is None:
                warn("Could not parse type name '%s'" % type_str)
                typ = sqltypes.NULLTYPE()
            else:
                type_name, type_args = m.groups()
                try:
                    type_class = _type_map[type_name.lower()]
                except KeyError:
                    warn("Did not recognize type '%s' of column '%s'" %
                         (type_name, name))
                    type_class = sqltypes.NULLTYPE
                if type_args:
                    typ = type_class(*[int(s.strip()) for s in type_args.split(',')])
                elif type_class is sqltypes.DECIMAL:
                    typ = type_class(
                        precision=row.numeric_precision,
                        scale=row.numeric_scale,
                    )
                else:
                    typ = type_class()
            res.append(dict(
                name=name,
                type=typ,
                nullable=nullable,
                default=default,
            ))
        return res

    def get_indexes(self, conn, table_name, schema=None, **kw):
        # Maps names to a bool indicating whether the index is unique.

        if not self._is_v2plus:
            # v1.1 or earlier.
            # Bad: the table name is not properly escaped.
            # Oh well. Hoping 1.1 won't be around for long.
            rows = conn.execute('''
SELECT "Name" as index_name,
       "Column" as column_name,
       "Unique" as unique,
       "Implicit" as implicit
FROM [SHOW INDEXES FROM "%s"."%s"]''' %
                                (schema or self.default_schema_name, table_name))
        else:
            # v2.0: usable information schema.
            rows = conn.execute('''<commit_msg>Preserve length of strings<commit_after>import collections
import re
import threading
from sqlalchemy.dialects.postgresql.base import PGDialect
from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
from sqlalchemy.dialects.postgresql import INET
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.util import warn
import sqlalchemy.sql as sql

import sqlalchemy.types as sqltypes

from .stmt_compiler import CockroachCompiler, CockroachIdentifierPreparer

# Map type names (as returned by information_schema) to sqlalchemy type
# objects.
#
# TODO(bdarnell): test more of these. The stock test suite only covers
# a few basic ones.
_type_map = {
    'bool': sqltypes.BOOLEAN,  # introspection returns "BOOL" not boolean
    'boolean': sqltypes.BOOLEAN,

    'bigint': sqltypes.INT,
    'int': sqltypes.INT,
    'int2': sqltypes.INT,
    'int4': sqltypes.INT,
    'int64': sqltypes.INT,
    'int8': sqltypes.INT,
    'integer': sqltypes.INT,
    'smallint': sqltypes.INT,

    'double precision': sqltypes.FLOAT,
    'float': sqltypes.FLOAT,
    'float4': sqltypes.FLOAT,
    'float8': sqltypes.FLOAT,
    'real': sqltypes.FLOAT,

    'dec': sqltypes.DECIMAL,
    'decimal': sqltypes.DECIMAL,
    'numeric': sqltypes.DECIMAL,

    'date': sqltypes.DATE,

    'time': sqltypes.Time,
    'time without time zone': sqltypes.Time,

    'timestamp': sqltypes.TIMESTAMP,
    'timestamptz': sqltypes.TIMESTAMP,
    'timestamp with time zone': sqltypes.TIMESTAMP,
    'timestamp without time zone': sqltypes.TIMESTAMP,

    'interval': sqltypes.Interval,

    'char': sqltypes.VARCHAR,
    'char varying': sqltypes.VARCHAR,
    'character': sqltypes.VARCHAR,
    'character varying': sqltypes.VARCHAR,
    'string': sqltypes.VARCHAR,
    'text': sqltypes.Text,
    'varchar': sqltypes.VARCHAR,

    'blob': sqltypes.BLOB,
    'bytea': sqltypes.BLOB,
    'bytes': sqltypes.BLOB,

    'json': sqltypes.JSON,
    'jsonb': sqltypes.JSON,

    'uuid': UUID,

    'inet': INET,
}


class _SavepointState(threading.local):
    """Hack to override names used in savepoint statements.

    To get the Session to do the right thing with transaction retries,
    we use the begin_nested() method, which executes a savepoint. We
    need to transform the savepoint statements that are a part of this
    retry loop, while leaving other savepoints alone. Unfortunately
    the interface leaves us with no way to pass this information along
    except via a thread-local variable.
    """
    def __init__(self):
        self.cockroach_restart = False


savepoint_state = _SavepointState()


class CockroachDBDialect(PGDialect_psycopg2):
    name = 'cockroachdb'
    supports_comments = False
    supports_sequences = False
    statement_compiler = CockroachCompiler
    preparer = CockroachIdentifierPreparer

    def __init__(self, *args, **kwargs):
        super(CockroachDBDialect, self).__init__(*args,
                                                 use_native_hstore=False,
                                                 server_side_cursors=False,
                                                 **kwargs)

    def initialize(self, connection):
        # Bypass PGDialect's initialize implementation, which looks at
        # server_version_info and performs postgres-specific queries
        # to detect certain features on the server. Set the attributes
        # by hand and hope things don't change out from under us too
        # often.
        super(PGDialect, self).initialize(connection)
        self.implicit_returning = True
        self.supports_native_enum = False
        self.supports_smallserial = False
        self._backslash_escapes = False
        self._has_native_jsonb = True
        sversion = connection.scalar("select version()")
        self._is_v2plus = " v2." in sversion
        self._is_v21plus = self._is_v2plus and (" v2.0." not in sversion)
        self._has_native_json = self._is_v2plus
        self._has_native_jsonb = self._is_v2plus

    def _get_server_version_info(self, conn):
        # PGDialect expects a postgres server version number here,
        # although we've overridden most of the places where it's
        # used.
        return (9, 5, 0)

    def get_table_names(self, conn, schema=None, **kw):
        # Upstream implementation needs correlated subqueries.

        if not self._is_v2plus:
            # v1.1 or earlier.
            return [row.Table for row in conn.execute("SHOW TABLES")]

        # v2.0+ have a good information schema. Use it.
        return [row.table_name for row in conn.execute(
            "SELECT table_name FROM information_schema.tables WHERE table_schema=%s",
            (schema or self.default_schema_name,))]

    def has_table(self, conn, table, schema=None):
        # Upstream implementation needs pg_table_is_visible().
        return any(t == table for t in self.get_table_names(conn, schema=schema))

    # The upstream implementations of the reflection functions below depend on
    # correlated subqueries which are not yet supported.
    def get_columns(self, conn, table_name, schema=None, **kw):
        if not self._is_v2plus:
            # v1.1.
            # Bad: the table name is not properly escaped.
            # Oh well. Hoping 1.1 won't be around for long.
            rows = conn.execute('SHOW COLUMNS FROM "%s"."%s"' %
                                (schema or self.default_schema_name, table_name))
        else:
            # v2.0 or later. Information schema is usable.
            rows = conn.execute(
                'SELECT column_name, data_type, is_nullable::bool, column_default, '
                'numeric_precision, numeric_scale, character_maximum_length '
                'FROM information_schema.columns '
                'WHERE table_schema = %s AND table_name = %s AND NOT is_hidden::bool',
                (schema or self.default_schema_name, table_name),
            )

        res = []
        for row in rows:
            name, type_str, nullable, default = row[:4]
            # When there are type parameters, attach them to the
            # returned type object.
            m = re.match(r'^(\w+(?: \w+)*)(?:\(([0-9, ]*)\))?$', type_str)
            if m is None:
                warn("Could not parse type name '%s'" % type_str)
                typ = sqltypes.NULLTYPE()
            else:
                type_name, type_args = m.groups()
                try:
                    type_class = _type_map[type_name.lower()]
                except KeyError:
                    warn("Did not recognize type '%s' of column '%s'" %
                         (type_name, name))
                    type_class = sqltypes.NULLTYPE
                if type_args:
                    typ = type_class(*[int(s.strip()) for s in type_args.split(',')])
                elif type_class is sqltypes.DECIMAL:
                    typ = type_class(
                        precision=row.numeric_precision,
                        scale=row.numeric_scale,
                    )
                elif type_class is sqltypes.VARCHAR:
                    typ = type_class(length=row.character_maximum_length)
                else:
                    typ = type_class()
            res.append(dict(
                name=name,
                type=typ,
                nullable=nullable,
                default=default,
            ))
        return res

    def get_indexes(self, conn, table_name, schema=None, **kw):
        # Maps names to a bool indicating whether the index is unique.

        if not self._is_v2plus:
            # v1.1 or earlier.
            # Bad: the table name is not properly escaped.
            # Oh well. Hoping 1.1 won't be around for long.
            rows = conn.execute('''
SELECT "Name" as index_name,
       "Column" as column_name,
       "Unique" as unique,
       "Implicit" as implicit
FROM [SHOW INDEXES FROM "%s"."%s"]''' %
                                (schema or self.default_schema_name, table_name))
        else:
            # v2.0: usable information schema.
            rows = conn.execute('''
{'fraction_non_alphanumeric': 0.09135277186835189, 'fraction_numerical': 0.005506194468777374, 'mean_word_length': 2.194044764188649, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2626489', 'n_tokens_mistral': 4710, 'n_tokens_neox': 4513, 'n_words': 1323}