<commit_before>#
# This source file is part of the EdgeDB open source project.
#
# Copyright 2016-present MagicStack Inc. and the EdgeDB authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#


import collections
import hashlib
import pathlib
import pickle
import typing

import asyncpg
import immutables

from edb import errors

from edb.server import defines
from edb.server.pgsql import compiler as pg_compiler
from edb.server.pgsql import intromech

from edb.lang import edgeql
from edb.lang import graphql
from edb.lang.common import debug

from edb.lang.edgeql import ast as qlast
from edb.lang.edgeql import compiler as ql_compiler
from edb.lang.edgeql import quote as ql_quote

from edb.lang.ir import staeval as ireval

from edb.lang.schema import database as s_db
from edb.lang.schema import ddl as s_ddl
from edb.lang.schema import delta as s_delta
from edb.lang.schema import deltas as s_deltas
from edb.lang.schema import schema as s_schema
from edb.lang.schema import types as s_types

from edb.server.pgsql import delta as pg_delta
from edb.server.pgsql import dbops as pg_dbops

from . import config
from . import dbstate
from . import errormech
from . import sertypes


class CompilerDatabaseState(typing.NamedTuple):

    dbver: int
    con_args: dict
    schema: s_schema.Schema


class CompileContext(typing.NamedTuple):

    state: dbstate.CompilerConnectionState
    output_format: pg_compiler.OutputFormat
    single_query_mode: bool
    legacy_mode: bool
    graphql_mode: bool


EMPTY_MAP = immutables.Map()


class Compiler:

    _connect_args: dict
    _dbname: typing.Optional[str]
    _cached_db: typing.Optional[CompilerDatabaseState]

    def __init__(self, connect_args: dict, data_dir: str):
        self._connect_args = connect_args
        self._data_dir = pathlib.Path(data_dir)
        self._dbname = None
        self._cached_db = None
        self._cached_std_schema = None
        self._current_db_state = None

    async def _get_database(self, dbver: int) -> CompilerDatabaseState:
        if self._cached_db is not None and self._cached_db.dbver == dbver:
            return self._cached_db

        self._cached_db = None

        con_args = self._connect_args.copy()
        con_args['user'] = defines.EDGEDB_SUPERUSER
        con_args['database'] = self._dbname

        con = await asyncpg.connect(**con_args)
        try:
            im = intromech.IntrospectionMech(con)
            schema = await im.readschema(
                schema=self._get_std_schema(),
                exclude_modules=s_schema.STD_MODULES)

            db = CompilerDatabaseState(
                dbver=dbver,
                con_args=con_args,
                schema=schema)

            self._cached_db = db
            return db
        finally:
            await con.close()

    def _get_std_schema(self):
        if self._cached_std_schema is not None:
            return self._cached_std_schema

        with open(self._data_dir / 'stdschema.pickle', 'rb') as f:
            try:
                self._cached_std_schema = pickle.load(f)
            except Exception as e:
                raise RuntimeError(
                    'could not load std schema pickle') from e

        return self._cached_std_schema

    def _hash_sql(self, sql: bytes, **kwargs: bytes):
        h = hashlib.sha1(sql)
        for param, val in kwargs.items():
            h.update(param.encode('latin1'))
            h.update(val)
        return h.hexdigest().encode('latin1')

    def _new_delta_context(self, ctx: CompileContext):
        current_tx = ctx.state.current_tx()
        config = current_tx.get_config()

        context = s_delta.CommandContext()
        context.testmode = bool(config.get('__internal_testmode'))

        return context

    def _process_delta(self, ctx: CompileContext, delta, schema):
        """Adapt and process the delta command."""

        if debug.flags.delta_plan:
            debug.header('Delta Plan')
            debug.dump(delta, schema=schema)

        delta = pg_delta.CommandMeta.adapt(delta)
        context = self._new_delta_context(ctx)
        schema, _ = delta.apply(schema, context)

        if debug.flags.delta_pgsql_plan:
            debug.header('PgSQL Delta Plan')<commit_msg>backend/compiler: Load std schema once when the worker is spawned<commit_after>#
# This source file is part of the EdgeDB open source project.
#
# Copyright 2016-present MagicStack Inc. and the EdgeDB authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#


import collections
import hashlib
import pathlib
import pickle
import typing

import asyncpg
import immutables

from edb import errors

from edb.server import defines
from edb.server.pgsql import compiler as pg_compiler
from edb.server.pgsql import intromech

from edb.lang import edgeql
from edb.lang import graphql
from edb.lang.common import debug

from edb.lang.edgeql import ast as qlast
from edb.lang.edgeql import compiler as ql_compiler
from edb.lang.edgeql import quote as ql_quote

from edb.lang.ir import staeval as ireval

from edb.lang.schema import database as s_db
from edb.lang.schema import ddl as s_ddl
from edb.lang.schema import delta as s_delta
from edb.lang.schema import deltas as s_deltas
from edb.lang.schema import schema as s_schema
from edb.lang.schema import types as s_types

from edb.server.pgsql import delta as pg_delta
from edb.server.pgsql import dbops as pg_dbops

from . import config
from . import dbstate
from . import errormech
from . import sertypes


class CompilerDatabaseState(typing.NamedTuple):

    dbver: int
    con_args: dict
    schema: s_schema.Schema


class CompileContext(typing.NamedTuple):

    state: dbstate.CompilerConnectionState
    output_format: pg_compiler.OutputFormat
    single_query_mode: bool
    legacy_mode: bool
    graphql_mode: bool


EMPTY_MAP = immutables.Map()


class Compiler:

    _connect_args: dict
    _dbname: typing.Optional[str]
    _cached_db: typing.Optional[CompilerDatabaseState]

    def __init__(self, connect_args: dict, data_dir: str):
        self._connect_args = connect_args
        self._data_dir = pathlib.Path(data_dir)
        self._dbname = None
        self._cached_db = None
        self._std_schema = self._get_std_schema()
        self._current_db_state = None

    async def _get_database(self, dbver: int) -> CompilerDatabaseState:
        if self._cached_db is not None and self._cached_db.dbver == dbver:
            return self._cached_db

        self._cached_db = None

        con_args = self._connect_args.copy()
        con_args['user'] = defines.EDGEDB_SUPERUSER
        con_args['database'] = self._dbname

        con = await asyncpg.connect(**con_args)
        try:
            im = intromech.IntrospectionMech(con)
            schema = await im.readschema(
                schema=self._std_schema,
                exclude_modules=s_schema.STD_MODULES)

            db = CompilerDatabaseState(
                dbver=dbver,
                con_args=con_args,
                schema=schema)

            self._cached_db = db
            return db
        finally:
            await con.close()

    def _get_std_schema(self):
        with open(self._data_dir / 'stdschema.pickle', 'rb') as f:
            try:
                return pickle.load(f)
            except Exception as e:
                raise RuntimeError(
                    'could not load std schema pickle') from e

    def _hash_sql(self, sql: bytes, **kwargs: bytes):
        h = hashlib.sha1(sql)
        for param, val in kwargs.items():
            h.update(param.encode('latin1'))
            h.update(val)
        return h.hexdigest().encode('latin1')

    def _new_delta_context(self, ctx: CompileContext):
        current_tx = ctx.state.current_tx()
        config = current_tx.get_config()

        context = s_delta.CommandContext()
        context.testmode = bool(config.get('__internal_testmode'))

        return context

    def _process_delta(self, ctx: CompileContext, delta, schema):
        """Adapt and process the delta command."""

        if debug.flags.delta_plan:
            debug.header('Delta Plan')
            debug.dump(delta, schema=schema)

        delta = pg_delta.CommandMeta.adapt(delta)
        context = self._new_delta_context(ctx)
        schema, _ = delta.apply(schema, context)

        if debug.flags.delta_pgsql_plan:
            debug.header('PgSQL Delta Plan')
{'fraction_non_alphanumeric': 0.07459033836986594, 'fraction_numerical': 0.0023409236007661204, 'mean_word_length': 2.7960420032310176, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3494047', 'n_tokens_mistral': 2806, 'n_tokens_neox': 2735, 'n_words': 831}