def generate_columns_descriptions()

in src/package/dataplexutils/metadata/wizard.py [0:0]

• 55 lines of code
• 5 McCabe index (conditional complexity)

    def generate_columns_descriptions(self, table_fqn,documentation_uri=None):
        """Generates metadata on the columns.

        Args:
            table_fqn: The fully qualified name of the table
            (e.g., 'project.dataset.table')

        Returns:
          None.

        Raises:
            NotFound: If the specified table does not exist.
        """
        try:
            logger.info(f"Generating metadata for columns in table {table_fqn}.")
            self._table_exists(table_fqn)
            table_schema_str, table_schema = self._get_table_schema(table_fqn)
            table_sample = self._get_table_sample(
                table_fqn, constants["DATA"]["NUM_ROWS_TO_SAMPLE"]
            )

            # Get additional information
            table_quality = self._get_table_quality(
                self._client_options._use_data_quality, table_fqn
            )
            table_profile = self._get_table_profile(
                self._client_options._use_profile, table_fqn
            )
            table_sources_info = self._get_table_sources_info(
                self._client_options._use_lineage_tables, table_fqn
            )
            job_sources_info = self._get_job_sources(
                self._client_options._use_lineage_processes, table_fqn
            )

            if documentation_uri == "":
                documentation_uri = None

            prompt_manager = PromptManager(
                PromtType.PROMPT_TYPE_TABLE, self._client_options
            )
            # Get prompt
            prompt_manager = PromptManager(
                PromtType.PROMPT_TYPE_COLUMN, self._client_options
            )
            column_description_prompt = prompt_manager.get_promtp()
            # We need to generate a new schema with the updated column
            # descriptions and then swap it
            updated_schema = []
            for column in table_schema:
                column_description_prompt_expanded = column_description_prompt.format(
                    column_name=column.name,
                    table_fqn=table_fqn,
                    table_schema_str=table_schema_str,
                    table_sample=table_sample,
                    table_profile=table_profile,
                    table_quality=table_quality,
                    table_sources_info=table_sources_info,
                    job_sources_info=job_sources_info,
                )
                #logger.info(f"Prompt used is: {column_description_prompt_expanded}.")
                column_description = self._llm_inference(
                    column_description_prompt_expanded,
                    documentation_uri=documentation_uri,
                )
                updated_schema.append(
                    self._get_updated_column(column, column_description)
                )
                logger.info(f"Generated column description: {column_description}.")
            self._update_table_schema(table_fqn, updated_schema)
        except Exception as e:
            logger.error(f"Generation of column description table {table_fqn} failed.")
            raise e(
                message=f"Generation of column description table {table_fqn} failed."
            )