Python sql_connection.TRN类代码示例

    def update(self, md_template):
        r"""Update values in the template

        Parameters
        ----------
        md_template : DataFrame
            The metadata template file contents indexed by samples Ids

        Raises
        ------
        QiitaDBError
            If md_template and db do not have the same sample ids
            If md_template and db do not have the same column headers
            If self.can_be_updated is not True
        """
        with TRN:
            # Clean and validate the metadata template given
            new_map = self._clean_validate_template(md_template, self.study_id, self.columns_restrictions)
            # Retrieving current metadata
            sql = "SELECT * FROM qiita.{0}".format(self._table_name(self.id))
            TRN.add(sql)
            current_map = self._transform_to_dict(TRN.execute_fetchindex())
            current_map = pd.DataFrame.from_dict(current_map, orient="index")

            # simple validations of sample ids and column names
            samples_diff = set(new_map.index).difference(current_map.index)
            if samples_diff:
                raise QiitaDBError(
                    "The new template differs from what is stored "
                    "in database by these samples names: %s" % ", ".join(samples_diff)
                )
            columns_diff = set(new_map.columns).difference(current_map.columns)
            if columns_diff:
                raise QiitaDBError(
                    "The new template differs from what is stored "
                    "in database by these columns names: %s" % ", ".join(columns_diff)
                )

            # here we are comparing two dataframes following:
            # http://stackoverflow.com/a/17095620/4228285
            current_map.sort(axis=0, inplace=True)
            current_map.sort(axis=1, inplace=True)
            new_map.sort(axis=0, inplace=True)
            new_map.sort(axis=1, inplace=True)
            map_diff = (current_map != new_map).stack()
            map_diff = map_diff[map_diff]
            map_diff.index.names = ["id", "column"]
            changed_cols = map_diff.index.get_level_values("column").unique()

            if not self.can_be_updated(columns=set(changed_cols)):
                raise QiitaDBError(
                    "The new template is modifying fields that cannot be "
                    "modified. Try removing the target gene fields or "
                    "deleting the processed data. You are trying to modify: %s" % ", ".join(changed_cols)
                )

            for col in changed_cols:
                self.update_category(col, new_map[col].to_dict())

            self.generate_files()

 def preprocessed_data(self):
     with TRN:
         sql = """SELECT preprocessed_data_id
                  FROM qiita.prep_template_preprocessed_data
                  WHERE prep_template_id=%s"""
         TRN.add(sql, [self.id])
         return TRN.execute_fetchflatten()

    def __getitem__(self, key):
        r"""Returns the value of the metadata category `key`

        Parameters
        ----------
        key : str
            The metadata category

        Returns
        -------
        obj
            The value of the metadata category `key`

        Raises
        ------
        KeyError
            If the metadata category `key` does not exists

        See Also
        --------
        get
        """
        with TRN:
            key = key.lower()
            if key not in self._get_categories():
                # The key is not available for the sample, so raise a KeyError
                raise KeyError(
                    "Metadata category %s does not exists for sample %s"
                    " in template %d" % (key, self._id, self._md_template.id))

            sql = """SELECT {0} FROM qiita.{1}
                     WHERE sample_id=%s""".format(key, self._dynamic_table)
            TRN.add(sql, [self._id])
            return TRN.execute_fetchlast()

    def setitem(self, column, value):
        """Sets `value` as value for the given `column`

        Parameters
        ----------
        column : str
            The column to update
        value : str
            The value to set. This is expected to be a str on the assumption
            that psycopg2 will cast as necessary when updating.

        Raises
        ------
        QiitaDBColumnError
            If the column does not exist in the table
        """
        with TRN:
            # Check if the column exist in the table
            if column not in self._get_categories():
                raise QiitaDBColumnError("Column %s does not exist in %s" %
                                         (column, self._dynamic_table))

            sql = """UPDATE qiita.{0}
                     SET {1}=%s
                     WHERE sample_id=%s""".format(self._dynamic_table, column)
            TRN.add(sql, [value, self._id])

    def test_context_manager_multiple_2(self):
        self.assertEqual(TRN._contexts_entered, 0)

        def tester():
            self.assertEqual(TRN._contexts_entered, 1)
            with TRN:
                self.assertEqual(TRN._contexts_entered, 2)
                sql = """SELECT EXISTS(
                        SELECT * FROM qiita.test_table WHERE int_column=%s)"""
                TRN.add(sql, [2])
                self.assertTrue(TRN.execute_fetchlast())
            self.assertEqual(TRN._contexts_entered, 1)

        with TRN:
            self.assertEqual(TRN._contexts_entered, 1)
            sql = """INSERT INTO qiita.test_table (str_column, int_column)
                         VALUES (%s, %s) RETURNING str_column, int_column"""
            args = [["insert1", 1], ["insert2", 2], ["insert3", 3]]
            TRN.add(sql, args, many=True)
            tester()
            self.assertEqual(TRN._contexts_entered, 1)
            self._assert_sql_equal([])

        self.assertEqual(TRN._contexts_entered, 0)
        self._assert_sql_equal([("insert1", True, 1), ("insert2", True, 2), ("insert3", True, 3)])
        self.assertEqual(TRN._connection.get_transaction_status(), TRANSACTION_STATUS_IDLE)

 def _check_id(self, id_):
     r"""Checks that the MetadataTemplate id_ exists on the database"""
     with TRN:
         sql = "SELECT EXISTS(SELECT * FROM qiita.{0} WHERE {1}=%s)".format(
             self._table, self._id_column)
         TRN.add(sql, [id_])
         return TRN.execute_fetchlast()

    def status(self):
        """The status of the prep template

        Returns
        -------
        str
            The status of the prep template

        Notes
        -----
        The status of a prep template is inferred by the status of the
        processed data generated from this prep template. If no processed
        data has been generated with this prep template; then the status
        is 'sandbox'.
        """
        with TRN:
            sql = """SELECT processed_data_status
                    FROM qiita.processed_data_status pds
                      JOIN qiita.processed_data pd
                        USING (processed_data_status_id)
                      JOIN qiita.preprocessed_processed_data ppd_pd
                        USING (processed_data_id)
                      JOIN qiita.prep_template_preprocessed_data pt_ppd
                        USING (preprocessed_data_id)
                    WHERE pt_ppd.prep_template_id=%s"""
            TRN.add(sql, [self._id])

            return infer_status(TRN.execute_fetchindex())

    def test_create_templates_from_qiime_mapping_file_reverse_linker(self):
        with TRN:
            TRN.add("SELECT last_value FROM "
                    "qiita.prep_template_prep_template_id_seq")
            curr_id = TRN.execute_fetchflatten()[0]
        obs_st, obs_pt = create_templates_from_qiime_mapping_file(
            StringIO(QIIME_MAP_WITH_REVERSE_LINKER_PRIMER),
            self.new_study, "16S")

        # Be green: clean the environment
        for template in [obs_st, obs_pt]:
            for _, fp in template.get_filepaths():
                self._clean_up_files.append(fp)

        self.assertEqual(obs_st.id, self.new_study.id)
        self.assertEqual(obs_pt.id, curr_id + 1)

        # Check that each template has the correct columns
        exp = {"physical_specimen_location", "physical_specimen_remaining",
               "dna_extracted", "sample_type", "host_subject_id", "latitude",
               "longitude", "taxon_id", "scientific_name",
               "collection_timestamp", "description"}
        self.assertEqual(set(obs_st.categories()), exp)

        exp = {"barcode", "primer", "center_name", "run_prefix", "platform",
               "library_construction_protocol", "instrument_model",
               "experiment_design_description", "reverselinkerprimer"}
        self.assertEqual(set(obs_pt.categories()), exp)

    def exists(cls, sample_id, md_template):
        r"""Checks if already exists a MetadataTemplate for the provided object

        Parameters
        ----------
        sample_id : str
            The sample id
        md_template : MetadataTemplate
            The metadata template to which the sample belongs to

        Returns
        -------
        bool
            True if already exists. False otherwise.
        """
        with TRN:
            cls._check_subclass()
            sql = """SELECT EXISTS(
                        SELECT * FROM qiita.{0}
                        WHERE sample_id=%s AND {1}=%s
                    )""".format(
                cls._table, cls._id_column
            )
            TRN.add(sql, [sample_id, md_template.id])
            return TRN.execute_fetchlast()

    def update_category(self, category, samples_and_values):
        """Update an existing column

        Parameters
        ----------
        category : str
            The category to update
        samples_and_values : dict
            A mapping of {sample_id: value}

        Raises
        ------
        QiitaDBUnknownIDError
            If a sample_id is included in values that is not in the template
        QiitaDBColumnError
            If the column does not exist in the table. This is implicit, and
            can be thrown by the contained Samples.
        ValueError
            If one of the new values cannot be inserted in the DB due to
            different types
        """
        with TRN:
            if not set(self.keys()).issuperset(samples_and_values):
                missing = set(self.keys()) - set(samples_and_values)
                table_name = self._table_name(self._id)
                raise QiitaDBUnknownIDError(missing, table_name)

            for k, v in viewitems(samples_and_values):
                sample = self[k]
                sample.setitem(category, v)

            try:
                TRN.execute()
            except ValueError as e:
                # catching error so we can check if the error is due to
                # different column type or something else

                value_types = set(type_lookup(type(value)) for value in viewvalues(samples_and_values))

                sql = """SELECT udt_name
                         FROM information_schema.columns
                         WHERE column_name = %s
                            AND table_schema = 'qiita'
                            AND (table_name = %s OR table_name = %s)"""
                TRN.add(sql, [category, self._table, self._table_name(self._id)])
                column_type = TRN.execute_fetchlast()

                if any([column_type != vt for vt in value_types]):
                    value_str = ", ".join([str(value) for value in viewvalues(samples_and_values)])
                    value_types_str = ", ".join(value_types)

                    raise ValueError(
                        'The new values being added to column: "%s" are "%s" '
                        '(types: "%s"). However, this column in the DB is of '
                        'type "%s". Please change the values in your updated '
                        "template or reprocess your template." % (category, value_str, value_types_str, column_type)
                    )

                raise e

    def _common_creation_steps(cls, md_template, obj_id):
        r"""Executes the common creation steps

        Parameters
        ----------
        md_template : DataFrame
            The metadata template file contents indexed by sample ids
        obj_id : int
            The id of the object being created
        """
        with TRN:
            cls._check_subclass()

            # Get some useful information from the metadata template
            sample_ids = md_template.index.tolist()
            headers = sorted(md_template.keys().tolist())

            # Insert values on template_sample table
            values = [[obj_id, s_id] for s_id in sample_ids]
            sql = """INSERT INTO qiita.{0} ({1}, sample_id)
                     VALUES (%s, %s)""".format(cls._table, cls._id_column)
            TRN.add(sql, values, many=True)

            # Insert rows on *_columns table
            datatypes = get_datatypes(md_template.ix[:, headers])
            # psycopg2 requires a list of tuples, in which each tuple is a set
            # of values to use in the string formatting of the query. We have
            # all the values in different lists (but in the same order) so use
            # zip to create the list of tuples that psycopg2 requires.
            values = [[obj_id, h, d] for h, d in zip(headers, datatypes)]
            sql = """INSERT INTO qiita.{0} ({1}, column_name, column_type)
                     VALUES (%s, %s, %s)""".format(cls._column_table,
                                                   cls._id_column)
            TRN.add(sql, values, many=True)

            # Create table with custom columns
            table_name = cls._table_name(obj_id)
            column_datatype = ["%s %s" % (col, dtype)
                               for col, dtype in zip(headers, datatypes)]
            sql = """CREATE TABLE qiita.{0} (
                        sample_id varchar NOT NULL, {1},
                        CONSTRAINT fk_{0} FOREIGN KEY (sample_id)
                            REFERENCES qiita.study_sample (sample_id)
                            ON UPDATE CASCADE
                     )""".format(table_name, ', '.join(column_datatype))
            TRN.add(sql)

            # Insert values on custom table
            values = as_python_types(md_template, headers)
            values.insert(0, sample_ids)
            values = [list(v) for v in zip(*values)]
            sql = """INSERT INTO qiita.{0} (sample_id, {1})
                     VALUES (%s, {2})""".format(
                table_name, ", ".join(headers),
                ', '.join(["%s"] * len(headers)))
            TRN.add(sql, values, many=True)

            # Execute all the steps
            TRN.execute()

    def test_post_commit_funcs_error(self):
        def func():
            raise ValueError()

        with self.assertRaises(RuntimeError):
            with TRN:
                TRN.add("SELECT 42")
                TRN.add_post_commit_func(func)

 def tester():
     self.assertEqual(TRN._contexts_entered, 1)
     with TRN:
         self.assertEqual(TRN._contexts_entered, 2)
         sql = """SELECT EXISTS(
                 SELECT * FROM qiita.test_table WHERE int_column=%s)"""
         TRN.add(sql, [2])
         self.assertTrue(TRN.execute_fetchlast())
     self.assertEqual(TRN._contexts_entered, 1)

    def test_context_manager_execute(self):
        with TRN:
            sql = """INSERT INTO qiita.test_table (str_column, int_column)
                 VALUES (%s, %s) RETURNING str_column, int_column"""
            args = [["insert1", 1], ["insert2", 2], ["insert3", 3]]
            TRN.add(sql, args, many=True)
            self._assert_sql_equal([])

        self._assert_sql_equal([("insert1", True, 1), ("insert2", True, 2), ("insert3", True, 3)])
        self.assertEqual(TRN._connection.get_transaction_status(), TRANSACTION_STATUS_IDLE)

 def test_execute_return(self):
     with TRN:
         sql = """INSERT INTO qiita.test_table (str_column, int_column)
                  VALUES (%s, %s) RETURNING str_column, int_column"""
         TRN.add(sql, ["test_insert", 2])
         sql = """UPDATE qiita.test_table SET bool_column = %s
                  WHERE str_column = %s RETURNING int_column"""
         TRN.add(sql, [False, "test_insert"])
         obs = TRN.execute()
         self.assertEqual(obs, [[["test_insert", 2]], [[2]]])