getDatasetInfoAndSchema

Instructions

Implementation Reference

• index.js:198-237 (handler)

  The main handler function for the 'getDatasetInfoAndSchema' tool. It applies rate limiting, fetches dataset information and schema from Axiom's internal API using fetch, validates the response with datasetInfoSchema, converts the fields schema to TypeScript type definitions using helper functions, and returns the data as a formatted text content block.

  async ({ dataset }) => {
    const remainingTokens = datasetsLimiter.tryRemoveTokens(1);
    if (!remainingTokens) {
      throw new Error("Rate limit exceeded for dataset operations");
    }
  
    try {
      // Axiom client does not provide access to internal routes. We need to hit the API directly.
      const response = await fetch(
        `${config.internalUrl}/datasets/${dataset}/info`,
        {
          headers: {
            Authorization: `Bearer ${config.token}`,
            "X-AXIOM-ORG-ID": config.orgId,
          },
        }
      );
  
      const rawData = await response.json();
  
      // Validate the response data
      const data = datasetInfoSchema.parse(rawData);
  
      // Convert the fields to type definitions string
      const typeDefsString = getStringifiedSchema(
        convertSchemaToJSON(data.fields)
      );
  
      return {
        content: [
          {
            type: "text",
            text: JSON.stringify({ ...data, fields: typeDefsString }), // Override the fields with the type definitions
          },
        ],
      };
    } catch (error) {
      throw new Error(`Failed to list datasets: ${error.message}`);
    }
  }

• index.js:192-238 (registration)

  Registration of the 'getDatasetInfoAndSchema' tool with the MCP server, including tool name, description, input schema (dataset name), and the inline handler function.

  server.tool(
    "getDatasetInfoAndSchema",
    "Get dataset info and schema",
    {
      dataset: z.string().describe("The dataset to get info and schema for"),
    },
    async ({ dataset }) => {
      const remainingTokens = datasetsLimiter.tryRemoveTokens(1);
      if (!remainingTokens) {
        throw new Error("Rate limit exceeded for dataset operations");
      }
  
      try {
        // Axiom client does not provide access to internal routes. We need to hit the API directly.
        const response = await fetch(
          `${config.internalUrl}/datasets/${dataset}/info`,
          {
            headers: {
              Authorization: `Bearer ${config.token}`,
              "X-AXIOM-ORG-ID": config.orgId,
            },
          }
        );
  
        const rawData = await response.json();
  
        // Validate the response data
        const data = datasetInfoSchema.parse(rawData);
  
        // Convert the fields to type definitions string
        const typeDefsString = getStringifiedSchema(
          convertSchemaToJSON(data.fields)
        );
  
        return {
          content: [
            {
              type: "text",
              text: JSON.stringify({ ...data, fields: typeDefsString }), // Override the fields with the type definitions
            },
          ],
        };
      } catch (error) {
        throw new Error(`Failed to list datasets: ${error.message}`);
      }
    }
  );

• index.js:175-190 (schema)

  Zod schema used to validate the raw dataset info response from the Axiom API before processing.

  const datasetInfoSchema = z.object({
    compressedBytes: z.number(),
    compressedBytesHuman: z.string(),
    created: z.string(),
    fields: fieldsSchema,
    inputBytes: z.number(),
    inputBytesHuman: z.string(),
    maxTime: z.string(),
    minTime: z.string(),
    name: z.string(),
    numBlocks: z.number(),
    numEvents: z.number(),
    numFields: z.number(),
    quickQueries: z.null(),
    who: z.string(),
  });

• helpers/convertToSchemaJSON.js:20-49 (helper)

  Helper function that converts the flat array of Axiom dataset fields (with dotted notation for nested fields) into a nested JSON object suitable for TypeScript type generation.

  function convertSchemaToJSON(fields) {
    // Validate fields
    const validatedFields = fieldsSchema.parse(fields);
  
    const defs = {};
  
    validatedFields.forEach((field) => {
      const type = field.type || "any"; // Directly use the type from the field
      const path = field.name.split(".");
  
      let current = defs;
  
      path.forEach((key, index) => {
        // Ensure the current level is initialized as an object or the final type
        if (!current[key]) {
          current[key] = index === path.length - 1 ? type : {};
        } else if (
          index === path.length - 1 &&
          typeof current[key] === "object"
        ) {
          // If the final level was previously initialized as an object, overwrite with the type
          current[key] = type;
        }
  
        current = current[key] || {};
      });
    });
  
    return defs;
  }

• helpers/convertToSchemaJSON.js:57-72 (helper)

  Helper function that recursively converts the nested JSON schema object into a formatted multi-line string representing TypeScript type/interface definitions.

  function getStringifiedSchema(defs, indent = 2) {
    const entries = Object.entries(defs);
    const spaces = " ".repeat(indent);
  
    return `{
  ${entries
    .map(([key, value]) => {
      if (typeof value === "string") {
        return `${spaces}${key}: ${value};`;
      }
  
      return `${spaces}${key}: ${getStringifiedSchema(value, indent + 2)};`;
    })
    .join("\n")}
  ${" ".repeat(indent - 2)}}`;
  }